2021广东工业智造创新大赛-智能算法赛：瓷砖表面瑕疵质检

1. Dataset

# class DataPrefetcher(object):
#     def __init__(self, loader):
#         super(DataPrefetcher, self).__init__()
#         self.data = loader
#         self.loader = iter(self.data)
#         self.stream = torch.cuda.Stream()
#         self.preload()
#
#     def preload(self):
#         try:
#             self.next_imgs, self.next_masks, self.next_weights = next(self.loader)
#         except StopIteration:
#             self.loader = iter(self.data)
#             self.next_imgs, self.next_masks, self.next_weights = next(self.loader)
#         with torch.cuda.stream(self.stream):
#             self.next_imgs = self.next_imgs.cuda(non_blocking=True)
#             self.next_masks = self.next_masks.cuda(non_blocking=True)
#             self.next_weights = self.next_weights.cuda(non_blocking=True)
#             self.next_imgs = self.next_imgs.float() / 255.0
#             self.next_masks = self.next_masks.float() / 255.0
#             self.next_weights = self.next_weights.float() / 255.0
#
#     def __len__(self):
#         return len(self.loader)
#
#     def next(self):
#         torch.cuda.current_stream().wait_stream(self.stream)
#         imgs = self.next_imgs
#         masks = self.next_masks
#         weights = self.next_weights
#         self.preload()
#         return imgs, masks, weights


class DataPrefetcher(object):
    def __init__(self, loader):
        super(DataPrefetcher, self).__init__()
        self.data = loader
        self.loader = iter(self.data)
        self.stream = torch.cuda.Stream()
        self.preload()

    def preload(self):
        try:
            self.next_imgs = next(self.loader)
        except StopIteration:
            # self.loader = iter(self.data)
            # self.next_imgs = next(self.loader)
            self.next_imgs = None
            return
        with torch.cuda.stream(self.stream):
            self.next_imgs = self.next_imgs.cuda(non_blocking=True)

            self.next_imgs = self.next_imgs.float()

    def __len__(self):
        return len(self.loader)

    def next(self):
        torch.cuda.current_stream().wait_stream(self.stream)
        imgs = self.next_imgs
        self.preload()
        return imgs

class data_prefetcher():
    def __init__(self, loader):
        self.loader = iter(loader)
        self.stream = torch.cuda.Stream()
        self.preload()

    def preload(self):
        try:
            self.data = next(self.loader)
        except StopIteration:
            self.data = None
            return
        with torch.cuda.stream(self.stream):
            self.data = self.data.cuda(non_blocking=True)
            # self.next_target = self.next_target.cuda(non_blocking=True)
            # With Amp, it isn't necessary to manually convert data to half.
            # if args.fp16:
            #     self.next_input = self.next_input.half()
            # else:
            self.data = self.data.float()
            # self.data = self.da.sub_(self.mean).div_(self.std)

• train_dataset+val_dataset：一共5388张图片
• test_dataset：一共1762张图片

类别名			图像数量  边界框数量

边异常            419        519       
角异常            1779       1976      
白色点瑕疵        1259       1945      
浅色块瑕疵        643        981       
深色点块瑕疵      2913       7986      
光圈瑕疵          246        302    

train_id_2_category_dict = {0: "边异常", # 
                            1: "角异常", # 橙色
                            2: "白色点瑕疵", # 黄色
                            3: "浅色块瑕疵", # 橙红色
                            4: "深色点块瑕疵", # 绿色
                            5: "光圈瑕疵" # 绿色}

1.1 train_dataset

Get 4850 valid images in total
# 只有两种分辨率的输入图像，其中6000x8192的图像有3557张，3500x4096的图像有1293张
height_width =  [[6000, 8192], [3500, 4096]]

1.2 val_dataset：

Get 538 valid images in total 
# 只有两种分辨率的输入图像，其中6000x8192的图像有396张，3500x4096的图像有142张
height_width = [[6000, 8192], [3500, 4096]]