Update dataset.py

dataset.py

@@ -1,93 +1,93 @@
-import os
-import numpy as np
-import torch
-import cv2
-from torch.utils.data import Dataset
-import albumentations as A
-from albumentations.pytorch import ToTensorV2
-
-
-class CocoSegmentationDataset(Dataset):
-
-    def __init__(self, coco, image_folder,
-                 category_name=None,
-                 transform=None):
-
-        self.coco = coco
-        self.image_folder = image_folder
-        self.transform = transform
-
-        if category_name:
-            self.cat_ids = self.coco.getCatIds(catNms=[category_name])
-            self.img_ids = self.coco.getImgIds(catIds=self.cat_ids)
-        else:
-            # Use all categories and all images if no specific category is provided
-            self.cat_ids = self.coco.getCatIds()
-            self.img_ids = self.coco.getImgIds()
-
-    def __len__(self):
-        return len(self.img_ids)
-
-    def __getitem__(self, index):
-
-        img_id = self.img_ids[index]
-        img_info = self.coco.loadImgs(img_id)[0]
-        img_path = os.path.join(self.image_folder, img_info['file_name'])
-
-        # Load image with OpenCV (BGR to RGB)
-        image = cv2.imread(img_path)
-        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-
-        # Fetch annotations for the image. If self.cat_ids is everything, it gets all annotations.
-        ann_ids = self.coco.getAnnIds(
-            imgIds=img_info['id'],
-            catIds=self.cat_ids,
-            iscrowd=None
-        )
-
-        anns = self.coco.loadAnns(ann_ids)
-        mask = np.zeros((img_info['height'], img_info['width']))
-
-        for ann in anns:
-            mask = np.maximum(mask, self.coco.annToMask(ann))
-
-        if self.transform:
-            augmented = self.transform(image=image, mask=mask)
-            image = augmented['image']
-            mask = augmented['mask']
-        
-        if not isinstance(mask, torch.Tensor):
-            mask = torch.from_numpy(mask).float()
-            
-        if mask.ndim == 2:
-            mask = mask.unsqueeze(0)
-
-        return image, mask
-
-
-def get_train_transforms(image_size=256):
-    return A.Compose([
-        A.LongestMaxSize(max_size=image_size),
-        A.PadIfNeeded(min_height=image_size, min_width=image_size, border_mode=cv2.BORDER_CONSTANT, fill=(123.675, 116.28, 103.53), fill_mask=0),
-        A.HorizontalFlip(p=0.5),
-        A.VerticalFlip(p=0.3),
-        A.RandomBrightnessContrast(p=0.4),
-        A.Affine(
-            scale=(0.9, 1.1),
-            rotate=(-15, 15),
-            translate_percent=(0.05, 0.05),
-            p=0.5
-        ),
-        A.GaussianBlur(p=0.2),
-        A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
-        ToTensorV2(),
-    ])
-
-
-def get_val_transforms(image_size=256):
-    return A.Compose([
-        A.LongestMaxSize(max_size=image_size),
-        A.PadIfNeeded(min_height=image_size, min_width=image_size, border_mode=cv2.BORDER_CONSTANT, fill=(123.675, 116.28, 103.53), fill_mask=0),
-        A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
-        ToTensorV2(),
-    ])
+import os
+import numpy as np
+import torch
+import cv2
+from torch.utils.data import Dataset
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+
+
+class CocoSegmentationDataset(Dataset):
+
+    def __init__(self, coco, image_folder,
+                 category_name=None,
+                 transform=None):
+
+        self.coco = coco
+        self.image_folder = image_folder
+        self.transform = transform
+
+        if category_name:
+            self.cat_ids = self.coco.getCatIds(catNms=[category_name])
+            self.img_ids = self.coco.getImgIds(catIds=self.cat_ids)
+        else:
+            # Use all categories and all images if no specific category is provided
+            self.cat_ids = self.coco.getCatIds()
+            self.img_ids = self.coco.getImgIds()
+
+    def __len__(self):
+        return len(self.img_ids)
+
+    def __getitem__(self, index):
+
+        img_id = self.img_ids[index]
+        img_info = self.coco.loadImgs(img_id)[0]
+        img_path = os.path.join(self.image_folder, img_info['file_name'])
+
+        # Load image with OpenCV (BGR to RGB)
+        image = cv2.imread(img_path)
+        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+        # Fetch annotations for the image. If self.cat_ids is everything, it gets all annotations.
+        ann_ids = self.coco.getAnnIds(
+            imgIds=img_info['id'],
+            catIds=self.cat_ids,
+            iscrowd=None
+        )
+
+        anns = self.coco.loadAnns(ann_ids)
+        mask = np.zeros((img_info['height'], img_info['width']))
+
+        for ann in anns:
+            mask = np.maximum(mask, self.coco.annToMask(ann))
+
+        if self.transform:
+            augmented = self.transform(image=image, mask=mask)
+            image = augmented['image']
+            mask = augmented['mask']
+        
+        if not isinstance(mask, torch.Tensor):
+            mask = torch.from_numpy(mask).float()
+            
+        if mask.ndim == 2:
+            mask = mask.unsqueeze(0)
+
+        return image, mask
+
+
+def get_train_transforms(image_size=256):
+    return A.Compose([
+        A.LongestMaxSize(max_size=image_size),
+        A.PadIfNeeded(min_height=image_size, min_width=image_size, border_mode=cv2.BORDER_CONSTANT, value=(123.675, 116.28, 103.53), mask_value=0),
+        A.HorizontalFlip(p=0.5),
+        A.VerticalFlip(p=0.3),
+        A.RandomBrightnessContrast(p=0.4),
+        A.Affine(
+            scale=(0.9, 1.1),
+            rotate=(-15, 15),
+            translate_percent=(0.05, 0.05),
+            p=0.5
+        ),
+        A.GaussianBlur(p=0.2),
+        A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
+        ToTensorV2(),
+    ])
+
+
+def get_val_transforms(image_size=256):
+    return A.Compose([
+        A.LongestMaxSize(max_size=image_size),
+        A.PadIfNeeded(min_height=image_size, min_width=image_size, border_mode=cv2.BORDER_CONSTANT, value=(123.675, 116.28, 103.53), mask_value=0),
+        A.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
+        ToTensorV2(),
+    ])