fixed bug when reading from annotation file

cocoDataset.py

@@ -4,31 +4,53 @@ from skimage import io
 from random import random
 
 class CocoDataSet(Dataset):
-    def __init__(self, img_dir, label_dir, percentage_without_person=0.0):
-        x = [str(file) for file in list(Path(img_dir).glob("*.jpg"))]
-        y = [str(file) for file in list(Path(label_dir).glob("*.txt"))]
-        self.person_id = 1
+    def __init__(self, img_dir, label_dir, percentage_without_person=0.0, **kwargs):
+        """"
+        img_dir and label_dir as relative paths
+        pecent_without_person as a float in range 0.0-1.0
+        """
+        img_dir = f"{Path().resolve()}\{img_dir}"
+        label_dir = f"{Path().resolve()}\{label_dir}"
+
+        self.person_images = 0
+        self.non_person_images = 0
+
+        x = []
+        y = []
+
+        img_names = [file.stem for file in list(Path(img_dir).glob("*.jpg"))]
+        annotation_file_names = [file.stem for file in list(Path(label_dir).glob("*.txt"))]
+        # If both img and annotation file exists add to (x,y)
+        for img_name in list(set(img_names) & set(annotation_file_names)):
+            x.append(f"{img_dir}\{img_name}.jpg")
+            y.append(f"{label_dir}\{img_name}.txt")
+
+        self.person_id = 0
         self.x = []
         self.y = []
 
-        if percentage_without_person < 100.0:
+        if percentage_without_person < 1.0:
     
-            for image, labels in zip(x,y):
-                other = []
-                persons = []
-                for bb in labels:
-                    if (bb[0] == self.person_id):
-                        persons.append(bb)
+            for i, image, annotation_file in zip(range(len(x)),x,y):
+                print(f"{i}/{len(x)}", end="\r")
+
+                labels = self.read_annotation_file(annotation_file)
+                # with open(label_file) as file:
+                #     labels = [line.strip(" \n").split(" ") for line in file.readlines()]
+
+                persons = [label for label in labels if label[0] == self.person_id]
 
                 # include if there is person in the image
                 if len(persons) > 0:
                     self.x.append(image)
                     self.y.append(persons)
+                    self.person_images += 1
                 # include one element with prob. if no person is on the image
-                elif len(other) > 0 and percentage_without_person > 0.0:
-                    if percentage_without_person < random():
+                elif percentage_without_person > 0.0:
+                    if percentage_without_person > random():
                         self.x.append(image)
-                        self.y.append(other[0]) # one element with label 0 as confidence
+                        self.y.append([]) # one element with label 0 as confidence
+                        self.non_person_images += 1
         else:
             self.x = x                    
             self.y = y                    
@@ -37,20 +59,46 @@ class CocoDataSet(Dataset):
     def __len__(self):
         return len(self.x)
 
+
     def __getitem__(self, index):
         x = io.imread(self.x[index])
         y = None
 
-        with open(self.y[index], "r") as file:
-            y = [line.strip(" \n").split(" ") for line in file.readlines()]
+        # with open(self.y[index], "r") as file:
+        #     y = [line.strip(" \n").split(" ") for line in file.readlines()]        
+        y = self.read_annotation_file(self.y[index])
 
-        # set confidence/probability to 1 for person and 0 otherwise      
+        # set confidence/probability to 1 for person and 0 otherwise
+        indexes_to_remove = []
         for i, bb in enumerate(y):
-            y[i][0] = 1 if bb[0] == self.person_id else 0
+            y[i][0] = int(bb[0] == self.person_id)
+            if y[i][0] == 0:
+                indexes_to_remove.append(i)
         
+        offset = 0
+        for i in indexes_to_remove:
+            del y[i-offset]
+            offset += 1
+
         return [x,y]
 
+    def read_annotation_file(self, file):
+        """
+        Read data from an annotation/label file and returns the correct format which is 
+        [c,x,y,w,h] where c is an int and (x,y,w,h) are floats
+        """
+        data = None
+        with open(file) as file:
+            data = [[float(i) for i in line.strip(" \n").split(" ")] for line in file.readlines()]
+        data = [[int(point[0]), *point[1:]] for point in data]
+        return data
+
 if __name__ == "__main__":
-    test_dataset = CocoDataSet("./data/test/images", "./data/test/labels", 100)
-    print(test_dataset[0])
+    from PIL import Image
+    # test_dataset = CocoDataSet("./data/test/images", "./data/test/labels", 1)
+    test_dataset = CocoDataSet("./data/train2014", "./data/labels/train2014", 1)
+    print(test_dataset.x[0])
+    print(test_dataset[0][1])
+    img = Image.fromarray(test_dataset[0][0], 'RGB')
+    img.show()
         
\ No newline at end of file