UNET多分类dataset参考

 try:
            head_mask_read = cv2.imread(head_mask,0)
            tail_mask_read = cv2.imread(tail_mask,0)
            img_read = cv2.imread(img)
            # top, bottom = 0, 28
            # left, right = 0, 16
            # mask_read = cv2.copyMakeBorder(mask_read, top, bottom, left, right, cv2.BORDER_CONSTANT, value=0)
            # # 对图像使用彩色图像 padding，这里假设原始图像是彩色的
            # img_read = cv2.copyMakeBorder(img_read, top, bottom, left, right, cv2.BORDER_CONSTANT, value=[0, 0, 0])


        except:
            img_read = np.zeros((900,1200,3))
            head_mask_read = np.zeros((900,1200,1))
            tail_mask_read = np.zeros((900,1200,1))

        head_mask = np.where(head_mask_read > 125,1, 0)
        tail_mask = np.where(tail_mask_read > 125,1, 0)
        tail_mask = np.where(head_mask == 1, 0, tail_mask)
        background = np.logical_not(head_mask | tail_mask)
        multi_mask = np.stack([head_mask, tail_mask, background], axis=0)
        # Convert to torch tensor
        multi_mask_tensor = torch.from_numpy(multi_mask).float()

        # if not self.is_val:
        #     sample = self.augmentation(image=img_read, mask=mask_read)
        #     img_read, mask_read = sample['image'], sample['mask']
            # cv2.imshow("img", img_read)
            # cv2.imshow("mask", mask_read)
            # cv2.waitKey(0)


        img = transforms1(img_read)
        mask = multi_mask_tensor
        return img,mask