refactor augmentation to class, update classes.json, fix small bugs

augmentation.py

@@ -0,0 +1,160 @@
+import concurrent.futures
+import os.path
+import shutil
+import time
+from datetime import datetime
+from pathlib import Path
+
+import albumentations as A
+import cv2
+import numpy as np
+
+from constants import (data_images_dir, data_labels_dir, processed_images_dir, processed_labels_dir, processed_dir)
+from dto.imageLabel import ImageLabel
+
+
+class Augmentator:
+    def __init__(self):
+        self.total_files_processed = 0
+        self.total_images_to_process = 0
+
+        self.correct_margin = 0.0005
+        self.correct_min_bbox_size = 0.01
+
+        self.transform = A.Compose([
+            # Flips, rotations and brightness
+            A.HorizontalFlip(p=0.6),
+            A.RandomBrightnessContrast(p=0.4, brightness_limit=(-0.1, 0.1), contrast_limit=(-0.1, 0.1)),
+            A.Affine(p=0.7, scale=(0.8, 1.2), rotate=(-20, 20), shear=(-10, 10), translate_percent=0.2),
+
+            # Weather
+            A.RandomFog(p=0.3, fog_coef_range=(0, 0.3)),
+            A.RandomShadow(p=0.2),
+
+            # Image Quality/Noise
+            A.MotionBlur(p=0.2, blur_limit=(3, 5)),
+
+            # Color Variations
+            A.HueSaturationValue(p=0.4, hue_shift_limit=10, sat_shift_limit=10, val_shift_limit=10)
+        ], bbox_params=A.BboxParams(format='yolo'))
+
+    def correct_bboxes(self, labels):
+        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 - self.correct_margin) - (x + half_width), (x - half_width) - self.correct_margin, 0)
+            w = bboxes[2] + 2*w_diff
+            if w < self.correct_min_bbox_size:
+                continue
+            h_diff = min((1 - self.correct_margin) - (y + half_height), ((y - half_height) - self.correct_margin), 0)
+            h = bboxes[3] + 2 * h_diff
+            if h < self.correct_min_bbox_size:
+                continue
+            res.append([x, y, w, h, bboxes[4]])
+        return res
+        pass
+
+    def augment_inner(self, img_ann: ImageLabel) -> [ImageLabel]:
+        results = []
+        labels = self.correct_bboxes(img_ann.labels)
+        if len(labels) == 0 and len(img_ann.labels) != 0:
+            print('no labels but was!!!')
+        results.append(ImageLabel(
+                image=img_ann.image,
+                labels=img_ann.labels,
+                image_path=os.path.join(processed_images_dir, Path(img_ann.image_path).name),
+                labels_path=os.path.join(processed_labels_dir, Path(img_ann.labels_path).name)
+            )
+        )
+        for i in range(7):
+            try:
+                res = self.transform(image=img_ann.image, bboxes=labels)
+                path = Path(img_ann.image_path)
+                name = f'{path.stem}_{i + 1}'
+                img = ImageLabel(
+                    image=res['image'],
+                    labels=res['bboxes'],
+                    image_path=os.path.join(processed_images_dir, f'{name}{path.suffix}'),
+                    labels_path=os.path.join(processed_labels_dir, f'{name}.txt')
+                )
+                results.append(img)
+            except Exception as e:
+                print(f'Error during transformation: {e}')
+        return results
+
+    def read_labels(self, labels_path) -> [[]]:
+        with open(labels_path, 'r') as f:
+            rows = f.readlines()
+            arr = []
+            for row in rows:
+                str_coordinates = row.split(' ')
+                class_num = str_coordinates.pop(0)
+                coordinates = [float(n.replace(',', '.')) for n in str_coordinates]
+                # noinspection PyTypeChecker
+                coordinates.append(class_num)
+                arr.append(coordinates)
+            return arr
+
+    def augment_annotation(self, image_file):
+        try:
+            image_path = os.path.join(data_images_dir, image_file.name)
+            labels_path = os.path.join(data_labels_dir, f'{Path(str(image_path)).stem}.txt')
+            image = cv2.imdecode(np.fromfile(image_path, dtype=np.uint8), cv2.IMREAD_UNCHANGED)
+
+            img_ann = ImageLabel(
+                image_path=image_path,
+                image=image,
+                labels_path=labels_path,
+                labels=self.read_labels(labels_path)
+            )
+            try:
+                results = self.augment_inner(img_ann)
+                for annotation in results:
+                    cv2.imencode('.jpg', annotation.image)[1].tofile(annotation.image_path)
+                    with open(annotation.labels_path, 'w') as f:
+                        lines = [f'{l[4]} {round(l[0], 5)} {round(l[1], 5)} {round(l[2], 5)} {round(l[3], 5)}\n' for l in
+                                 annotation.labels]
+                        f.writelines(lines)
+                        f.close()
+
+                print(f'{datetime.now():{'%Y-%m-%d %H:%M:%S'}}: {self.total_files_processed + 1}/{self.total_images_to_process} : {image_file.name} has augmented')
+            except Exception as e:
+                print(e)
+            self.total_files_processed += 1
+        except Exception as e:
+            print(f'Error appeared in thread for {image_file.name}: {e}')
+
+    def augment_annotations(self, from_scratch=False):
+        self.total_files_processed = 0
+
+        if from_scratch:
+            shutil.rmtree(processed_dir)
+
+        os.makedirs(processed_images_dir, exist_ok=True)
+        os.makedirs(processed_labels_dir, exist_ok=True)
+
+
+        processed_images = set(f.name for f in os.scandir(processed_images_dir))
+        images = []
+        with os.scandir(data_images_dir) as imd:
+            for image_file in imd:
+                if image_file.is_file() and image_file.name not in processed_images:
+                    images.append(image_file)
+        self.total_images_to_process = len(images)
+
+        with concurrent.futures.ThreadPoolExecutor() as executor:
+            executor.map(self.augment_annotation, images)
+
+
+if __name__ == '__main__':
+    augmentator = Augmentator()
+    while True:
+        augmentator.augment_annotations()
+        print('All processed, waiting for 2 minutes...')
+        time.sleep(300)

