reorganizing files

add train
some catches

dto/annotationClass.py

@@ -1,4 +1,5 @@
 import json
+from os.path import dirname, join
 
 
 class AnnotationClass:
@@ -9,7 +10,8 @@ class AnnotationClass:
 
     @staticmethod
     def read_json():
-        with open('classes.json', 'r', encoding='utf-8') as f:
+        classes_path = join(dirname(dirname(__file__)), 'classes.json')
+        with open(classes_path, 'r', encoding='utf-8') as f:
             j = json.loads(f.read())
             return {cl['Id']: AnnotationClass(id=cl['Id'], name=cl['Name'], color=cl['Color']) for cl in j}
 

dto/imageLabel.py

@@ -0,0 +1,32 @@
+import cv2
+from matplotlib import pyplot as plt
+
+
+class ImageLabel:
+    def __init__(self, image_path, image, labels_path, labels):
+        self.image_path = image_path
+        self.image = image
+        self.labels_path = labels_path
+        self.labels = labels
+
+    def visualize(self, annotation_classes):
+        img = cv2.cvtColor(self.image.copy(), cv2.COLOR_BGR2RGB)
+        height, width, channels = img.shape
+        for label in self.labels:
+            class_num = int(label[-1])
+            x_c = float(label[0])
+            y_c = float(label[1])
+            w = float(label[2])
+            h = float(label[3])
+            x_min = x_c - w / 2
+            y_min = y_c - h / 2
+            x_max = x_min + w
+            y_max = y_min + h
+            color = annotation_classes[class_num].color_tuple
+
+            cv2.rectangle(img, (int(x_min * width), int(y_min * height)), (int(x_max * width), int(y_max * height)),
+                          color=color, thickness=3)
+        plt.figure(figsize=(12, 12))
+        plt.axis('off')
+        plt.imshow(img)
+        plt.show()

preprocessing.py

@@ -3,8 +3,8 @@ import time
 from pathlib import Path
 import albumentations as A
 import cv2
-from matplotlib import pyplot as plt
-from AnnotationClass import AnnotationClass
+from dto.annotationClass import AnnotationClass
+from dto.imageLabel import ImageLabel
 
 labels_dir = 'labels'
 images_dir = 'images'
@@ -14,36 +14,6 @@ current_labels_dir = os.path.join(current_dataset_dir, 'labels')
 annotation_classes = AnnotationClass.read_json()
 
 
-class ImageLabel:
-    def __init__(self, image_path, image, labels_path, labels):
-        self.image_path = image_path
-        self.image = image
-        self.labels_path = labels_path
-        self.labels = labels
-
-    def visualize(self):
-        img = cv2.cvtColor(self.image.copy(), cv2.COLOR_BGR2RGB)
-        height, width, channels = img.shape
-        for label in self.labels:
-            class_num = int(label[-1])
-            x_c = float(label[0])
-            y_c = float(label[1])
-            w = float(label[2])
-            h = float(label[3])
-            x_min = x_c - w / 2
-            y_min = y_c - h / 2
-            x_max = x_min + w
-            y_max = y_min + h
-            color = annotation_classes[class_num].color_tuple
-
-            cv2.rectangle(img, (int(x_min * width), int(y_min * height)), (int(x_max * width), int(y_max * height)),
-                          color=color, thickness=3)
-        plt.figure(figsize=(12, 12))
-        plt.axis('off')
-        plt.imshow(img)
-        plt.show()
-
-
 def image_processing(img_ann: ImageLabel) -> [ImageLabel]:
     transforms = [
         A.Compose([A.HorizontalFlip(always_apply=True)],
@@ -85,13 +55,16 @@ def write_result(img_ann: ImageLabel, show_image=False):
     os.makedirs(os.path.dirname(img_ann.labels_path), exist_ok=True)
 
     if show_image:
-        img_ann.visualize()
+        img_ann.visualize(annotation_classes)
 
     cv2.imwrite(img_ann.image_path, img_ann.image)
+    print(f'{img_ann.image_path} written')
+
     with open(img_ann.labels_path, 'w') as f:
         lines = [f'{ann[4]} {round(ann[0], 5)} {round(ann[1], 5)} {round(ann[2], 5)} {round(ann[3], 5)}\n' for ann in img_ann.labels]
         f.writelines(lines)
         f.close()
+    print(f'{img_ann.labels_path} written')
 
 
 def read_labels(labels_path) -> [[]]:
@@ -129,6 +102,7 @@ def main():
             continue
 
         for image in images:
+            try:
                 image_path = os.path.join(images_dir, image)
                 labels_path = os.path.join(labels_dir, f'{Path(image_path).stem}.txt')
                 process_image(ImageLabel(
@@ -137,6 +111,8 @@ def main():
                     labels_path=labels_path,
                     labels=read_labels(labels_path)
                 ))
+            except FileNotFoundError:
+                print(f'No labels file {labels_path} found')
 
 
 if __name__ == '__main__':

tests/imagelabel_visualize_test.py

@@ -0,0 +1,25 @@
+from pathlib import Path
+import cv2
+import os.path
+
+from dto.annotationClass import AnnotationClass
+from dto.imageLabel import ImageLabel
+from preprocessing import read_labels
+
+images_dir = '../images'
+labels_dir = '../labels'
+annotation_classes = AnnotationClass.read_json()
+
+
+image = os.listdir(images_dir)[0]
+image_path = os.path.join(images_dir, image)
+labels_path = os.path.join(labels_dir, f'{Path(image_path).stem}.txt')
+
+img = ImageLabel(
+    image_path=image_path,
+    image=cv2.imread(image_path),
+    labels_path=labels_path,
+    labels=read_labels(labels_path)
+)
+
+img.visualize(annotation_classes)

train.py

@@ -0,0 +1,4 @@
+import os
+
+
+current_dataset_dir = os.path.join('datasets', 'zombobase-current')