proper augmentations. Add 7 more images per 1 image

make dirs if not exists
add visualize

AnnotationClass.py

@@ -0,0 +1,21 @@
+import json
+
+
+class AnnotationClass:
+    def __init__(self, id, name, color):
+        self.id = id
+        self.name = name
+        self.color = color
+
+    @staticmethod
+    def read_json():
+        with open('classes.json', '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}
+
+    @property
+    def color_tuple(self):
+        color = self.color[3:]
+        lv = len(color)
+        xx = range(0, lv, lv // 3)
+        return tuple(int(color[i:i + lv // 3], 16) for i in xx)

preprocessing.py

@@ -1,38 +1,67 @@
 import os.path
 import time
-from array import *
 from pathlib import Path
 import albumentations as A
 import cv2
-
+from matplotlib import pyplot as plt
+from AnnotationClass import AnnotationClass
 
 labels_dir = 'labels'
 images_dir = 'images'
 current_dataset_dir = os.path.join('datasets', 'zombobase-current')
 current_images_dir = os.path.join(current_dataset_dir, 'images')
 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)], bbox_params=A.BboxParams(format='yolo')),
-        A.Compose([A.RandomBrightnessContrast(always_apply=True)], bbox_params=A.BboxParams(format='yolo')),
-        A.Compose([A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.2, rotate_limit=-40, always_apply=True)],
+        A.Compose([A.HorizontalFlip(always_apply=True)],
                   bbox_params=A.BboxParams(format='yolo')),
-        A.Compose([A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.2, rotate_limit=-20, always_apply=True)],
+        A.Compose([A.RandomBrightnessContrast(always_apply=True)],
                   bbox_params=A.BboxParams(format='yolo')),
-        A.Compose([A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.2, rotate_limit=20, always_apply=True)],
+        A.Compose([A.SafeRotate(limit=90, always_apply=True)],
                   bbox_params=A.BboxParams(format='yolo')),
-        A.Compose([A.ShiftScaleRotate(shift_limit=0.1, scale_limit=0.2, rotate_limit=40, always_apply=True)],
+        A.Compose([A.SafeRotate(limit=90, always_apply=True),
+                   A.RandomBrightnessContrast(always_apply=True)],
+                  bbox_params=A.BboxParams(format='yolo')),
+        A.Compose([A.ShiftScaleRotate(scale_limit=0.2, always_apply=True),
+                   A.VerticalFlip(always_apply=True),],
+                  bbox_params=A.BboxParams(format='yolo')),
+        A.Compose([A.ShiftScaleRotate(scale_limit=0.2, always_apply=True)],
+                  bbox_params=A.BboxParams(format='yolo')),
+        A.Compose([A.SafeRotate(limit=90, always_apply=True),
+                   A.RandomBrightnessContrast(always_apply=True)],
                   bbox_params=A.BboxParams(format='yolo'))
     ]
 
@@ -41,16 +70,23 @@ def image_processing(img_ann: ImageLabel) -> [ImageLabel]:
         res = transform(image=img_ann.image, bboxes=img_ann.labels)
         path = Path(img_ann.image_path)
         name = f'{path.stem}_{i+1}'
-        results.append(ImageLabel(
+        img = ImageLabel(
             image=res['image'],
             labels=res['bboxes'],
             image_path=os.path.join(current_images_dir, f'{name}{path.suffix}'),
             labels_path=os.path.join(current_labels_dir, f'{name}.txt')
-        ))
+        )
+        results.append(img)
     return results
 
 
-def write_result(img_ann: ImageLabel):
+def write_result(img_ann: ImageLabel, show_image=False):
+    os.makedirs(os.path.dirname(img_ann.image_path), exist_ok=True)
+    os.makedirs(os.path.dirname(img_ann.labels_path), exist_ok=True)
+
+    if show_image:
+        img_ann.visualize()
+
     cv2.imwrite(img_ann.image_path, img_ann.image)
     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]
@@ -75,7 +111,6 @@ def process_image(img_ann):
     results = image_processing(img_ann)
     for res_ann in results:
         write_result(res_ann)
-
     write_result(ImageLabel(
         image=img_ann.image,
         labels=img_ann.labels,
@@ -105,4 +140,4 @@ def main():
 
 
 if __name__ == '__main__':
-    main()
+    main()