From 292809593c028f6e8a16570eb049143d660af2f0 Mon Sep 17 00:00:00 2001 From: zxsanny Date: Thu, 15 May 2025 06:10:10 +0300 Subject: [PATCH] refactor augmentation to class, update classes.json, fix small bugs --- augmentation.py | 160 +++++++++++++++++++++++++++++++++++++++++++ classes.json | 36 +++++----- preprocess-train.py | 7 +- preprocessing.py | 162 -------------------------------------------- train.py | 13 ++-- 5 files changed, 188 insertions(+), 190 deletions(-) create mode 100644 augmentation.py delete mode 100644 preprocessing.py diff --git a/augmentation.py b/augmentation.py new file mode 100644 index 0000000..fd80988 --- /dev/null +++ b/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) diff --git a/classes.json b/classes.json index eff5acc..b6e1ac5 100644 --- a/classes.json +++ b/classes.json @@ -1,18 +1,18 @@ - [ - { "Id": 0, "Name": "Armored-Vehicle", "Color": "#80FF0000" }, - { "Id": 1, "Name": "Truck", "Color": "#8000FF00" }, - { "Id": 2, "Name": "Vehicle", "Color": "#800000FF" }, - { "Id": 3, "Name": "Artillery", "Color": "#80FFFF00" }, - { "Id": 4, "Name": "Shadow", "Color": "#80FF00FF" }, - { "Id": 5, "Name": "Trenches", "Color": "#8000FFFF" }, - { "Id": 6, "Name": "Military-men", "Color": "#80188021" }, - { "Id": 7, "Name": "Tyre-tracks", "Color": "#80800000" }, - { "Id": 8, "Name": "Additional-armored-tank", "Color": "#80008000" }, - { "Id": 9, "Name": "Smoke", "Color": "#80000080" }, - { "Id": 10, "Name": "Plane", "Color": "#80000080" }, - { "Id": 11, "Name": "Moto", "Color": "#80808000" }, - { "Id": 12, "Name": "Camouflage-net", "Color": "#80800080" }, - { "Id": 13, "Name": "Camouflage-branches", "Color": "#80008080" }, - { "Id": 14, "Name": "Roof", "Color": "#80008080" }, - { "Id": 15, "Name": "Building", "Color": "#80008080" } - ] +[ + { "Id": 0, "Name": "ArmorVehicle", "ShortName": "Броня", "Color": "#FF0000" }, + { "Id": 1, "Name": "Truck", "ShortName": "Вантаж.", "Color": "#00FF00" }, + { "Id": 2, "Name": "Vehicle", "ShortName": "Машина", "Color": "#0000FF" }, + { "Id": 3, "Name": "Atillery", "ShortName": "Арта", "Color": "#FFFF00" }, + { "Id": 4, "Name": "Shadow", "ShortName": "Тінь", "Color": "#FF00FF" }, + { "Id": 5, "Name": "Trenches", "ShortName": "Окопи", "Color": "#00FFFF" }, + { "Id": 6, "Name": "MilitaryMan", "ShortName": "Військов", "Color": "#188021" }, + { "Id": 7, "Name": "TyreTracks", "ShortName": "Накати", "Color": "#800000" }, + { "Id": 8, "Name": "AdditArmoredTank", "ShortName": "Танк.захист", "Color": "#008000" }, + { "Id": 9, "Name": "Smoke", "ShortName": "Дим", "Color": "#000080" }, + { "Id": 10, "Name": "Plane", "ShortName": "Літак", "Color": "#000080" }, + { "Id": 11, "Name": "Moto", "ShortName": "Мото", "Color": "#808000" }, + { "Id": 12, "Name": "CamouflageNet", "ShortName": "Сітка", "Color": "#800080" }, + { "Id": 13, "Name": "CamouflageBranches", "ShortName": "Гілки", "Color": "#2f4f4f" }, + { "Id": 14, "Name": "Roof", "ShortName": "Дах", "Color": "#1e90ff" }, + { "Id": 15, "Name": "Building", "ShortName": "Будівля", "Color": "#ffb6c1" } +] \ No newline at end of file diff --git a/preprocess-train.py b/preprocess-train.py index d23fd0d..a85dc6d 100644 --- a/preprocess-train.py +++ b/preprocess-train.py @@ -1,6 +1,5 @@ -from preprocessing import preprocess_annotations +from augmentation import Augmentator from train import train_dataset, convert2rknn -preprocess_annotations() -train_dataset('2024-10-01') -convert2rknn() \ No newline at end of file +Augmentator().augment_annotations(from_scratch=True) +train_dataset(from_scratch=True) \ No newline at end of file diff --git a/preprocessing.py b/preprocessing.py deleted file mode 100644 index 1a104a9..0000000 --- a/preprocessing.py +++ /dev/null @@ -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() diff --git a/train.py b/train.py index 16c5023..efaeab3 100644 --- a/train.py +++ b/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) - # create_yaml() + if from_scratch: + 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()