一、程序结构展示
下图展示的是winclip代码中的datasets目录,上篇笔记主要介绍其中的_init_.py文件内容,这次的笔记就是下面的dataset.py文件的注释。
二、代码作用
这段代码定义了一个名为 CLIPDataset
的自定义数据集类,用于加载和处理数据集。
-
导入库:导入了
os
、cv2
、numpy
和torch
库,以及Dataset
类从torch.utils.data
模块中。这些库用于文件和路径操作、图像处理、数据处理和深度学习框架 PyTorch 的相关操作。 -
设备判断:使用
torch.cuda.is_available()
函数判断当前设备是否支持 CUDA,如果支持,则将device
设置为 CUDA 设备;否则,设置为 CPU 设备。 -
CLIPDataset
类:- 这个类继承了 PyTorch 的
Dataset
类,用于定义自定义数据集。 __init__
方法用于初始化数据集,接受加载函数、类别、阶段(训练或测试)、k_shot
值和实验索引作为参数。在初始化过程中,调用了load_dataset
方法加载数据集。load_dataset
方法根据阶段加载数据集,返回图像路径、标签路径、标签和类型。__len__
方法返回数据集的长度。__getitem__
方法根据索引获取单个样本,读取图像和标签,并进行一些预处理,包括调整大小和生成图像名称。
- 这个类继承了 PyTorch 的
-
数据预处理:在
__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