add VOC eval

kahnchana · Mar 11, 2023 · e74faee · e74faee
1 parent 2a22b82
commit e74faee
Show file tree

Hide file tree

Showing 4 changed files with 279 additions and 157 deletions.
diff --git a/src/evaluation/datasets.py b/src/evaluation/datasets.py
@@ -1,6 +1,8 @@
 import numpy as np
 from PIL import Image
 
+from .utils import default_seg_target_transform
+
 voc_label_map = {
     0: 'background',
     1: 'aeroplane',
@@ -24,6 +26,68 @@
     19: 'train',
     20: 'television monitor'
 }
+voc_extended_classes = (
+    "background",
+    "aeroplane",
+    "bicycle",
+    "bird",
+    "boat",
+    "bottle",
+    "bus",
+    "car",
+    "cat",
+    "chair",
+    "cow",
+    "table",
+    "dog",
+    "horse",
+    "motorbike",
+    "person",
+    "potted plant",
+    "sheep",
+    "sofa",
+    "train",
+    "television monitor",
+    "bag",
+    "bed",
+    "bench",
+    "book",
+    "building",
+    "cabinet",
+    "ceiling",
+    "cloth",
+    "computer",
+    "cup",
+    "door",
+    "fence",
+    "floor",
+    "flower",
+    "food",
+    "grass",
+    "ground",
+    "keyboard",
+    "light",
+    "mountain",
+    "mouse",
+    "curtain",
+    "platform",
+    "sign",
+    "plate",
+    "road",
+    "rock",
+    "shelves",
+    "sidewalk",
+    "sky",
+    "snow",
+    "bedclothes",
+    "track",
+    "tree",
+    "truck",
+    "wall",
+    "water",
+    "window",
+    "wood",
+)
 
 
 class SegDataset:
@@ -32,77 +96,18 @@ class SegDataset:
 
 class PascalDataset(SegDataset):
 
-    CLASSES = (
-        "background",
-        "aeroplane",
-        "bicycle",
-        "bird",
-        "boat",
-        "bottle",
-        "bus",
-        "car",
-        "cat",
-        "chair",
-        "cow",
-        "table",
-        "dog",
-        "horse",
-        "motorbike",
-        "person",
-        "potted plant",
-        "sheep",
-        "sofa",
-        "train",
-        "television monitor",
-        "bag",
-        "bed",
-        "bench",
-        "book",
-        "building",
-        "cabinet",
-        "ceiling",
-        "cloth",
-        "computer",
-        "cup",
-        "door",
-        "fence",
-        "floor",
-        "flower",
-        "food",
-        "grass",
-        "ground",
-        "keyboard",
-        "light",
-        "mountain",
-        "mouse",
-        "curtain",
-        "platform",
-        "sign",
-        "plate",
-        "road",
-        "rock",
-        "shelves",
-        "sidewalk",
-        "sky",
-        "snow",
-        "bedclothes",
-        "track",
-        "tree",
-        "truck",
-        "wall",
-        "water",
-        "window",
-        "wood",
-    )
-
-    def __init__(self, root, split="train", transform=None):
+    def __init__(self, root, split="train", transform=None, target_transform=None):
         self.root = root
         assert split in ["train", "val"], f"invalid split: {split}"
         self.split = split
         anno_path = f"{root}/ImageSets/Segmentation/{split}.txt"
         self.file_list = self.get_file_list(anno_path)
         self.label_remap = voc_label_map
+
         self.transform = transform
+        if target_transform is None:
+            target_transform = default_seg_target_transform
+        self.target_transform = target_transform
 
     @staticmethod
     def get_file_list(path):
@@ -112,7 +117,7 @@ def get_file_list(path):
         return files
 
     @staticmethod
-    def load_image(path, get_arr=True):
+    def load_image(path, get_arr=False):
         if get_arr:
             return np.array(Image.open(path))
         else:
@@ -121,7 +126,7 @@ def load_image(path, get_arr=True):
     def __len__(self):
         return len(self.file_list)
 
-    def __getitem__(self, item, get_arr=True):
+    def __getitem__(self, item, get_arr=False):
         if isinstance(item, int):
             item = item % len(self.file_list)
             cur_idx = self.file_list[item]
@@ -133,19 +138,21 @@ def __getitem__(self, item, get_arr=True):
         seg_path = f"{self.root}/SegmentationClass/{cur_idx}.png"
 
         if self.transform is not None:
-            img = self.load_image(img_path, get_arr=False)
-            img = self.transform(img)
-        else:
-            img = self.load_image(img_path, get_arr=get_arr)
+            get_arr = False
+
+        img = self.load_image(img_path, get_arr=get_arr)
         seg = self.load_image(seg_path, get_arr=get_arr)
-        seg_labels = np.unique(seg)
-        seg_labels = list(set(seg_labels) - {0, 255})
 
-        return img, seg, seg_labels, cur_idx
+        if self.transform is not None:
+            img = self.transform(img)
+        if self.target_transform is not None:
+            seg = self.target_transform(seg)
+
+        return img, seg
 
 
 if __name__ == '__main__':
-    data_dir = "/Users/kanchana/Downloads/VOCdevkit/VOC2012"
+    data_dir = "/raid/datasets/pascal_voc/VOC2012"
     ds = PascalDataset(data_dir)
 
     from collections import defaultdict
@@ -154,3 +161,5 @@ def __getitem__(self, item, get_arr=True):
     for i in range(50):
         for j in ds[i][2]:
             l_set[j] += 1
+
+    print(l_set)