论文代码阅读(一)

model/utils.py

数据预处理功能:加载函数图像,并进行归一化,读取数据等

import numpy as np
from collections import OrderedDict # 从python标准库的collections模块中导入了OrderedDict类
import os
import glob # 用于查找文件和目录匹配的通配符模式
import cv2 # 导入openCV库
import torch.utils.data as data # 导入了PyTorch中的数据加载模块,其中包括data模块,它提供了用于构建和加载自定义数据集的工具
import sys
OS_NAME = sys.platform # 获取了当前操作系统的名称,将其存储在名为OS_NAME的变量中
SEP = '\\' if OS_NAME == "win32" else '/' # 如果操作系统是Windows("win32"),则分隔符是反斜杠("")

rng = np.random.RandomState(2020) # 创建了一个随机数生成器(rng),并设置了种子值(2020)

def np_load_frame(filename, resize_height, resize_width, c):
    #TODO:定义加载图像并将其转化为numpy数组和归一化处理的函数
    """
    Load image path and convert it to numpy.ndarray. Notes that the color channels are BGR and the color space
    is normalized from [0, 255] to [-1, 1].颜色空间标准化至[-1,1]区间

    :param filename: the full path of image图像文件的完整路径
    :param resize_height: resized height期望的图像高度
    :param resize_width: resized width期望的图像宽度
    :return: numpy.ndarray
    """
    if c == 1: # c为1表示加载灰度图像
        image_decoded = np.repeat(cv2.imread(filename, 0)[:, :, None], 3, axis=-1)
    else:
        image_decoded = cv2.imread(filename)
    image_resized = cv2.resize(image_decoded, (resize_width, resize_height)) # 将加载的图像进行大小调整(resize)以适应指定的宽度和高度
    image_resized = image_resized.astype(dtype=np.float32) # 将调整大小后的图像数据类型转换为 np.float32
    image_resized = (image_resized / 127.5) - 1.0 # 归一化处理
    return image_resized # 返回归一化后的图像数据作为NumPy数组

class DataLoader(data.Dataset):
    #TODO:自定义数据集类,用于加载视频数据集
    def __init__(self, video_folder, transform, resize_height, resize_width, time_step=4, num_pred=1, c=1): # __init__ 方法:这是类的构造函数,用于初始化数据加载器的属性和参数
        self.c = c # 图像通道数,通常为1(灰度图像)或3(彩色图像)
        self.dir = video_folder # 视频数据集的文件夹路径
        self.transform = transform # 对加载的图像数据进行变换的函数
        self.videos = OrderedDict()
        self._resize_height = resize_height # 期望的图像高度
        self._resize_width = resize_width # 期望的图像宽度
        self._time_step = time_step # 时间步长,表示每个样本包含的连续帧数
        self._num_pred = num_pred # 预测的帧数,表示每个样本需要预测的未来帧数
        self.setup() # 该方法用于初始化数据集,通常在构造函数中调用。它可能会遍历指定的视频文件夹,准备视频文件列表以供后续加载。
        self.samples = self.get_all_samples() # 该方法用于获取所有样本的信息,例如每个样本的视频路径、帧数等。这些信息将用于后续的数据加载
        
        
    def setup(self):
        # TODO:自定义DataLoader类中的 setup 方法,它用于初始化数据集
        videos = glob.glob(os.path.join(self.dir, '*'))
        for idx, video in enumerate(sorted(videos)):
            video_name = video.split(SEP)[-1]
            self.videos[video_name] = {}
            self.videos[video_name]['path'] = video
            self.videos[video_name]['frame'] = glob.glob(os.path.join(video, '*.jpg'))
            self.videos[video_name]['frame'].sort() # 列表中的文件路径按文件名排序,以确保帧的顺序是正确的
            self.videos[video_name]['idx'] = idx # 视频文件夹在数据集中的索引,通常是按照文件夹的顺序分配的
            self.videos[video_name]['length'] = len(self.videos[video_name]['frame']) # 视频文件夹中图像帧的总数,即视频的总帧数
            
            
    def get_all_samples(self):
        #TODO:定义DataLoader类中的 get_all_samples 方法,作用是获取数据集中的所有样本(图像帧的文件路径列表)
        frames = []
        videos = glob.glob(os.path.join(self.dir, '*'))
        for video in sorted(videos):
            video_name = video.split(SEP)[-1] #从视频文件夹的路径中提取出文件夹的名称,作为表示视频的唯一标识符
            for i in range(len(self.videos[video_name]['frame'])-self._time_step):
                frames.append(self.videos[video_name]['frame'][i])
                # (1)对于每个视频文件夹,遍历该文件夹中的图像帧路径列表(self.videos[video_name]['frame']);
                # (2)通过 len(self.videos[video_name]['frame']) 获取帧数;
                # (3)为了生成连续的帧序列作为样本,从第0帧开始,依次遍历到倒数第 self._time_step + 1 帧;
                # (4)将每个样本的帧路径添加到名为 frames 的列表中,以便后续加载和使用.
        return frames               

    def __getitem__(self, index):
        #TODO:特殊的getitem方法,根据索引获取数据集中的一个样本。根据索引,加载相应的视频帧数据并进行预处理(如图像变换),然后将样本返回。
        video_name = self.samples[index].split(SEP)[-2] # 从 samples 列表中获取指定索引 index 处的样本路径,然后提取出视频文件夹的名称 video_name。这个名称通常表示样本属于哪个视频序列。
        frame_name = int(self.samples[index].split(SEP)[-1].split('.')[-2]) # 从样本路径中提取出帧图像的名称,并将其转换为整数

        batch = [] # 创建一个空列表 batch 用于存储样本中的图像帧数据
        for i in range(self._time_step+self._num_pred): # 循环遍历样本中的帧图像序列;
            # self._time_step 表示样本中连续的帧数,
            # self._num_pred 表示预测的帧数
            image = np_load_frame(self.videos[video_name]['frame'][frame_name+i], self._resize_height, self._resize_width, self.c) # 根据样本中的帧序号,从视频文件夹中加载相应的图像帧,并进行预处理
            #TODO:利用前面写的np_load_frame 函数将图像加载并进行归一化处理
            if self.transform is not None: # 如果定义了图像变换器transform
                batch.append(self.transform(image)) # 则将加载的图像进行变化,并将变换后的图像添加到batch列表

        return np.concatenate(batch, axis=0),self.transform(np.array([[self.videos[video_name]['idx']]])) # 返回样本数据,其中包括连接在一起的图像帧数据以及样本所属的视频序列的索引

    def __len__(self): # 一个特殊方法,用于返回数据集的总样本数量
        return len(self.samples)

model/final_future_prediction_ped2.py

网络模型架构

import numpy as np # 提供多维数组和函数,用于科学计算和数据处理
import cv2 # 计算机视觉库,用于图像处理与识别等功能
import torch # 构建和训练神经网络
import torch.nn as nn # 导入了PyTorch中的神经网络模块,包括各种层和损失函数
import torch.nn.functional as F # 导入了PyTorch中的函数式模块,其中包括用于神经网络的函数,如激活函数和池化操作在前向传播中使用

import torch.nn.modules.conv as conv # 导入卷积层模块,用于处理图像和特征映射,学习图像中的特征并用于分类、检测和分割

