diff --git a/constants.py b/constants.py new file mode 100644 index 0000000..1f790ff --- /dev/null +++ b/constants.py @@ -0,0 +1,13 @@ +import os +from datetime import datetime + +from dto.annotationClass import AnnotationClass + +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() + + +prefix = 'zombobase-' +today_dataset = os.path.join('datasets', f'{prefix}{datetime.now():%Y-%m-%d}') \ No newline at end of file diff --git a/preprocessing.py b/preprocessing.py index f476a57..c5bd525 100644 --- a/preprocessing.py +++ b/preprocessing.py @@ -3,15 +3,11 @@ import time from pathlib import Path import albumentations as A import cv2 -from dto.annotationClass import AnnotationClass +from constants import current_images_dir, current_labels_dir, annotation_classes from dto.imageLabel import ImageLabel 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() def image_processing(img_ann: ImageLabel) -> [ImageLabel]: @@ -37,16 +33,19 @@ def image_processing(img_ann: ImageLabel) -> [ImageLabel]: results = [] for i, transform in enumerate(transforms): - res = transform(image=img_ann.image, bboxes=img_ann.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(current_images_dir, f'{name}{path.suffix}'), - labels_path=os.path.join(current_labels_dir, f'{name}.txt') - ) - results.append(img) + try: + res = transform(image=img_ann.image, bboxes=img_ann.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(current_images_dir, f'{name}{path.suffix}'), + labels_path=os.path.join(current_labels_dir, f'{name}.txt') + ) + results.append(img) + except Exception as e: + print(f'Error during transformtation: {e}') return results @@ -74,7 +73,7 @@ def read_labels(labels_path) -> [[]]: for row in rows: str_coordinates = row.split(' ') class_num = str_coordinates.pop(0) - coordinates = [float(n) for n in str_coordinates] + coordinates = [float(n.replace(',', '.')) for n in str_coordinates] coordinates.append(class_num) arr.append(coordinates) return arr @@ -111,8 +110,18 @@ def main(): labels_path=labels_path, labels=read_labels(labels_path) )) - except FileNotFoundError: - print(f'No labels file {labels_path} found') + except Exception as e: + print(f'Error appeared {e}') + + try: + os.remove(image_path) + except OSError: + pass + + try: + os.remove(labels_path) + except OSError: + pass if __name__ == '__main__': diff --git a/tests/imagelabel_visualize_test.py b/tests/imagelabel_visualize_test.py index 67c7890..c101ad5 100644 --- a/tests/imagelabel_visualize_test.py +++ b/tests/imagelabel_visualize_test.py @@ -1,19 +1,16 @@ -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() +images_dir = '' -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') +image_path = 'test01.jpg' +labels_path = 'test01.txt' img = ImageLabel( image_path=image_path, diff --git a/tests/test01.jpg b/tests/test01.jpg new file mode 100644 index 0000000..099192f Binary files /dev/null and b/tests/test01.jpg differ diff --git a/tests/test01.txt b/tests/test01.txt new file mode 100644 index 0000000..feade35 --- /dev/null +++ b/tests/test01.txt @@ -0,0 +1 @@ +0 0.3809 0.49269 0.21636 0.39129 diff --git a/train.py b/train.py index 8105dbd..787c3c7 100644 --- a/train.py +++ b/train.py @@ -1,4 +1,86 @@ import os +import shutil +from datetime import datetime +from pathlib import Path + +from ultralytics import YOLO +from constants import current_images_dir, current_labels_dir, annotation_classes, today_dataset, prefix + +yaml_name = 'data.yaml' +yaml_path = os.path.join(today_dataset, yaml_name) +train_set = 70 +valid_set = 20 +test_set = 10 -current_dataset_dir = os.path.join('datasets', 'zombobase-current') +def form_dataset(): + os.makedirs(today_dataset, exist_ok=True) + images = os.listdir(current_images_dir) + + train_size = int(len(images) * train_set / 100.0) + valid_size = int(len(images) * valid_set / 100.0) + + move_annotations(images[:train_size], 'train') + move_annotations(images[train_size:train_size + valid_size], 'valid') + move_annotations(images[train_size + valid_size:], 'test') + + create_yaml() + + +def move_annotations(images, folder): + destination_images = os.path.join(today_dataset, folder, 'images') + os.makedirs(destination_images, exist_ok=True) + destination_labels = os.path.join(today_dataset, folder, 'labels') + os.makedirs(destination_labels, exist_ok=True) + for image_name in images: + image_path = os.path.join(current_images_dir, image_name) + label_name = f'{Path(image_name).stem}.txt' + label_path = os.path.join(current_labels_dir, label_name) + os.replace(image_path, os.path.join(destination_images, image_name)) + os.replace(label_path, os.path.join(destination_labels, label_name)) + + +def create_yaml(): + lines = ['names:'] + for c in annotation_classes: + lines.append(f'- {annotation_classes[c].name}') + lines.append(f'nc: {len(annotation_classes)}') + lines.append(f'test: test/images') + lines.append(f'train: train/images') + lines.append(f'val: valid/images') + lines.append('') + + with open(yaml_path, 'w', encoding='utf-8') as f: + f.writelines([f'{line}\n' for line in lines]) + + +def get_recent_model(): + date_sets = [] + datasets = [next((file for file in os.listdir(os.path.join('datasets', d)) if file.endswith('pt')), None) + for d in os.listdir('datasets')] + + # date_str = d.replace(prefix, '') + # if date_str == 'current' or date_str == f'{datetime.now():%Y-%m-%d}': + # continue + # if len(date_sets) == 0: + # return None + + recent = max(date_sets) + return os.path.join('datasets', f'{prefix}{recent}', f'{prefix}{recent}.pt') + + +def retrain(): + model = YOLO(get_recent_model() or 'yolov10x.yaml') + model.train(data=yaml_path, save=True, cache=True) + + +def revert_to_current(date): + def revert_dir(dir): + os.listdir(os.path.join(current_images_dir, 'images')) + + date_dataset = f'{prefix}{date}' + revert_dir(os.path.join(date_dataset, 'test')) + +form_dataset() +create_yaml() +retrain() diff --git a/yolov10x.yaml b/yolov10x.yaml new file mode 100644 index 0000000..539af6b --- /dev/null +++ b/yolov10x.yaml @@ -0,0 +1,40 @@ +# Parameters +nc: 50 # number of classes +scales: # model compound scaling constants, i.e. 'model=yolov8n.yaml' will call yolov8.yaml with scale 'n' + # [depth, width, max_channels] + x: [1.00, 1.25, 512] + +# YOLOv8.0n backbone +backbone: + # [from, repeats, module, args] + - [-1, 1, Conv, [64, 3, 2]] # 0-P1/2 + - [-1, 1, Conv, [128, 3, 2]] # 1-P2/4 + - [-1, 3, C2f, [128, True]] + - [-1, 1, Conv, [256, 3, 2]] # 3-P3/8 + - [-1, 6, C2f, [256, True]] + - [-1, 1, SCDown, [512, 3, 2]] # 5-P4/16 + - [-1, 6, C2fCIB, [512, True]] + - [-1, 1, SCDown, [1024, 3, 2]] # 7-P5/32 + - [-1, 3, C2fCIB, [1024, True]] + - [-1, 1, SPPF, [1024, 5]] # 9 + - [-1, 1, PSA, [1024]] # 10 + +# YOLOv8.0n head +head: + - [-1, 1, nn.Upsample, [None, 2, "nearest"]] + - [[-1, 6], 1, Concat, [1]] # cat backbone P4 + - [-1, 3, C2fCIB, [512, True]] # 13 + + - [-1, 1, nn.Upsample, [None, 2, "nearest"]] + - [[-1, 4], 1, Concat, [1]] # cat backbone P3 + - [-1, 3, C2f, [256]] # 16 (P3/8-small) + + - [-1, 1, Conv, [256, 3, 2]] + - [[-1, 13], 1, Concat, [1]] # cat head P4 + - [-1, 3, C2fCIB, [512, True]] # 19 (P4/16-medium) + + - [-1, 1, SCDown, [512, 3, 2]] + - [[-1, 10], 1, Concat, [1]] # cat head P5 + - [-1, 3, C2fCIB, [1024, True]] # 22 (P5/32-large) + + - [[16, 19, 22], 1, v10Detect, [nc]] # Detect(P3, P4, P5)