WinClip非官方复现代码学习笔记4

一、程序结构展示

下图展示的是winclip代码中的datasets目录，上篇笔记主要介绍其中的_init_.py文件内容，这次的笔记就是下面的dataset.py文件的注释。

二、代码作用

这段代码定义了一个名为 CLIPDataset 的自定义数据集类，用于加载和处理数据集。

1. 导入库：导入了 os、cv2、numpy 和 torch 库，以及 Dataset 类从 torch.utils.data 模块中。这些库用于文件和路径操作、图像处理、数据处理和深度学习框架 PyTorch 的相关操作。

2. 设备判断：使用 torch.cuda.is_available() 函数判断当前设备是否支持 CUDA，如果支持，则将 device 设置为 CUDA 设备；否则，设置为 CPU 设备。

3. CLIPDataset 类：

   • 这个类继承了 PyTorch 的 Dataset 类，用于定义自定义数据集。
   • __init__ 方法用于初始化数据集，接受加载函数、类别、阶段（训练或测试）、k_shot 值和实验索引作为参数。在初始化过程中，调用了 load_dataset 方法加载数据集。
   • load_dataset 方法根据阶段加载数据集，返回图像路径、标签路径、标签和类型。
   • __len__ 方法返回数据集的长度。
   • __getitem__ 方法根据索引获取单个样本，读取图像和标签，并进行一些预处理，包括调整大小和生成图像名称。

4. 数据预处理：在 __getitem__ 方法中，首先根据索引获取图像路径、标签、标签和类型。然后使用 OpenCV 的 cv2.imread() 函数读取图像和标签，cv2.resize() 函数调整图像和标签的大小为 (1024, 1024)，并对标签进行最近邻插值以保持二值化。最后生成图像名称，返回图像、标签、标签、图像名称和类型。

三、逐行注释

import os
import cv2
import numpy as np
import torch
from torch.utils.data import Dataset
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  # 判断运行的设备是什么
# 主要目的就是要将数据集变成可以被用的格式
class CLIPDataset(Dataset):#这个类继承了pytorch里面的Dataset类
    def __init__(self, load_function, category, phase, k_shot,
                 experiment_indx=0):  # 方法用于初始化数据集，接受加载函数、类别、阶段（训练或测试）、k_shot值和实验索引作为参数

        self.load_function = load_function
        self.phase = phase

        assert k_shot in [0, 1, 5, 10]
        assert experiment_indx in [0, 1, 2]

        self.category = category

        # load datasets
        self.img_paths, self.gt_paths, self.labels, self.types = self.load_dataset(k_shot,
                                                                                   experiment_indx)  # self.labels => good : 0, anomaly : 1

    def load_dataset(self, k_shot, experiment_indx):  # 根据阶段加载数据集，返回图像路径、标签路径、标签和类型

        (train_img_tot_paths, train_gt_tot_paths, train_tot_labels, train_tot_types), \
        (test_img_tot_paths, test_gt_tot_paths, test_tot_labels, test_tot_types) = self.load_function(self.category,
                                                                                                      k_shot,
                                                                                                      experiment_indx)

        if self.phase == 'train':

            return train_img_tot_paths, \
                   train_gt_tot_paths, \
                   train_tot_labels, \
                   train_tot_types
        else:
            return test_img_tot_paths, test_gt_tot_paths, test_tot_labels, test_tot_types

    def __len__(self):
        return len(self.img_paths)  # 返回数据集的长度

    def __getitem__(self, idx):  # 根据索引获取单个样本，读取图像和标签，并进行一些预处理，包括调整大小和生成图像名称
        img_path, gt, label, img_type = self.img_paths[idx], self.gt_paths[idx], self.labels[idx], self.types[idx]
        img = cv2.imread(img_path, cv2.IMREAD_COLOR)

        if gt == 0:
            gt = np.zeros([img.shape[0], img.shape[0]])
        else:
            gt = cv2.imread(gt, cv2.IMREAD_GRAYSCALE)
            gt[gt > 0] = 255

        img = cv2.resize(img, (1024, 1024))
        gt = cv2.resize(gt, (1024, 1024), interpolation=cv2.INTER_NEAREST)
# 如果标签是0，表示正常，将其定义为与图像大小相同的全零矩阵；否则，将标签读取为灰度图，并将非零值设置为255。
        # 然后将图像和标签调整为固定大小(1024x1024)，并将标签进行最近邻插值以保持二值化。
        img_name = f'{self.category}-{img_type}-{os.path.basename(img_path[:-4])}'

        return img, gt, label, img_name, img_type