class AddCoords2d(nn.Module):
    #TODO(定义AddCoords2d模块):用于在输入张量中添加坐标信息通道,
    # 便于模型能够利用图像像素的位置信息进行搜索,有助于模型更好地理解和处理图像中的空间结构
    def __init__(self):
        super(AddCoords2d, self).__init__() # 初始化AddCoords2d 模块,创建4个属性,用于存储不同位置信息通道的值,用于表示图像中每个像素位置的坐标信息
        self.xx_channel, self.yy_channel = None, None
        self.xx_channel_, self.yy_channel_ = None, None

    def forward(self, input_tensor, feature_size=32): # 前向传播函数,input_tensor是输入张量,包含模型处理的数据

        batch_size_shape, channel_in_shape, dim_y, dim_x = input_tensor.shape # 获取输入张量的形状信息,依次为:批量大小、输入通道数、图像高度、图像宽度

        if self.xx_channel is None:
            xx_ones = torch.ones([1, 1, 1, dim_x], dtype=torch.float32).cuda()
            yy_ones = torch.ones([1, 1, 1, dim_y], dtype=torch.float32).cuda() # 创建用于生成坐标信息通道的基本张量,都包含1

            xx_range = (torch.arange(dim_y, dtype=torch.float32)/(dim_y-1)).cuda()
            yy_range = (torch.arange(dim_x, dtype=torch.float32)/(dim_x-1)).cuda() # 创建用于生成坐标信息通道的范围range范围,xx_range 和 yy_range。它们表示图像的水平和垂直坐标范围,并将它们归一化到[0,1]范围

            xx_range_ = (torch.arange(feature_size, dtype=torch.float32) / (feature_size - 1)).view(feature_size,1).repeat(1, dim_y//feature_size).flatten().cuda()
            yy_range_ = (torch.arange(feature_size, dtype=torch.float32) / (feature_size - 1)).view(feature_size,1).repeat(1, dim_y//feature_size).flatten().cuda() # 创建用于生成坐标信息通道的范围张量 xx_range_ 和 yy_range_,它们用于生成附加坐标信息通道,也被归一化到 [0, 1] 范围内

            xx_range = xx_range - xx_range_
            yy_range = yy_range - yy_range_

            xx_range = xx_range[None, None, :, None]
            yy_range = yy_range[None, None, :, None]

            xx_channel = torch.matmul(xx_range, xx_ones)
            yy_channel = torch.matmul(yy_range, yy_ones) # 计算坐标信息通道 xx_channel 和 yy_channel,并通过矩阵乘法将范围张量与基本张量相乘
            yy_channel = yy_channel.permute(0, 1, 3, 2) # 进行维度置换,以匹配输入张量的形状

            self.xx_channel = xx_channel * 2 - 1 # 对生成的坐标信息通道进行缩放和平移,将范围从 [0, 1] 映射到 [-1, 1]
            self.yy_channel = yy_channel * 2 - 1

            xx_range_ = xx_range_[None, None, :, None]
            yy_range_ = yy_range_[None, None, :, None]

            xx_channel_ = torch.matmul(xx_range_, xx_ones)
            yy_channel_ = torch.matmul(yy_range_, yy_ones)
            yy_channel_ = yy_channel_.permute(0, 1, 3, 2)

            self.xx_channel_ = xx_channel_ * 2 - 1
            self.yy_channel_ = yy_channel_ * 2 - 1
            
        xx_channel = self.xx_channel.repeat(batch_size_shape, 1, 1, 1)
        yy_channel = self.yy_channel.repeat(batch_size_shape, 1, 1, 1)
        xx_channel_ = self.xx_channel_.repeat(batch_size_shape, 1, 1, 1)
        yy_channel_ = self.yy_channel_.repeat(batch_size_shape, 1, 1, 1) # 将生成的坐标信息通道进行扩展,以匹配输入张量的批量大小

        out = torch.cat([input_tensor, xx_channel, xx_channel_, yy_channel, yy_channel_], dim=1) # 将它们与输入张量 input_tensor 连接在一起,形成一个新的输出张量 out
        return out

class CoordConv2d(conv.Conv2d): # 该类继承了conv.Conv2d,意味着它是一个卷积层,可以使用卷积操作处理输入数据
    #TODO(自定义CoordConv2d类):一个带有坐标信息通道的卷积层模块,
    #在卷积操作之前将输入张量添加坐标信息,以提供更多的空间信息给神经网络模型。这有助于模型更好地理解图像的空间结构
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, bias=False): # CoordConv2d 类的构造函数,用于初始化卷积层
        super(CoordConv2d, self).__init__(in_channels, out_channels, kernel_size, stride, padding) # 初始化父类conv.Conv2d
        self.addcoords2d = AddCoords2d()
        #TODO(创建名为addcoords2d 的 AddCoords2d 类的实例),该实例用于添加坐标信息通道
        self.conv = nn.Conv2d(in_channels + 4, out_channels, kernel_size, stride=stride, padding=padding, bias=bias)

    def forward(self, input_tensor):
        out = self.addcoords2d(input_tensor) # 调用 addcoords2d(input_tensor) 将输入张量传递给 AddCoords2d 模块,以添加坐标信息通道
        out = self.conv(out) # 将添加了坐标信息通道的张量 out 传递给卷积层 conv 进行卷积操作
        return out

class ResBlock(nn.Module):
    #TODO(自定义ResBlock类):用于构建神经网络中的残差块
    def __init__(self, in_channels, out_channels, mid_channels=None, bn=False): # 布尔参数值bn,用于指示是否应该在残差块中包含批量归一化层
        super(ResBlock, self).__init__()

        if mid_channels is None:
            mid_channels = out_channels

        # layers列表中定义了残差块的一系列层
        layers = [
            nn.ReLU(),
            nn.Conv2d(in_channels, mid_channels,
                      kernel_size=3, stride=1, padding=1),
            nn.ReLU(),
            nn.Conv2d(mid_channels, out_channels,
                      kernel_size=1, stride=1, padding=0)
        ]
        if bn:
            layers.insert(2, nn.BatchNorm2d(out_channels)) # 在残差块的第二个卷积层之前插入一个批量归一化层 (nn.BatchNorm2d)
        self.convs = nn.Sequential(*layers) # 所有这些层被组织成一个序列 convs,以便在前向传播中依次应用

    def forward(self, x): # 对输入张量x进行前向计算
        return x + self.convs(x)
        #TODO(skip connection)
        # 输入张量 x 首先通过序列 convs 进行前向计算,即应用了一系列的卷积和激活层
        # 将 x 与 convs(x) 相加,得到残差块的输出。这个操作是残差连接的关键部分,通过将输入张量与卷积操作的结果相加,
        # 可以实现跳跃连接(skip connection),有助于减轻梯度消失问题


class Gradient_Loss(nn.Module):
    #TODO(自定义Gradient_Loss):用于计算生成图像和目标图像之间梯度的损失
    #通过测量图像梯度的差异来评估图像指令。该损失函数常用于图像生成任务,如图像超分辨率、风格迁移等
    def __init__(self, channels=1, alpha=1): # alpha表示损失的超参数,默认为1
        super(Gradient_Loss, self).__init__()
        self.alpha = alpha
        filter = torch.FloatTensor([[-1., 1.]]).cuda() # 创建卷积核filter,用于计算图像在水平和垂直方向上的梯度差异
        self.filter_x = filter.view(1, 1, 1, 2).repeat(1, channels, 1, 1)
        self.filter_y = filter.view(1, 1, 2, 1).repeat(1, channels, 1, 1) # 创建两个卷积核filter_x 和 filter_y,分别用于水平核垂直方向上的梯度计算
        #这些卷积核会根据输入图像的通道数进行复制

    def forward(self, gen_frames, gt_frames): # 前向传播方法,用于计算梯度损失
        gen_frames_x = nn.functional.pad(gen_frames, (1, 0, 0, 0))
        gen_frames_y = nn.functional.pad(gen_frames, (0, 0, 1, 0)) # 首先,使用nn.functional.pad函数进行零填充操作,以便正确计算卷积
        gt_frames_x = nn.functional.pad(gt_frames, (1, 0, 0, 0))
        gt_frames_y = nn.functional.pad(gt_frames, (0, 0, 1, 0))

        gen_dx = nn.functional.conv2d(gen_frames_x, self.filter_x)
        gen_dy = nn.functional.conv2d(gen_frames_y, self.filter_y) # 计算生成图像和目标图像在水平和垂直方向上的梯度差异

        gt_dx = nn.functional.conv2d(gt_frames_x, self.filter_x)
        gt_dy = nn.functional.conv2d(gt_frames_y, self.filter_y) # 计算目标图像的梯度差异

        grad_diff_x = torch.abs(gt_dx - gen_dx)
        grad_diff_y = torch.abs(gt_dy - gen_dy) # 计算梯度差异的绝对值

        return grad_diff_x ** self.alpha + grad_diff_y ** self.alpha # 根据超参数计算梯度损失。返回的损失值是梯度差异的绝对值的指数幂


class Smooth_Loss(nn.Module):
    #TODO(自定义Smooth_Loss类):用于计算生成图像的平滑损失,通过测量图像的平滑程度来评估图像的质量
    #该损失函数常用于图像生成任务,如图像去噪、图像超分辨率等
    def __init__(self, channels=2, ks=7, alpha=1): # ks表示平滑滤波器的内核大小,alpha表示损失的超参数
        super(Smooth_Loss, self).__init__()
        self.alpha = alpha
        self.ks = ks
        filter = torch.FloatTensor([[-1 / (ks - 1)] * ks]).cuda() # 创建一个平滑滤波器filter,用于进行平滑操作。滤波器中间位置设为1,周围为-1 / (ks - 1),以实现平滑效果
        filter[0, ks // 2] = 1
        self.filter_x = filter.view(1, 1, 1, ks).repeat(1, channels, 1, 1)
        self.filter_y = filter.view(1, 1, ks, 1).repeat(1, channels, 1, 1) # 创建两个卷积核 filter_x 和 filter_y,分别用于水平和垂直方向上的平滑操作

    def forward(self, gen_frames): # 前向传播方法,用于计算平滑损失
        gen_frames_x = nn.functional.pad(gen_frames, (self.ks // 2, self.ks // 2, 0, 0))
        gen_frames_y = nn.functional.pad(gen_frames, (0, 0, self.ks // 2, self.ks // 2)) # 首先,通过对生成图像 gen_frames 进行水平和垂直方向上的零填充(padding)操作,以便能够正确计算卷积操作。填充的大小为 ks // 2,确保卷积后图像的大小与原始图像相同
        gen_dx = nn.functional.conv2d(gen_frames_x, self.filter_x)
        gen_dy = nn.functional.conv2d(gen_frames_y, self.filter_y) # 分别对填充后的图像进行水平和垂直方向上的卷积操作,使用之前定义的卷积核 filter_x 和 filter_y
        smooth_xy = torch.abs(gen_dx) + torch.abs(gen_dy) # 计算生成图像在水平和垂直方向上的平滑程度,即梯度的绝对值之和

        return torch.mean(smooth_xy ** self.alpha) #计算平均的平滑程度,并根据超参数 alpha 计算平滑损失。返回的损失值是平滑程度的平均值的指数幂。

    #     filter = torch.FloatTensor([[-1 / (ks**2 - 1)] * ks]*ks).cuda()
    #     filter[ks//2, ks // 2] = 1
    #
    #     self.filter = filter.view(1, 1, ks, ks).repeat(1, channels, 1, 1)
    #
    # def forward(self, gen_frames):
    #     gen_frames = nn.functional.pad(gen_frames, (self.ks // 2, self.ks // 2, self.ks // 2, self.ks // 2))
    #     smooth = nn.functional.conv2d(gen_frames, self.filter).abs()
    #
    #     return (smooth ** self.alpha).mean()
    

class Test_Loss(nn.Module):
    #TODO(自定义Test_Loss类):用于计算图像的测试损失,生成图像进行卷积操作并提取最大值来评估图像的某种特定特征
    # 这种损失函数可能用于某些图像处理任务中,但具体的应用场景需要根据问题而定
    def __init__(self, channels=1, ks=(16, 8), alpha=1): # ks 表示卷积核的大小(默认为(16, 8)),alpha 表示损失的超参数(默认为1)
        super(Test_Loss, self).__init__()
        self.alpha = alpha
        self.ks = ks
        self.c = channels
        self.filter = torch.ones((1,1,ks[0], ks[1]),dtype=torch.float32).cuda().repeat(1, channels, 1, 1)/(ks[0]*ks[1])
        # 在构造函数中,创建了一个卷积核 filter,其值为全1矩阵,大小为 (channels, ks[0], ks[1]),用于进行卷积操作。
        # 这个卷积核会根据输入图像的通道数进行复制,并除以卷积核大小的乘积,以进行平均操作

    def forward(self, gen_frames):
        shape = gen_frames.size() # 获取生成图像 gen_frames 的形状
        b,w,h = shape[0], shape[-2], shape[-1]
        gen_frames = nn.functional.pad(gen_frames.abs().view(b,self.c, w,h), (self.ks[1], self.ks[1], self.ks[0], self.ks[0]))
        gen_dx = nn.functional.conv2d(gen_frames, self.filter).max() # 使用之前定义的卷积核 filter 对填充后的图像进行卷积操作,并计算卷积结果的最大值
        
        return gen_dx # 返回计算得到的测试损失值 gen_dx,这个值表示卷积操作后的最大值

                                                                        
class VectorQuantizer(nn.Module):
    #TODO(自定义VectorQuantizer模块):执行向量量化操作,将输入编码映射到最接近的字典向量,并计算量化引入的损失
    # 模块通常用于自动编码器等深度学习模型中,以实现数据压缩和特征学习
    def __init__(self,
                 num_embeddings=50, # 表示编码的数量,即向量的字典大小,默认为50
                 embedding_dim=256, # 表示每个向量的维度,默认为256
                 beta=0.25): # 表示向量量化的损失权重,默认为0.25
        super(VectorQuantizer, self).__init__()
        self.K = num_embeddings
        self.D = embedding_dim
        self.beta = beta

        self.embedding = nn.Embedding(self.K, self.D) # 创建了嵌入层(self.embedding),用于将输入的编码映射到字典中的向量。初始化字典中的向量为均匀分布
        self.embedding.weight.data.uniform_(-1 / self.K, 1 / self.K)

        self.triplet_loss = nn.TripletMarginLoss(margin=1.0, p=2) # 创建了一个三元组损失(self.triplet_loss),用于量化损失计算。这是一种用于训练向量量化模型的损失函数
        self.get_err = Test_Loss(ks=(2,3)) # 创建了一个自定义的测试损失模块(self.get_err),用于评估向量量化的效果。Test_Loss为上一个类

    def forward(self, latents, test=False):
        latents = latents.permute(0, 2, 3, 1).contiguous()  # [B x D x H x W] -> [B x H x W x D] # 将输入编码latens重新排列,将通道维度移到最后
        latents_shape = latents.shape
        flat_latents = latents.view(-1, self.D)  # [BHW x D]

        # Compute L2 distance between latents and embedding weights
        dist = torch.sum(flat_latents ** 2, dim=1, keepdim=True) + \
               torch.sum(self.embedding.weight ** 2, dim=1) - \
               2 * torch.matmul(flat_latents, self.embedding.weight.t())  # [BHW x K]

        # Get the encoding that has the min distance
        _, idx = dist.topk(2, dim=1, largest=False, sorted=True)
        encoding_inds = idx[:, 0].unsqueeze(1) # 找到与每个编码最接近的字典向量的索引

        # Convert to one-hot encodings 将找到的索引转换为 one-hot 编码
        device = latents.device
        encoding_one_hot = torch.zeros(encoding_inds.size(0), self.K, device=device)
        encoding_one_hot.scatter_(1, encoding_inds, 1)  # [BHW x K]

        # Quantize the latents
        quantized_latents = torch.matmul(encoding_one_hot, self.embedding.weight)  # [BHW, D]
        quantized_latents = quantized_latents.view(latents_shape)  # [B x H x W x D]
        
        # Compute the VQ Losses
        commitment_loss = F.mse_loss(quantized_latents.detach(), latents)
        embedding_loss = F.mse_loss(quantized_latents, latents.detach()) # 计算向量量化引入的损失,包括预测编码与量化编码之间的均方误差损失和量化编码与原始编码之间的均方误差损失

        # Add the residue back to the latents
        quantized_latents = latents + (quantized_latents - latents).detach() # 将量化后的编码恢复为原始形状,并将通道维度移回到正确的位置

        return quantized_latents.permute(0, 3, 1, 2).contiguous(), commitment_loss * self.beta, embedding_loss  # [B x D x H x W]
        # 返回量化后的编码、均方误差损失(commitment_loss)和编码误差损失(embedding_loss)


class Encoder(torch.nn.Module):
    #TODO(定义Encoder模块):用于构建编码器,该编码器可以用于提取输入数据的特征,并生成掩码。
    # 这通常用于图像处理任务中,以实现特征学习和图像分割等任务
    def __init__(self, t_length = 5, n_channel =3):
        super(Encoder, self).__init__()
        
        def Basic(intInput, intOutput): #TODO: 该函数定义了一个基本的卷积块,包括两个卷积层、批归一化和激活函数
            return torch.nn.Sequential(
                torch.nn.Conv2d(in_channels=intInput, out_channels=intOutput, kernel_size=3, stride=1, padding=1, bias=False),
                torch.nn.BatchNorm2d(intOutput),
                torch.nn.ReLU(inplace=True),
                torch.nn.Conv2d(in_channels=intOutput, out_channels=intOutput, kernel_size=3, stride=1, padding=1, bias=False),
                torch.nn.BatchNorm2d(intOutput),
                torch.nn.ReLU(inplace=True)
            )
        
        def Basic_(intInput, intOutput): #TODO: 定义了另一个基本的卷积块,包括两个卷积层和批归一化,但没有激活函数
            return torch.nn.Sequential(
                torch.nn.Conv2d(in_channels=intInput, out_channels=intOutput, kernel_size=3, stride=1, padding=1, bias=False),
                torch.nn.BatchNorm2d(intOutput),
                torch.nn.ReLU(inplace=True),
                torch.nn.Conv2d(in_channels=intOutput, out_channels=intOutput, kernel_size=3, stride=1, padding=1),
            )

        def Upsample(nc, intOutput): #TODO: 定义了一个上采样块,用于上采样特征图
            return torch.nn.Sequential(
                torch.nn.ConvTranspose2d(in_channels=nc, out_channels=intOutput, kernel_size=3, stride=2, padding=1,
                output_padding=1, bias=False),
                torch.nn.BatchNorm2d(intOutput),
                torch.nn.ReLU(inplace=True)
                )
        def Mask(intInput, intOutput, nc):#TODO: 定义了一个用于生成掩码的块,用于生成掩码张量
            # 在构造函数中,创建了多个卷积块和均值池化层,以构建编码器的不同层级。还创建了用于生成掩码的块
            return torch.nn.Sequential(
                    torch.nn.Conv2d(in_channels=intInput, out_channels=nc, kernel_size=3, stride=1, padding=1, bias=False),
                    torch.nn.BatchNorm2d(nc),
                    torch.nn.ReLU(inplace=True),
                    torch.nn.Conv2d(in_channels=nc, out_channels=nc, kernel_size=3, stride=1, padding=1, bias=False),
                    torch.nn.BatchNorm2d(nc),
                    torch.nn.ReLU(inplace=True),
                    torch.nn.Conv2d(in_channels=nc, out_channels=intOutput, kernel_size=3, stride=1, padding=1),
                    torch.nn.Sigmoid()
                    )
        
        self.moduleConv1 = Basic(n_channel*(t_length-1), 32)
        self.modulePool1 = torch.nn.MaxPool2d(kernel_size=2, stride=2)

        self.moduleConv2 = Basic(32, 64)
        self.modulePool2 = torch.nn.MaxPool2d(kernel_size=2, stride=2)
        
        self.moduleConv3 = Basic(64, 128)
        self.modulePool3 = torch.nn.MaxPool2d(kernel_size=2, stride=2)

        self.moduleConv4 = Basic_(128, 256)

        self.mask_ = Upsample(64, 32)
        self.mask = Mask(64, n_channel, 32)

        self.skip = Basic(128, 32) 
        self.skip_ = Basic(256, 32)
        
    def forward(self, x): # 用于执行编码器的前向传播操作
        tensorConv1 = self.moduleConv1(x)
        tensorPool1 = self.modulePool1(tensorConv1)

        tensorConv2 = self.moduleConv2(tensorPool1)
        tensorPool2 = self.modulePool2(tensorConv2)

        tensorConv3 = self.moduleConv3(tensorPool2)
        tensorPool3 = self.modulePool3(tensorConv3)

        tensorConv4 = self.moduleConv4(tensorPool3) # 对输入张量 x 执行一系列卷积和池化操作,以提取特征信息

        mask = self.mask(torch.cat([self.mask_(tensorConv2), tensorConv1], dim=1)) # 使用 mask_ 块生成掩码,并使用 Mask 块生成掩码张量
        
        return tensorConv4, self.skip(tensorConv3.detach()),self.skip_(tensorConv4.detach()), mask
        # 返回编码器的输出,以及两个跳跃连接(skip connection)的结果,以及生成的掩码


class Decoder(torch.nn.Module):
    #TODO(定义Decoder类):用于实现解码器部分的功能。
    # 该解码器用于从编码器的特征中生成图像或其他输出。这通常用于图像生成、分割或恢复等任务中
    def __init__(self, t_length=5, n_channel=3):
        super(Decoder, self).__init__()

        def Basic(intInput, intOutput): # 定义了一个基本的卷积块,包括两个卷积层、批归一化和激活函数
            return torch.nn.Sequential(
                CoordConv2d(intInput, intOutput, kernel_size=3, padding=1),
                torch.nn.BatchNorm2d(intOutput),
                torch.nn.ReLU(inplace=True),
                CoordConv2d(intOutput, intOutput, kernel_size=3, padding=1),
                torch.nn.BatchNorm2d(intOutput),
                torch.nn.ReLU(inplace=True)
            )

        def Gen(intInput, intOutput, nc): # 定义了一个用于生成图像的块,包括三个卷积层、批归一化和Tanh激活函数。这个块用于生成解码后的图像
            return torch.nn.Sequential(
                torch.nn.Conv2d(in_channels=intInput, out_channels=nc, kernel_size=3, stride=1, padding=1, bias=False),
                torch.nn.BatchNorm2d(nc),
                torch.nn.ReLU(inplace=True),
                torch.nn.Conv2d(in_channels=nc, out_channels=nc, kernel_size=3, stride=1, padding=1, bias=False),
                torch.nn.BatchNorm2d(nc),
                torch.nn.ReLU(inplace=True),
                torch.nn.Conv2d(in_channels=nc, out_channels=intOutput, kernel_size=3, stride=1, padding=1),
                torch.nn.Tanh()
            )

        def Upsample(nc, intOutput):# 定义了一个上采样块,用于上采样特征图
            return torch.nn.Sequential(
                torch.nn.ConvTranspose2d(in_channels=nc, out_channels=intOutput, kernel_size=3, stride=2, padding=1,
                                         output_padding=1, bias=False),
                torch.nn.BatchNorm2d(intOutput),
                torch.nn.ReLU(inplace=True)
            )

        self.moduleConv = Basic(256+32, 256)
        self.moduleUpsample4 = Upsample(256, 128)

        self.moduleDeconv3 = Basic(128+32, 128)
        self.moduleUpsample3 = Upsample(128, 64)

        self.moduleDeconv2 = Basic(64, 64)
        self.moduleUpsample2 = Upsample(64, 32)

        self.moduleDeconv1 = Gen(32, n_channel, 32)

    def forward(self, x, skip3, skip4): # 前向传播方法,用于执行解码器的前向传播操作
        tensorConv = self.moduleConv(torch.cat([x, skip4],dim=1)) # 对输入张量 x 执行一系列卷积和上采样操作,以进行特征解码
        tensorUpsample4 = self.moduleUpsample4(tensorConv)

        tensorDeconv3 = self.moduleDeconv3(torch.cat([tensorUpsample4, skip3],dim=1)) # 通过 torch.cat 将解码器的输出与跳跃连接(skip connection)的结果(skip3 和 skip4)连接在一起,以增强解码器的性能
        tensorUpsample3 = self.moduleUpsample3(tensorDeconv3)

        tensorDeconv2 = self.moduleDeconv2(tensorUpsample3)
        tensorUpsample2 = self.moduleUpsample2(tensorDeconv2)

        output = self.moduleDeconv1(tensorUpsample2) # 最终生成解码后的输出图像

        return output

class OffsetNet(torch.nn.Module):
    #TODO(自定义类OffsetNet):用于生成图像中像素的偏移信息,这些信息可以用于后续的图像处理任务,如图像合成或变换
    def __init__(self, t_length=5, n_channel=3, size=None, bn=True):
        super(OffsetNet, self).__init__()
        self.conv_offset1 = nn.Sequential( # 第一个卷积块,用于生成偏移的中间表示
            CoordConv2d(n_channel*(t_length-1), 64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            CoordConv2d(64, 64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
        )
        self.conv_offset2 = nn.Sequential( # 第二个卷积块,用于生成偏移的中间表示
            ResBlock(64, 64, bn=bn),
            nn.BatchNorm2d(64),
        )
        self.conv_offset3 = nn.Sequential( # 第三个卷积块,用于生成最终偏移值的卷积块
            nn.Conv2d(128, 64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 2, kernel_size=3, stride=1, padding=1),
            nn.Tanh(),
        )
        self.conv_offset4 = nn.Sequential( # 第四个卷积块,用于生成最终偏移值,生成的偏移值将在后续用于图像的像素偏移
            nn.Conv2d(128, 64, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 2, kernel_size=3, stride=2, padding=1),
            nn.Tanh(),
        )
        self.size = size

    def forward(self, x, max_offset1=2, max_offset2=3):
        bs, c, w, h = x.size()

        x1 = self.conv_offset1(x) # 输入张量 x 经过 conv_offset1 的处理,生成中间表示 x1
        x2 = torch.cat([self.conv_offset2(x1),x1], dim=1) # 输入张量x1经过conv_offset2的处理,生成中间表示x2
        offset1 = self.conv_offset4(x2)*2*max_offset1/w # 通过 conv_offset4,生成偏移张量 offset1
        offset2 = self.conv_offset3(x2)*2*max_offset2/w # 通过 conv_offset3,生成偏移张量 offset2,这些偏移张量将用于后续的像素偏移

        offset1 = F.interpolate(offset1, (w, h), mode='bilinear', align_corners=True).permute(0,2,3,1)
        offset2 = F.interpolate(offset2, (w, h), mode='bilinear', align_corners=True).permute(0,2,3,1)
        # 对 offset1 和 offset2 进行插值操作,以将它们的大小调整为与输入图像相同,并将其转置为适合处理的形状

        # 最后,通过与网格(grid)相加,计算最终的像素偏移,并将结果返回
        gridY = torch.linspace(-1, 1, steps=h).view(1, -1, 1, 1).expand(bs, w, h, 1)
        gridX = torch.linspace(-1, 1, steps=w).view(1, 1, -1, 1).expand(bs, w, h, 1)
        grid = torch.cat((gridX, gridY), dim=3).type(torch.float32).cuda()
        grid1 = torch.clamp(offset1 + grid, min=-1, max=1)
        grid2 = torch.clamp(offset2 + grid, min=-1, max=1)
        return grid1, grid2, offset1, offset2


class convAE(torch.nn.Module):
    #TODO(定义convAE类):用于构建图像重构、像素偏移计算和损失计算的网络,是一个用于图像处理任务的网络模型
    def __init__(self, n_channel=3,  t_length=5, memory_size=50, memory_dim=256, beta1=1, beta2=0.25, beta3=1, gamma1=1, gamma2=0.25):
        super(convAE, self).__init__()

        self.encoder = Encoder(t_length, 3) # 编码器 self.encoder 用于将输入图像编码为特征张量
        self.offset_net = OffsetNet(t_length, 3) # 偏移网络 self.offset_net 用于计算图像的像素偏移信息
        self.decoder = Decoder(t_length, 3) # 解码器 self.decoder 用于将编码的特征还原为重构图像

        self.bkg = nn.Parameter(torch.from_numpy(cv2.imread('./bkg_ped2.jpg').copy().transpose((2, 0, 1)).astype(np.float32)/127.5 - 1))
        # 背景图像 self.bkg 是一个可学习的参数,表示背景图像,用于在图像合成中使用

        self.vq_layer = VectorQuantizer(memory_size, memory_dim, beta2) # self.vq_layer 是一个向量量化器,用于执行向量量化和计算相应的损失
        self.beta = [beta1, beta2, beta3] # self.beta 是一组超参数,用于权衡不同损失项的重要性
        self.gamma = [gamma1, gamma2] # self.gamma 是一组超参数,用于权衡不同损失项的重要性
        self.mse1 = 0
        self.mse2 = 0
        self.vq_loss = 0
        self.commit_loss = 0
        self.commit_max = 0
        self.grad_loss = 0
        self.smooth_loss = 0
        self.offset_loss1 = 0
        self.offset_loss2 = 0
        self.err1 = 0
        self.err2 = 0
        self.err1_ = 0
        self.err2_ = 0
        self.mask_loss = 0

        self.loss_grad = Gradient_Loss(3)
        self.loss_smooth2 = Smooth_Loss(ks=3)
        self.loss_smooth1 = Smooth_Loss(ks=7)
        self.get_err1 = Test_Loss(ks=(7,18))
        self.get_err_grad = Test_Loss(ks=(16,14))

    def forward(self, x, test=False):
        grid1, grid2, offset1, offset2 = self.offset_net(x)

        z_e, skip3, skip4, mask = self.encoder(x)

        quantized_inputs, commit_loss, vq_loss = self.vq_layer(z_e, test)

        z_q = self.decoder(quantized_inputs,skip3, skip4)

        z_q_ = F.grid_sample(z_q, grid1, align_corners=True)
        z_q_t = F.grid_sample(z_q_, grid2, align_corners=True)
        mask1 = F.grid_sample(mask, grid1, align_corners=True)
        mask2 = F.grid_sample(mask1, grid2, align_corners=True)
        
        z_q_ = z_q_ * mask1 + (1 - mask1) * self.bkg.detach()
        z_q_t = z_q_t * mask2 + (1 - mask2) * self.bkg.detach()

        return z_q_t, z_q_, z_q, z_e, commit_loss, vq_loss, offset1, offset2, mask, mask1, mask2

    def loss_function(self, x, recon_x, z_q_, z_q, z_e, commit_loss, vq_loss, offset1=None, offset2=None, mask=None, mask1=None, mask2=None, compute_err=False):
        # 计算损失函数,包括重构损失、向量量化损失、偏移损失、梯度损失、平滑损失等
        # 根据参数 compute_err 的值,它可以选择计算额外的误差项
        self.vq_loss = vq_loss
        self.commit_loss = commit_loss

        if compute_err:
            self.grad_loss = self.get_err_grad(self.loss_grad(x, recon_x))
            self.err1 = self.get_err1(offset1[:,:,:,1].abs())
            self.mse2 = F.l1_loss(recon_x, x)
        else:
            self.mse2 = F.mse_loss(recon_x, x)
            self.mse1 = F.mse_loss(z_q_, x)
            self.grad_loss = self.beta[0]*(self.loss_grad(x, recon_x).mean() + self.loss_grad(x, z_q_).mean()*0.25)

            self.offset_loss1 = ((offset1 ** 2).sum(dim=-1) ** 0.5).mean()*0.4
            self.offset_loss2 = ((offset2 ** 2).sum(dim=-1) ** 0.5).mean()*0.4
            self.smooth_loss = self.beta[2]*(self.loss_smooth1(offset1.permute(0, 3, 1, 2)) + self.loss_smooth2(offset2.permute(0, 3, 1, 2)))

            return self.mse1 + self.mse2 + self.grad_loss + \
                   self.gamma[0] * (self.vq_loss + self.commit_loss) + \
                   self.gamma[1] * (self.offset_loss1 + self.offset_loss2 + self.smooth_loss) + \
                   F.mse_loss(self.bkg.repeat(x.size(0),1,1,1),x)

    def latest_losses(self):
        return { 'mse1': self.mse1, 'mse2': self.mse2, 'vq': self.vq_loss, 'commitment': self.commit_loss, 'offset1':self.offset_loss1, 'offset2':self.offset_loss2, 'err1_':self.err1_, 'err2_':self.err2_, 'err1':self.err1, 'err2':self.err2, 'grad':self.grad_loss, 'smooth':self.smooth_loss}
        # latest_losses 方法用于获取模型的最新损失值,
        # 包括重构损失、向量量化损失、偏移损失、梯度损失、平滑损失等,以便进行训练监控和日志记录

utils.py

计算AUC值等

import numpy as np
import os
import random
import cv2
from collections import OrderedDict # 提供了一个有序的字典数据结构,用于按照插入顺序来保存键值对
import glob # 用于查找文件路径模式的库,通常用于批量处理文件
from sklearn.metrics import roc_auc_score # 计算ROC曲线下的面积(ROC-AUC),通常用于二分类的性能评估
from multiprocessing import Process
from multiprocessing import Pool # Process 和 Pool:用于创建和管理多个进程,以实现并行计算和任务分发
from tqdm import tqdm # 创建进度条以显示程序的执行进度

import sys # 提供与Python解释器和系统交互的功能,例如获取命令行参数等

OS_NAME = sys.platform
SEP = '\\' if OS_NAME == "win32" else '/'

n=100 # 初始化变量n,设置为100
temp = [] # 创建一个空的列表temp,存储后续生成的四个数值的组合
for i in range(n+1): # 使用嵌套的3个for循环来生成4个数值的组合。这三个循环分别控制四个变量:i、j、k和n - i - j - k,它们的和为n
    for j in range(n+1-i):
        for k in range(n+1-i-j):
            temp.append([i,j,k,n -i-j-k])
temp = np.array(temp, dtype=np.float32) / n # 将temp列表转换为NumPy数组,并将其数据类型(dtype)转换为np.float32。每个数组元素都除以n,以将它们的值归一化到0到1之间

def get_bkg(tosize=None, c=3):
    #TODO:定义加载并处理背景图像的函数
    bkg_list = [] # 初始化一个空列表,用于存储后续加载的背景图像
    print(sorted(glob.glob("./bkg/*"))) # 使用glob.glob函数列出./bkg/中的所有文件,并进行循环遍历
    for img in sorted(glob.glob("./bkg/*")):
        arr = np.repeat(cv2.imread(img, 0)[:, :, None], c, axis=-1) # 使用cv2.imread读取图像,并转换为numpy图像,然后将其转换为指定的通道数c,通过在最后一个轴上复制实现
        if tosize is not None:
            arr = cv2.resize(arr, (tosize,tosize)) # cv2.resize函数将图像调整为指定大小
        bkg_list.append(arr[None]) # 对文件夹./bkg/中的所有文件执行相同的操作
    
    for img in sorted(glob.glob("./bkg_/*")):
        arr = np.repeat(cv2.imread(img, 0)[:, :, None], c, axis=-1) # 将所有处理后的图像合并为一个numpy数组,根据通道数c来决定是否添加一个额外的维度。
        if tosize is not None:
            arr = cv2.resize(arr, (tosize,tosize))
        bkg_list.append(arr[None])

    return np.concatenate(bkg_list, axis=0)[:,:,:,None] if c==1 else np.concatenate(bkg_list, axis=0) # 返回合并后的背景图像数组
    # 若c为1,则在数组的最后一个维度上添加一个额外的维度,变成四维数组。否则,将图像堆叠在一起,创建一个三维数组

def mp_get_bkg(root, save_path, n=4):
    #TODO:定义以多进程的方式批量处理图像数据的函数
    #将它们从指定目录中加载并保存到另一个指定的目录中。每个进程处理一部分图像片段,以提高处理速度
    print(save_path, os.listdir(".\\")) # 打印保存路径save_path以及当前目录下的文件列表
    if save_path not in os.listdir(".\\"): # 检查是否存在名为save_path的目录,如果不存在,则在当前目录下创建该目录
        os.mkdir(".\\"+save_path)
    save_path = ".\\"+save_path+"\\" # 更新save_path,将其设定为保存路径的完整路径,包括当前目录和目标目录名称
    patch_len = len(os.listdir(root))//n+1 # 计算每个进程需要处理的图像数量,将图像总数除以进程数量n,然后向上取整,得到patch_len,表示每个进程需要处理的图像片段数量
    p = [Process(target=gen_bkg, args=(root, save_path, (i*patch_len,(i+1)*patch_len))) for i in range(n)]
    #创建一个包含n个进程的进程列表p,每个进程都会调用函数gen_bkg,并传递以下参数
    # root:图像数据的根目录; save_path:保存背景图像的目标目录; (i*patch_len, (i+1)*patch_len):表示当前进程需要处理的图像片段的起始索引和结束索引
    for sub_p in p:
        sub_p.start() # 启动所有进程,使它们同时处理不同的图像片段


def gen_bkg(root, to_path, bound):
    #TODO:定义生成背景图像的函数,用到了高斯滤波卷积核
    #该背景图像基于输入的图像数据,通过对每个像素位置应用高斯滤波卷积核来估算背景值。生成的背景图像将保存在指定目录中
    size = 13 # 卷积核大小
    std = 0.3 * ((size - 1) * 0.5 - 1) + 0.8 # 标准差。 size和std将用于生成高斯滤波卷积核
    kernel = np.arange(size) - size // 2
    kernel = np.exp(-kernel ** 2 / (2 * std ** 2)) / (std * (2 * np.pi) ** 0.5)
    print("processing the range: ", bound) # 打印当前正在处理的图像范围,该范围由bound参数指定
    for name in tqdm(os.listdir(root)[bound[0]: bound[1]]): # 遍历指定范围内的图像数据,对每个图像执行以下操作
        bg = None # 创建一个空的bg(背景)数组,用于存储生成的背景图像
        all_path = glob.glob(root + "/" + name + "/*") # 获取目标图像路径列表all_path,该列表包含了同一类别下多个图像的文件路径
        frames = [] # 初始化一个空的frames列表,用于存储加载的图像帧
        bg = np.zeros((256,256))

        for i, img_path in enumerate(all_path):
            frames.append(cv2.resize(cv2.imread(img_path, 0).copy(), (256,256))[:,:,None])
        frames = np.concatenate(frames,axis=-1) # 使用numpy的concatenate函数将frames列表中的图像帧沿着最后一个轴(通道轴)进行拼接,以创建一个具有多个通道的图像数据

        for i in range(256): # 使用两层嵌套的循环遍历256x256的像素位置
            for j in range(256):
                cnt = np.zeros((256)) # 创建一个长度为256的数组cnt,用于存储每个像素值的计数
                for mu in range(256): # 对于每个像素值(mu)
                    cnt[mu] = (frames[i,j]==mu).sum() # 统计在当前位置 (i, j) 的所有图像帧中像素值等于mu的数量,并将结果存储在cnt数组中
                bg[i,j] = np.argmax(np.convolve(cnt, kernel, 'same')) # 使用高斯滤波卷积核kernel对cnt数组进行卷积操作,得到当前像素位置 (i, j) 的背景像素值,并将其存储在bg数组中
        cv2.imwrite(to_path+name+".jpg", bg.astype(np.uint8))
    print("the range: ", bound, " has been processed") # 打印完成处理指定范围的图像信息

def gen_bkg_(root, to_path, bound=None, hist=400,n_frame=500):
    #TODO:自定义生成背景图像的功能
    #从一组输入图像中估计背景,并将估计的背景保存到文件中
    #背景估计是通过使用背景减除技术来实现的,该技术从图像序列中学习并提取出背景信息
    for name in os.listdir(root): # 遍历指定目录(root)下的子目录,每个子目录代表一个场景
        backSub = cv2.createBackgroundSubtractorMOG2() # 使用该函数创建一个背景减除器对象,用于提取前景对象
        bg = None # 初始化背景图像,它是与输入图像大小相同的空图像,用于存储最终的背景估计
        w = None # 初始化权重图像,它也是与输入图像大小相同的空图像,用于存储最终的权重信息
        all_path = glob.glob(root+"/"+name+"/*")
        random.shuffle(all_path) # 随机打乱当前场景下的图像路径列表
        
        for i, img_path in enumerate(all_path): # 遍历当前场景下的每一帧图像
            frame = cv2.imread(img_path).copy()
            if bg is None:
                bg, w = np.zeros_like(frame,dtype=np.float32), np.zeros_like(frame,dtype=np.float32)[:,:,0][:,:,None]
            if i < hist: # 如果帧数小于hist
                backSub.apply(frame) # 则将当前帧应用于背景减除器,以便适应场景的背景
            else:
                break
        for i, img_path in enumerate(all_path):
            frame = cv2.imread(img_path).copy()
            if i < n_frame:
                bgMask = 1 - backSub.apply(frame).astype(np.float32)[:,:,None]/255 # 对于每一帧图像,计算背景掩码(bgMask),它表示前景像素为0,背景像素为1
                bg+=bgMask*frame # 将当前帧根据掩码与背景图像相加,同时更新权重图像
                w += bgMask
            else:
                break
        cv2.imwrite(to_path+name+".png", (bg/w).astype(np.uint8)) # 将估计的背景图像保存为PNG文件,并命名为场景名称

def anomaly_score_inv(psnr, max_psnr, min_psnr):
    #TODO:根据图像的峰值信噪比(PSNR,Peak Signal-to-Noise Ratio)计算异常分数,用于评估图像的质量或清晰度
    return (1.0 - ((psnr - min_psnr) / (max_psnr-min_psnr+1e-8)))
    # psnr:待计算异常分数的图像的PSNR值
    # max_psnr:PSNR的最大值,通常代表最高质量的图像
    # min_psnr:PSNR的最小值,通常代表最低质量的图像
    # 通过将输入的PSNR值减去最小PSNR值,然后除以PSNR范围(最大PSNR值与最小PSNR值之差),最后取1减去这个比值;
    # 操作将PSNR映射到区间[0, 1],其中0表示最差的图像质量(异常程度最高),1表示最好的图像质量(异常程度最低)

def anomaly_score_list_inv(psnr_list):
    #TODO:计算一组图像的异常分数列表:异常分数越高,表示图像越正常;越低表示越异常
    anomaly_score_list = list()
    max_ele = np.max(psnr_list)
    min_ele = np.min(psnr_list)

    for i in range(len(psnr_list)): # 遍历输入列表中的每个PSNR值
        anomaly_score_list.append(anomaly_score_inv(psnr_list[i], max_ele, min_ele)) # 对每个PSNR值调用 anomaly_score_inv 函数,传递给它当前的PSNR值、最大PSNR值和最小PSNR值
    return anomaly_score_list

def AUC(anomal_scores, labels):
    #TODO:计算二进制分类AUC,labels 作为真实标签,anomal_scores 作为预测分数
    try:
        frame_auc = roc_auc_score(y_true=labels, y_score=anomal_scores) # 将真实标签与预测分数传递给roc_auc_score函数
    except:
        frame_auc = roc_auc_score(y_true=labels, y_score=np.squeeze(anomal_scores)) # 若出现异常,则np.squeeze()函数将anomal_scores压缩为一维数组
    return frame_auc

def score_sum(list1, list2, alpha): # list1和list2的长度应该一样
    list_result = []
    for i in range(len(list1)):
        list_result.append((alpha*list1[i]+(1-alpha)*list2[i])) # 对两个列表进行加权求和
    return list_result

def sub_auc3(temp, list1, list2, list3, labels):
    rec = None
    maxauc = 0
    # temp 是一个包含不同的权重组合的列表。每个组合都是一个四元元组 [w1, w2, w3, w4],其中 w1、w2 和 w3 表示权重,w4 通常用于表示权重的互补部分,以保证权重之和等于1
    for comb in tqdm(temp): # 遍历 temp 中的每个权重组合,
        auc = roc_auc_score(y_true=labels, y_score=comb[0] * list1 + comb[1] * list2 + comb[2] * list3) # 使用这些权重组合对 list1、list2 和 list3 进行加权混合
        if auc > maxauc:
            maxauc = auc
            rec = comb
    return maxauc, rec # 返回最大的 ROC AUC 值 maxauc 和对应的权重组合 rec
    # 该组合表示在哪种权重下,分数列表 list1、list2 和 list3 的混合在预测二元标签 labels 时具有最大的 ROC AUC 值

def auc3_mp(list1, list2, list3, labels, weight=None):
    #TODO:定义一个并行计算函数
    #用于计算在不同的权重组合下,使用三个不同的分数列表 list1、list2 和 list3 来预测二元标签 labels 的最大 ROC AUC 值和对应的权重组合。
    if weight is None:
        temp3 = [] # 生成一个权重组合的列表 temp3,其中包含了不同的权重组合
        for i in range(n + 1):
            for j in range(n + 1 - i):
                temp3.append([i, j, n - i - j])
        temp3 = np.array(temp3, dtype=np.float32) / n
    else:
        weight = int(n * weight)
        temp3 = []
        for i in range(n - weight + 1):
            temp3.append([weight, i, n - i - weight])
        # for i in range(n + 1):
        #     temp3.append([weight, i, n - i])
        temp3 = np.array(temp3, dtype=np.float32) / n

    pool = Pool(8) # 使用多进程并行计算不同权重组合下的ROCAUC值
    length = len(temp3) // 8 +1
    res = []
    list1, list2, list3, labels = np.array(list1), np.array(list2), np.array(list3), np.array(labels)
    for i in range(8):
        res.append(pool.apply_async(sub_auc3, (temp3[length*i:length*(i+1)], list1, list2, list3, labels)))
    pool.close()
    pool.join()

    max_auc = 0
    rec = None
    for p in res:
        auc, hyp = p.get()
        if auc > max_auc:
            max_auc=auc
            rec = hyp
    return max_auc, rec

def sub_auc4(temp, list1, list2, list3, list4, labels):
    #TODO:定义计算不同权重组合的函数,以找到最大的 ROC AUC 值和对应的权重组合
    #使用四个不同的分数列表 list1、list2、list3 和 list4 来预测二元标签 labels 的最大 ROC AUC 值和对应的权重组合
    rec = None
    maxauc = 0
    for comb in tqdm(temp):
        auc = roc_auc_score(y_true=labels, y_score=comb[0] * list1 + comb[1] * list2 + comb[2] * list3 + comb[3] * list4)
        if auc > maxauc:
            maxauc = auc
            rec = comb
    return maxauc, rec

def auc4_mp(list1, list2, list3, list4, labels):
    pool = Pool(8)
    length = len(temp) // 8 +1
    res = []
    list1, list2, list3, list4, labels = np.array(list1), np.array(list2), np.array(list3), np.array(list4), np.array(labels)
    for i in range(8):
        res.append(pool.apply_async(sub_auc4, (temp[length*i:length*(i+1)], list1, list2, list3, list4, labels)))
    pool.close()
    pool.join()

    max_auc = 0
    rec = None
    for p in res:
        auc, hyp = p.get()
        if auc > max_auc:
            max_auc=auc
            rec = hyp
    return max_auc, rec


def conf_avg(x, size=69, n_conf="robust"):
#TODO(定义执行滑动平均操作的函数):根据指定的窗口大小和数据点数量控制方式来平滑输入数组 x。
# 计算输入数组 x 的滑动平均值,并返回结果数组
    if n_conf == "robust": # "robust":使用 size 的一半数据点
        n_conf = size//2
    elif n_conf == "average": # "average":使用 size 的 90% 数据点
        n_conf = int(size *0.9)
    elif n_conf == "robust+": # "robust+":使用 size 的 36% 数据点
        n_conf = round(size * 0.36)
    assert isinstance(n_conf ,int)

    a = x.copy()
    b = np.ones_like(x)
    weight = np.ones(size)
    base = (size+1)//2

    for i in range(x.shape[0] - size+1):
        a_ = a[i:i + size].copy()
        u = a_.mean()
        dif = abs(a_ - u)
        sot = np.argsort(dif)[:n_conf]
        mask = np.zeros_like(dif)
        mask[sot] = 1
        weight_ = weight * mask
        b[i+base] = np.sum(a_ * weight_) / weight_.sum()
    return b.tolist()

def filter(a, y, test_folder, R=0.01): # a为原始数据数组,y为标签数组,用于计算ROCAUC,R为滤波器的参数,默认为0.01
    #TODO:定义实现滤波操作的函数,对输入数据进行平滑操作处理
    temp = a.copy() # 创建一个临时变量 temp,并对其进行复制,以确保不会修改原始数据

    videos = OrderedDict() # 函数创建一个空的有序字典 videos,用于存储测试视频的信息,包括视频名称和长度
    videos_list = sorted(glob.glob(os.path.join(test_folder, '*'))) #
    for video in videos_list:
        video_name = video.split(SEP)[-1]
        videos[video_name] = {}
        videos[video_name]['length'] = len(glob.glob(os.path.join(video, '*.jpg')))

    R = R ** 2 # 将 R 的值平方,用于后续滤波器参数
    finished_len = 0
    outputs = np.zeros(len(temp))

    for video in videos_list:
        video_name = video.split(SEP)[-1]
        temp_len = videos[video_name]['length'] - 4

        n_iter = temp_len
        sz = (n_iter,)  # size of array
        z = temp[finished_len:finished_len + temp_len].copy()

        Q = 1e-5
        xhat = np.zeros(sz)
        P = np.zeros(sz)
        xhatminus = np.zeros(sz)
        Pminus = np.zeros(sz)
        K = np.zeros(sz)
        xhat[0] = 0.0
        P[0] = 1.0

        for k in range(1, n_iter):
            xhatminus[k] = xhat[k - 1]
            Pminus[k] = P[k - 1] + Q
            K[k] = Pminus[k] / (Pminus[k] + R)
            xhat[k] = xhatminus[k] + K[k] * (z[k] - xhatminus[k])
            P[k] = (1 - K[k]) * Pminus[k]

        xhat[0] = z[0]
        outputs[finished_len:finished_len + temp_len] = xhat
        finished_len += temp_len

    return AUC(outputs, y)

Evaluate_ped2.py

Train_ped2.py

scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs)

  • 这段代码定义了一个学习率调度器,使用余弦退火学习率策略。余弦退火是一种学习率调度策略,通过在训练的不同阶段动态调整学习率,有助于模型更快地收敛。
  • T_max 参数指定了一个循环的周期数,通常设置为总的训练周期数 args.epochs
  • 在模型训练过程中,我们往往不是采用固定的优化参数,例如学习率等,会随着训练轮数的增加进行调整。最简单常见的学习率调整策略就是阶梯式下降,例如每隔一段时间将学习率降低为原来的几分之一。PyTorch 中有学习率调度器 LRScheduler 来对各种不同的学习率调整方式进行抽象,但支持仍然比较有限,在 MMEngine 中,我们对其进行了拓展,实现了更通用的参数调度器,可以对学习率、动量等优化器相关的参数进行调整,并且支持多个调度器进行组合,应用更复杂的调度策略。
  • 学习率计划是一个预定义的框架,随着训练的进行,在历时或迭代之间调整学习率。两种最常见的学习率表技术是。

  • 恒定学习率:顾名思义,我们初始化一个学习率,在训练期间不改变它。
  • 学习率衰减:我们选择一个初始的学习率,然后按照一个调度表逐渐降低。
  • import numpy as np
    import os
    import sys
    import torch
    import torch.nn as nn # 导入PyTorch中的神经网络模块,包括各种层(如全连接层、卷积层等)和损失函数
    import torch.optim as optim # 导入了PyTorch中的优化器模块,包括各种优化算法(如Adam、SGD等),用于调整模型的权重以减小损失
    import torch.utils.data as data # 导入了PyTorch中的数据加载和处理工具,包括数据加载器和数据变换
    import torchvision.transforms as transforms # # 导入图像变换工具,进行预处理和增强
    from torch.autograd import Variable # 导入Variable类,早期版本的PyTorch中用于包装张量以进行自动微分
    from collections import OrderedDict
    import copy # 导入了Python标准库中的copy模块,用于进行对象的复制操作
    import time # 导入了Python标准库中的时间模块,用于进行时间相关的操作和测量。
    from model.utils import DataLoader   # 作者自定义的类数据集
    from model.final_future_prediction_ped2 import * # pred 2 的模型
    from utils import *
    import Evaluate_ped2  as Evaluate
    import argparse # 用于配置模型的超参数和其他设置
    
    def MNADTrain():
        parser = argparse.ArgumentParser(description="DMAD")
        parser.add_argument('--batch_size', type=int, default=8, help='batch size for training 设置训练的批次大小为8')
        parser.add_argument('--epochs', type=int, default=60, help='number of epochs for training')
        parser.add_argument('--h', type=int, default=256, help='height of input images')
        parser.add_argument('--w', type=int, default=256, help='width of input images')
        parser.add_argument('--c', type=int, default=3, help='channel of input images')
        parser.add_argument('--lr', type=float, default=2e-4, help='initial learning rate') # why not 3e-4 XD
        parser.add_argument('--dim', type=int, default=256, help='channel dimension of the memory items')
        parser.add_argument('--msize', type=int, default=200, help='number of the memory items内存中的项目数量')
        parser.add_argument('--num_workers', type=int, default=2, help='number of workers for the train loader')
        parser.add_argument('--pin_memory', default=True, help='pinned memory for faster training, use more cpu')
        parser.add_argument('--dataset_type', type=str, default='ped2', help='type of dataset: ped2, avenue, shanghai')
        parser.add_argument('--dataset_path', type=str, default='F:\\ABD_project\\Diversity_Measurable_cvpr_2023\\DataSet\\anomaly_detection\\',
                            help='directory of data 数据目录')
        parser.add_argument('--exp_dir', type=str, default='log', help='directory of log')
        parser.add_argument('--log_type', type=str, default='realtime', help='type of log: txt, realtime')
        args = parser.parse_args()
    
        np.random.seed(2021)
        torch.manual_seed(2021)
        torch.cuda.manual_seed(2021)
    
        torch.backends.cudnn.benchmark = False
        torch.backends.cudnn.deterministic = True
        os.environ["CUDA_VISIBLE_DEVICES"] = '0'
        torch.backends.cudnn.enabled = True # make sure to use cudnn for computational performance
    
        train_folder = args.dataset_path + args.dataset_type + "\\training\\frames"
    
        # Loading dataset 加载训练数据
        train_dataset = DataLoader(train_folder, transforms.Compose([transforms.ToTensor(),]),
                                   resize_height=args.h, resize_width=args.w, time_step=4, c=args.c)
    
        # train_batch = data.DataLoader(train_dataset, batch_size = args.batch_size, shuffle=True,
        #num_workers=args.num_workers, drop_last=True, pin_memory=args.pin_memory)
        train_batch = data.DataLoader(train_dataset,batch_size = args.batch_size,shuffle=False,num_workers=args.num_workers)
    
        # Model setting
        model = convAE(args.c, 5, args.msize, args.dim)
        #TODO:训练加载的模型,采用具有解耦权重衰减的Adamw优化器
        optimizer = torch.optim.AdamW([{'params': model.encoder.parameters()},
                                       {'params': model.decoder.parameters()},
                                       {'params': model.offset_net.parameters()},
                                       {'params': model.vq_layer.parameters()},
                                       {'params': model.bkg, "lr": 50*args.lr},]
                                      , lr=args.lr)
        scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs) # 定义使用余弦退火调度器
        model.cuda()
    
        # Report the training process
        log_dir = os.path.join('./exp', args.dataset_type, args.exp_dir)
        #TODO:创建一个文本日志文件,用于保存记录训练过程中的输出记录
        if not os.path.exists(log_dir):
            os.makedirs(log_dir) # 创建目录,确保日志文件可以被保存到目录下
        orig_stdout = sys.stdout # 保存原始的标准输出流对象
        f = open(os.path.join(log_dir, 'log.txt'),'w')
        if args.log_type == 'txt':
            sys.stdout = f
    
        # Training
        early_stop = {'idx' : 0,# 存储一个索引值,通常用于跟踪训练的迭代次数或轮数
                      'best_eval_auc' : 0}
        log_interval = 100 # 定义训练过程中每100次打印一次日志
        loss_dict = model.latest_losses()
        losses = {k + '_train': 0 for k, v in loss_dict.items()}
        epoch_losses = {k + '_train': 0 for k, v in loss_dict.items()}
        print("开始训练")
    
        for epoch in range(args.epochs):# 60轮
        #TODO:外部训练周期的循环,通常为epoch
            start_time = time.time()
            model.train() # 在训练模式下,模型会更新权重参数,以便进行梯度下降优化;测试时,不会更新权重参数
            for j,(imgs, _) in enumerate(train_batch): # j是当前批次的索引,'_'存储标签,通常在训练时不需要
                imgs = Variable(imgs).cuda()
                # print("图片:",imgs)
                outputs = model.forward(imgs[:,0:12]) # 调用深度学习的前向传播函数forward,并传递处理过的图像数据作为输入。前向传播将生成的模型输出用于计算损失函数
    
                optimizer.zero_grad() # 将优化器的梯度缓存清零,以准备计算梯度
                loss = model.loss_function(imgs[:,-3:], *outputs)
                # TODO:计算损失
                # print("loss:",loss)
                loss.backward() # 执行反向传播,计算损失函数关于模型权重的梯度
                optimizer.step() # 使用优化器来更新模型的权重参数,以减小损失函数
                ########################################
                # TODO:计算并累积损失值,以便在训练过程中跟踪模型性能的变化
                latest_losses = model.latest_losses()
                for key in latest_losses: # latest_losses 可能包含了不同类型的损失值,每个键代表一种损失类型
                    losses[key + '_train'] += float(latest_losses[key]) # 将 latest_losses 字典中的某个键对应的损失值转换为浮点数,然后将它加到 losses 字典中相应键的值上
                    epoch_losses[key + '_train'] += float(latest_losses[key])
    
                if j % log_interval == 0: # 减少输出的频率,以避免过多的训练日志
                    for key in latest_losses:
                        losses[key + '_train'] /= log_interval # 获取每个小批次的损失值
                    loss_string = ' '.join(['{}: {:.6f}'.format(k, v) for k, v in losses.items()]) # 将损失值格式化为一个字符串,包括损失类型和对应的平均损失值
                    #TODO:打印损失值
                    print('Train Epoch: {epoch} [{batch:5d}/{total_batch} ({percent:2d}%)]   time:'
                          ' {time:3.2f}   {loss}'
                          .format(epoch=epoch, batch=j * len(imgs),
                                  total_batch=len(train_batch) * len(imgs),
                                  percent=int(100. * j / len(train_batch)),
                                  time=time.time() - start_time,
                                  loss=loss_string))
                    start_time = time.time()
                    for key in latest_losses:
                        losses[key + '_train'] = 0 # 重置损失的累计值,以准备下一个批次的计算
    
            scheduler.step()
            if epoch>4:optimizer.param_groups[-1]['lr'] = args.lr*20
            print('----------------------------------------')
            print('Epoch:', epoch+1, '; Time:', time.time()-start_time)
            print('----------------------------------------')
    
            time_start = time.time()
    
            score = Evaluate.MNADEval(model=model)
    
            if score > early_stop['best_eval_auc']:
                early_stop['best_eval_auc'] = score
                early_stop['idx'] = 0
                torch.save(model, os.path.join(log_dir, 'model.pth'))
            else:
                early_stop['idx'] += 1
                print('Score drop! Model not saved')
    
            print('With {} epochs, auc score is: {}, best score is: {}, used time: {}'.format(epoch+1, score, early_stop['best_eval_auc'], time.time()-time_start))
    
    
        print('Training is finished')
    
        sys.stdout = orig_stdout
        f.close()
    
    
    if __name__=='__main__':
        MNADTrain()

  • 0
    点赞
  • 1
    收藏
    觉得还不错? 一键收藏
  • 1
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 1
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值