Add: dataset continual.py

dataset/continual.py

@@ -0,0 +1,119 @@
+import os
+import torch
+from torch.utils.data import Dataset
+from torchvision import transforms
+from PIL import Image, ImageOps
+import numpy as np
+import json
+import cv2
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+
+
+import logging
+
+logger = logging.getLogger()
+
+class ImageDataset(Dataset):
+    def __init__(self,
+                 data_root,
+                 meta_file="", # meta file path
+                 resize=240,
+                 mode="train",
+                 aug=False,
+                 with_fg_mask=False,
+                 test_class='None'
+                 ):
+
+        self.data_root = data_root
+        self.resize = resize
+        self.mode = mode
+        self.test_class = test_class
+        self.with_fg_mask = with_fg_mask
+        self.aug = aug
+
+        if isinstance(meta_file, str):
+            meta_info = json.load(open(meta_file, 'r'))
+        else:
+            meta_info = meta_file
+
+        self.data_list = []
+        if self.mode == "train":
+            meta_info = meta_info[mode]
+            for cls_name, data_list in meta_info.items():
+                self.data_list.extend(data_list)
+                # for data in data_list:
+                #     if data["anomaly"] == 0:
+                #         self.data_list.append(data)
+            self.class_names = list(meta_info.keys())
+        else:
+            meta_info = meta_info[mode][test_class]
+            self.data_list.extend(meta_info)
+            self.class_names = [test_class]
+
+        
+        self.resize_img_transform = transforms.Resize((self.resize, self.resize), interpolation=Image.BICUBIC)
+        self.resize_mask_transform = transforms.Resize((self.resize, self.resize), interpolation=Image.NEAREST)
+        self.aug_transform = A.Compose([
+            A.HorizontalFlip(p=0.2),
+            A.VerticalFlip(p=0.2),
+            A.ShiftScaleRotate(shift_limit=0.2, scale_limit=0, rotate_limit=0, p=0.2),
+            # A.Rotate(limit=30, p=0.5),
+            ToTensorV2()
+        ])
+
+    def __getitem__(self, idx):
+        data = self.data_list[idx]
+        img_path, mask_path, cls_name, anomaly = data["img_path"], data["mask_path"], data["cls_name"], data["anomaly"]
+        
+        img_path = os.path.join(self.data_root, img_path)
+        mask_path = os.path.join(self.data_root, mask_path)
+        if self.with_fg_mask:
+            fg_mask_path = img_path.replace(self.data_root, "sam2_fg_mask")
+            fg_mask_path = fg_mask_path[:-3] + "png"
+
+        image = Image.open(img_path).convert('RGB')
+        image = ImageOps.exif_transpose(image)
+        image = self.resize_img_transform(image)
+
+        if anomaly == 0:
+            mask = Image.fromarray(np.zeros((self.resize, self.resize)), mode='L')
+        else:
+            mask = np.array(Image.open(mask_path).convert('L')) > 0
+            mask = Image.fromarray(mask.astype(np.uint8) * 255, mode='L')
+            mask = self.resize_mask_transform(mask)
+
+        if self.with_fg_mask:
+            fg_mask = np.array(Image.open(fg_mask_path).convert('L')) > 0
+            fg_mask = Image.fromarray(fg_mask.astype(np.uint8) * 255, mode='L')
+            fg_mask = self.resize_mask_transform(fg_mask)
+            
+        else:
+            fg_mask = torch.zeros(1, self.resize, self.resize)
+
+        if self.mode == "train" and self.aug:
+            image = np.array(image).astype(np.float32)
+            mask = np.array(mask)
+            augmented = self.aug_transform(image=image, mask=mask)
+            image = augmented['image']
+            mask = augmented['mask']
+            if self.with_fg_mask:
+                fg_mask = np.array(fg_mask)
+                fg_mask = self.aug_transform(mask=fg_mask)['mask']
+
+        else:
+            image = transforms.ToTensor()(image)
+            mask = transforms.ToTensor()(mask)
+            if self.with_fg_mask:
+                fg_mask = transforms.ToTensor()(fg_mask)
+
+        return {"image": image, "mask": mask, "fg_mask": fg_mask, "cls_name": cls_name, "anomaly": anomaly}
+
+    def __len__(self):
+        return len(self.data_list)
+
+
+
+if __name__ == '__main__':
+
+    ds = ImageDataset(is_train=True)