classes.json

@@ -1,18 +1,18 @@
 [
-    { "Id": 0, "Name": "Armored-Vehicle", "Color": "#80FF0000" },
+      { "Id": 0, "Name": "ArmorVehicle", "ShortName": "Броня", "Color": "#FF0000" },
-    { "Id": 1, "Name": "Truck", "Color": "#8000FF00" },
+      { "Id": 1, "Name": "Truck", "ShortName": "Вантаж.", "Color": "#00FF00" },
-    { "Id": 2, "Name": "Vehicle", "Color": "#800000FF" },
+      { "Id": 2, "Name": "Vehicle", "ShortName": "Машина", "Color": "#0000FF" },
-    { "Id": 3, "Name": "Artillery", "Color": "#80FFFF00" },
+      { "Id": 3, "Name": "Atillery", "ShortName": "Арта", "Color": "#FFFF00" },
-    { "Id": 4, "Name": "Shadow", "Color": "#80FF00FF" },
+      { "Id": 4, "Name": "Shadow", "ShortName": "Тінь", "Color": "#FF00FF" },
-    { "Id": 5, "Name": "Trenches", "Color": "#8000FFFF" },
+      { "Id": 5, "Name": "Trenches", "ShortName": "Окопи", "Color": "#00FFFF" },
-    { "Id": 6, "Name": "Military-men", "Color": "#80188021" },
+      { "Id": 6, "Name": "MilitaryMan", "ShortName": "Військов", "Color": "#188021" },
-    { "Id": 7, "Name": "Tyre-tracks", "Color": "#80800000" },
+      { "Id": 7, "Name": "TyreTracks", "ShortName": "Накати", "Color": "#800000" },
-    { "Id": 8, "Name": "Additional-armored-tank", "Color": "#80008000" },
+      { "Id": 8, "Name": "AdditArmoredTank", "ShortName": "Танк.захист", "Color": "#008000" },
-    { "Id": 9, "Name": "Smoke", "Color": "#80000080" },
+      { "Id": 9, "Name": "Smoke", "ShortName": "Дим", "Color": "#000080" },
-    { "Id": 10, "Name": "Plane", "Color": "#80000080" },
+      { "Id": 10, "Name": "Plane", "ShortName": "Літак", "Color": "#000080" },
-    { "Id": 11, "Name": "Moto", "Color": "#80808000" },
+      { "Id": 11, "Name": "Moto", "ShortName": "Мото", "Color": "#808000" },
-    { "Id": 12, "Name": "Camouflage-net", "Color": "#80800080" },
+      { "Id": 12, "Name": "CamouflageNet", "ShortName": "Сітка", "Color": "#800080" },
-    { "Id": 13, "Name": "Camouflage-branches", "Color": "#80008080" },
+      { "Id": 13, "Name": "CamouflageBranches", "ShortName": "Гілки", "Color": "#2f4f4f" },
-    { "Id": 14, "Name": "Roof", "Color": "#80008080" },
+      { "Id": 14, "Name": "Roof", "ShortName": "Дах", "Color": "#1e90ff" },
-    { "Id": 15, "Name": "Building", "Color": "#80008080" }
+      { "Id": 15, "Name": "Building", "ShortName": "Будівля", "Color": "#ffb6c1" }
 ]

