fix augmentation - correct bboxes to be within borders+margin

Add classes = 80 as default number in yaml file

preprocessing.py

@@ -11,10 +11,35 @@ from constants import (data_images_dir, data_labels_dir, processed_images_dir, p
 from dto.imageLabel import ImageLabel
 
 
+def correct_bboxes(labels):
+    margin = 0.0005
+    min_size = 0.01
+    res = []
+    for bboxes in labels:
+        x = bboxes[0]
+        y = bboxes[1]
+        half_width = 0.5*bboxes[2]
+        half_height = 0.5*bboxes[3]
+
+        # calc how much bboxes are outside borders ( +small margin ).
+        # value should be negative. If it's positive, then put 0, as no correction
+        w_diff = min( (1 - margin) - (x + half_width), (x - half_width) - margin, 0 )
+        w = bboxes[2] + 2*w_diff
+        if w < min_size:
+            continue
+        h_diff = min( (1 - margin) - (y + half_height), ((y - half_height) - margin), 0)
+        h = bboxes[3] + 2 * h_diff
+        if h < min_size:
+            continue
+        res.append([x, y, w, h, bboxes[4]])
+    return res
+    pass
+
+
 def image_processing(img_ann: ImageLabel) -> [ImageLabel]:
     transforms = [
         A.Compose([A.HorizontalFlip(always_apply=True)],
-                  bbox_params=A.BboxParams(format='yolo')),
+                  bbox_params=A.BboxParams(format='yolo', )),
         A.Compose([A.RandomBrightnessContrast(always_apply=True)],
                   bbox_params=A.BboxParams(format='yolo')),
         A.Compose([A.SafeRotate(limit=90, always_apply=True)],
@@ -33,9 +58,12 @@ def image_processing(img_ann: ImageLabel) -> [ImageLabel]:
     ]
 
     results = []
+    labels = correct_bboxes(img_ann.labels)
+    if len(labels) == 0 and len(img_ann.labels) != 0:
+        print('no labels but was!!!')
     for i, transform in enumerate(transforms):
         try:
-            res = transform(image=img_ann.image, bboxes=img_ann.labels)
+            res = transform(image=img_ann.image, bboxes=labels)
             path = Path(img_ann.image_path)
             name = f'{path.stem}_{i + 1}'
             img = ImageLabel(
@@ -51,9 +79,6 @@ def image_processing(img_ann: ImageLabel) -> [ImageLabel]:
 
 
 def write_result(img_ann: ImageLabel):
-    os.makedirs(os.path.dirname(img_ann.image_path), exist_ok=True)
-    os.makedirs(os.path.dirname(img_ann.labels_path), exist_ok=True)
-
     cv2.imencode('.jpg', img_ann.image)[1].tofile(img_ann.image_path)
     print(f'{img_ann.image_path} written')
 
@@ -92,6 +117,8 @@ def process_image(img_ann):
 
 
 def main():
+    os.makedirs(processed_images_dir, exist_ok=True)
+    os.makedirs(processed_labels_dir, exist_ok=True)
     while True:
         processed_images = set(f.name for f in os.scandir(processed_images_dir))
         images = []