preprocess-train.py

@@ -1,6 +1,5 @@
-from preprocessing import preprocess_annotations
+from augmentation import Augmentator
 from train import train_dataset, convert2rknn
 
-preprocess_annotations()
+Augmentator().augment_annotations(from_scratch=True)
-train_dataset('2024-10-01')
+train_dataset(from_scratch=True)
-convert2rknn()

preprocessing.py

@@ -1,162 +0,0 @@
-import concurrent.futures
-import os.path
-import time
-from pathlib import Path
-
-import albumentations as A
-import cv2
-import numpy as np
-
-from constants import (data_images_dir, data_labels_dir, processed_images_dir, processed_labels_dir)
-from dto.imageLabel import ImageLabel
-
-total_files_processed = 0
-transform = A.Compose([
-    # Flips, rotations and brightness
-    A.HorizontalFlip(),
-    A.RandomBrightnessContrast(brightness_limit=(-0.05, 0.05), contrast_limit=(-0.05, 0.05)),
-    A.Affine(p=0.7, scale=(0.8, 1.2), rotate=25, translate_percent=0.1),
-
-    # Weather
-    A.RandomFog(p=0.2, fog_coef_range=(0, 0.3)),
-    A.RandomShadow(p=0.2),
-
-    # Image Quality/Noise
-    A.MotionBlur(p=0.2, blur_limit=(3, 5)),
-
-    # Color Variations
-    A.HueSaturationValue(p=0.3, hue_shift_limit=8, sat_shift_limit=8, val_shift_limit=8)
-], bbox_params=A.BboxParams(format='yolo'))
-
-
-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]:
-    results = []
-    labels = correct_bboxes(img_ann.labels)
-    if len(labels) == 0 and len(img_ann.labels) != 0:
-        print('no labels but was!!!')
-    results.append(ImageLabel(
-            image=img_ann.image,
-            labels=img_ann.labels,
-            image_path=os.path.join(processed_images_dir, Path(img_ann.image_path).name),
-            labels_path=os.path.join(processed_labels_dir, Path(img_ann.labels_path).name)
-        )
-    )
-    for i in range(7):
-        try:
-            res = transform(image=img_ann.image, bboxes=labels)
-            path = Path(img_ann.image_path)
-            name = f'{path.stem}_{i + 1}'
-            img = ImageLabel(
-                image=res['image'],
-                labels=res['bboxes'],
-                image_path=os.path.join(processed_images_dir, f'{name}{path.suffix}'),
-                labels_path=os.path.join(processed_labels_dir, f'{name}.txt')
-            )
-            results.append(img)
-        except Exception as e:
-            print(f'Error during transformation: {e}')
-    return results
-
-
-def write_result(img_ann: ImageLabel):
-    cv2.imencode('.jpg', img_ann.image)[1].tofile(img_ann.image_path)
-    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()
-    global total_files_processed
-    print(f'{total_files_processed}. {img_ann.labels_path} written')
-
-
-def read_labels(labels_path) -> [[]]:
-    with open(labels_path, 'r') as f:
-        rows = f.readlines()
-        arr = []
-        for row in rows:
-            str_coordinates = row.split(' ')
-            class_num = str_coordinates.pop(0)
-            coordinates = [float(n.replace(',', '.')) for n in str_coordinates]
-            # noinspection PyTypeChecker
-            coordinates.append(class_num)
-            arr.append(coordinates)
-        return arr
-
-
-def preprocess_annotations():
-    global total_files_processed  # Indicate that we're using the global counter
-    total_files_processed = 0
-
-    os.makedirs(processed_images_dir, exist_ok=True)
-    os.makedirs(processed_labels_dir, exist_ok=True)
-
-    processed_images = set(f.name for f in os.scandir(processed_images_dir))
-    images = []
-    with os.scandir(data_images_dir) as imd:
-        for image_file in imd:
-            if image_file.is_file() and image_file.name not in processed_images:
-                images.append(image_file)
-    with concurrent.futures.ThreadPoolExecutor() as executor:
-        executor.map(process_image_file, images)
-
-
-def process_image_file(image_file):
-    try:
-        image_path = os.path.join(data_images_dir, image_file.name)
-        labels_path = os.path.join(data_labels_dir, f'{Path(str(image_path)).stem}.txt')
-        image = cv2.imdecode(np.fromfile(image_path, dtype=np.uint8), cv2.IMREAD_UNCHANGED)
-
-        img_ann = ImageLabel(
-            image_path=image_path,
-            image=image,
-            labels_path=labels_path,
-            labels=read_labels(labels_path)
-        )
-        try:
-            results = image_processing(img_ann)
-            for res_ann in results:
-                write_result(res_ann)
-        except Exception as e:
-            print(e)
-        global total_files_processed
-        total_files_processed += 1
-    except Exception as e:
-        print(f'Error appeared in thread for {image_file.name}: {e}')
-
-
-def main():
-    while True:
-        preprocess_annotations()
-        print('All processed, waiting for 2 minutes...')
-        time.sleep(300)
-
-
-if __name__ == '__main__':
-    main()

train.py

@@ -38,6 +38,7 @@ DEFAULT_CLASS_NUM = 80
 total_files_copied = 0
 
 def form_dataset(from_date: datetime):
+
     makedirs(today_dataset, exist_ok=True)
     images = []
     old_images = []
@@ -67,7 +68,6 @@ def form_dataset(from_date: datetime):
     copy_annotations(images[:train_size], 'train')
     copy_annotations(images[train_size:train_size + valid_size], 'valid')
     copy_annotations(images[train_size + valid_size:], 'test')
-    create_yaml()
 
 
 def copy_annotations(images, folder):
@@ -174,21 +174,22 @@ def get_latest_model():
 
 
 def train_dataset(existing_date=None, from_scratch=False):
-    latest_date, latest_model = get_latest_model()
+    latest_date, latest_model = get_latest_model() if not from_scratch else None, None
 
     if existing_date is not None:
         cur_folder = f'{prefix}{existing_date}'
         cur_dataset = path.join(datasets_dir, f'{prefix}{existing_date}')
     else:
-        # form_dataset(latest_date)
+        if from_scratch:
-        # create_yaml()
+            shutil.rmtree(today_dataset)
+        form_dataset(latest_date)
+        create_yaml()
         cur_folder = today_folder
         cur_dataset = today_dataset
 
     model_name = latest_model if latest_model is not None and path.isfile(latest_model) and not from_scratch else 'yolo11m.yaml'
     print(f'Initial model: {model_name}')
     model = YOLO(model_name)
-    model.info['author'] = 'LLC Azaion'
 
     yaml = abspath(path.join(cur_dataset, 'data.yaml'))
     results = model.train(data=yaml,
@@ -222,7 +223,7 @@ def validate(model_path):
 
 
 if __name__ == '__main__':
-    # model_path = train_dataset(from_scratch=True)
+    model_path = train_dataset(from_scratch=True)
     # validate(path.join('runs', 'detect', 'train7', 'weights', 'best.pt'))
     # form_data_sample(500)
     # convert2rknn()