Merge remote-tracking branch 'origin/main'

# Conflicts:
#	requirements.txt

orangepi5/00 install-os.md

@@ -1,5 +0,0 @@
-1. Download latest release from here https://joshua-riek.github.io/ubuntu-rockchip-download/boards/orangepi-5.html
-   f.e. https://github.com/Joshua-Riek/ubuntu-rockchip/releases/download/v2.3.2/ubuntu-22.04-preinstalled-desktop-arm64-orangepi-5.img.xz
-   but look to the more recent version on ubuntu 22.04
-   
-2. Write the image to the microsd using https://bztsrc.gitlab.io/usbimager/ (sudo ./usbimager on linux) (or use BalenaEtcher)

orangepi5/convert-install.sh

@@ -1,36 +0,0 @@
-mkdir rknn-convert
-cd rknn-convert
-
-# Install converter PT to ONNX
-git clone https://github.com/airockchip/ultralytics_yolov8
-cd ultralytics_yolov8
-sudo apt install python3.12-venv
-python3 -m venv env
-source env/bin/activate
-pip install .
-pip install onnx
-cp ultralytics/cfg/default.yaml ultralytics/cfg/default_backup.yaml
-sed -i -E "s/(model: ).+( #.+)/\1azaion.pt\2/" ultralytics/cfg/default.yaml
-cd ..
-deactivate
-
-# Install converter ONNX to RKNN
-wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O miniconda.sh
-chmod +x miniconda.sh
-bash miniconda.sh -b -p $HOME/miniconda
-
-source ~/miniconda/bin/activate
-conda create -n toolkit2 -y python=3.11
-conda activate toolkit2
-git clone https://github.com/rockchip-linux/rknn-toolkit2.git
-cd rknn-toolkit2/rknn-toolkit2/packages
-pip install -r requirements_cp311-1.6.0.txt
-pip install rknn_toolkit2-1.6.0+81f21f4d-cp311-cp311-linux_x86_64.whl
-pip install "numpy<2.0"
-cd ../../../
-git clone https://github.com/airockchip/rknn_model_zoo.git
-
-sed -i -E "s#(DATASET_PATH = ').+(')#\1/azaion/data-sample/azaion_subset.txt\2 #" rknn_model_zoo/examples/yolov8/python/convert.py
-
-conda deactivate
-conda deactivate

orangepi5/convert.sh

@@ -1,19 +0,0 @@
-# PT to ONNX
-cd rknn-convert/ultralytics_yolov8/
-cp --verbose /azaion/models/azaion.pt .
-source env/bin/activate
-pip install onnx
-export PYTHONPATH=./
-python ./ultralytics/engine/exporter.py
-cp --verbose azaion.onnx ../
-cd ..
-deactivate
-cp --verbose azaion.onnx /azaion/models/
-
-#  ONNX to RKNN
-source ~/miniconda/bin/activate
-conda activate toolkit2
-cd rknn_model_zoo/examples/yolov8/python
-python convert.py ../../../../azaion.onnx rk3588 i8 /azaion/models/azaion.rknn
-conda deactivate
-conda deactivate

preprocessing-cuda.py

@@ -0,0 +1,269 @@
+import os
+import time
+from datetime import datetime, timedelta
+from pathlib import Path
+import cv2
+import numpy as np
+import torch
+import kornia.augmentation as K
+import kornia.utils as KU
+from torch.utils.data import Dataset, DataLoader
+import concurrent.futures
+
+from constants import (data_images_dir, data_labels_dir, processed_images_dir, processed_labels_dir)
+from dto.imageLabel import ImageLabel
+
+# Configurable parameters
+num_augmented_images = 7
+augmentation_probability = 0.5 # general probability for augmentations, can be adjusted per augmentation
+RESIZE_SIZE = (1080, 1920) # Resize images to Full HD 1920x1080 (height, width)
+processed_images_dir = processed_images_dir + '_cuda'
+processed_labels_dir = processed_labels_dir + '_cuda'
+
+# Ensure directories exist
+os.makedirs(processed_images_dir, exist_ok=True)
+os.makedirs(processed_labels_dir, exist_ok=True)
+
+# Custom Augmentations
+class RandomFog(K.AugmentationBase2D):
+    def __init__(self, fog_coef_range=(0, 0.3), p=augmentation_probability, same_on_batch=True):
+        super().__init__(p=p, same_on_batch=same_on_batch)
+        self.fog_coef_range = fog_coef_range
+
+    def compute_transformation(self, input_shape: torch.Size, params: dict) -> dict:
+        return {"fog_factor": torch.rand(input_shape[0], device=self.device) * (self.fog_coef_range[1] - self.fog_coef_range[0]) + self.fog_coef_range[0]}
+
+    def apply_transform(self, input: torch.Tensor, params: dict, transform: dict) -> torch.Tensor:
+        fog_factor = transform['fog_factor'].view(-1, 1, 1, 1)
+        return input * (1.0 - fog_factor) + fog_factor
+
+class RandomShadow(K.AugmentationBase2D):
+    def __init__(self, shadow_factor_range=(0.2, 0.8), p=augmentation_probability, same_on_batch=True):
+        super().__init__(p=p, same_on_batch=same_on_batch)
+        self.shadow_factor_range = shadow_factor_range
+
+    def compute_transformation(self, input_shape: torch.Size, params: dict) -> dict:
+        batch_size, _, height, width = input_shape
+        x1 = torch.randint(0, width, (batch_size,), device=self.device)
+        y1 = torch.randint(0, height, (batch_size,), device=self.device)
+        x2 = torch.randint(x1, width, (batch_size,), device=self.device)
+        y2 = torch.randint(y1, height, (batch_size,), device=self.device)
+        shadow_factor = torch.rand(batch_size, device=self.device) * (self.shadow_factor_range[1] - self.shadow_factor_range[0]) + self.shadow_factor_range[0]
+        return {"x1": x1, "y1": y1, "x2": x2, "y2": y2, "shadow_factor": shadow_factor}
+
+    def apply_transform(self, input: torch.Tensor, params: dict, transform: dict) -> torch.Tensor:
+        batch_size, _, height, width = input.size()
+        mask = torch.zeros_like(input, device=self.device)
+        for b in range(batch_size):
+            mask[b, :, transform['y1'][b]:transform['y2'][b], transform['x1'][b]:transform['x2'][b]] = 1
+        shadow_factor = transform['shadow_factor'].view(-1, 1, 1, 1)
+        return input * (1.0 - mask) + input * mask * shadow_factor
+
+
+class ImageDataset(Dataset):
+    def __init__(self, images_dir, labels_dir):
+        self.images_dir = images_dir
+        self.labels_dir = labels_dir
+        self.image_filenames = [f for f in os.listdir(images_dir) if os.path.isfile(os.path.join(images_dir, f))]
+        self.resize = K.Resize(RESIZE_SIZE) # Add resize transform here
+
+    def __len__(self):
+        return len(self.image_filenames)
+
+    def __getitem__(self, idx):
+        image_filename = self.image_filenames[idx]
+        image_path = os.path.join(self.images_dir, image_filename)
+        label_path = os.path.join(self.labels_dir, Path(image_filename).stem + '.txt')
+
+        image_np = cv2.imread(image_path)
+        if image_np is None:
+            raise FileNotFoundError(f"Error reading image: {image_path}")
+        image_np = cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB) # Convert to RGB for Kornia
+        image = KU.image_to_tensor(image_np, keepdim=False).float()  # HWC -> CHW, and to tensor, convert to float here!
+        image = self.resize(image) # Resize image here to fixed size
+
+        print(f"Image shape after resize (index {idx}, filename {image_filename}): {image.shape}") # DEBUG PRINT
+
+        labels = []
+        if os.path.exists(label_path):
+            labels = self._read_labels(label_path)
+        return image, labels, image_filename
+
+    def _read_labels(self, labels_path):
+        labels = []
+        with open(labels_path, 'r') as f:
+            for row in f.readlines():
+                str_coordinates = row.strip().split(' ')
+                class_num = int(str_coordinates[0])
+                coordinates = [float(n) for n in str_coordinates[1:]] # x_center, y_center, width, height (normalized YOLO)
+                labels.append([*coordinates, class_num])
+        return labels
+
+def yolo_to_xyxy(bboxes_yolo, image_width, image_height):
+    bboxes_xyxy = []
+    for bbox in bboxes_yolo:
+        x_center, y_center, w, h, class_id = bbox
+        x_min = int((x_center - w / 2) * image_width)
+        y_min = int((y_center - h / 2) * image_height)
+        x_max = int((x_center + w / 2) * image_width)
+        y_max = int((y_center + h / 2) * image_height)
+        bboxes_xyxy.append([x_min, y_min, x_max, y_max, class_id])
+    return torch.tensor(bboxes_xyxy) if bboxes_xyxy else torch.empty((0, 5))
+
+def xyxy_to_yolo(bboxes_xyxy, image_width, image_height):
+    bboxes_yolo = []
+    for bbox in bboxes_xyxy:
+        x_min, y_min, x_max, y_max, class_id = bbox
+        x_center = ((x_min + x_max) / 2) / image_width
+        y_center = ((y_min + y_max) / 2) / image_height
+        w = (x_max - x_min) / image_width
+        h = (y_max - y_min) / image_height
+        bboxes_yolo.append([x_center, y_center, w, h, int(class_id)])
+    return bboxes_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]
+
+        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
+
+def process_image_and_labels(image, labels_yolo, image_filename, geometric_pipeline, intensity_pipeline, device):
+    image = image.float() / 255.0
+    original_height, original_width = RESIZE_SIZE[0], RESIZE_SIZE[1] # Use fixed resize size (Height, Width)
+
+    processed_image_labels = []
+
+    # 1. Original image and labels
+    current_labels_yolo_corrected = correct_bboxes(labels_yolo)
+    processed_image_labels.append(ImageLabel(
+        image=KU.tensor_to_image(image.byte()), # Convert back to numpy uint8 for saving
+        labels=current_labels_yolo_corrected,
+        image_path=os.path.join(processed_images_dir, image_filename),
+        labels_path=os.path.join(processed_labels_dir, Path(image_filename).stem + '.txt')
+    ))
+
+    # 2-8. Augmented images
+    for i in range(num_augmented_images):
+        img_batch = image.unsqueeze(0).to(device) # BCHW
+        bboxes_xyxy = yolo_to_xyxy(labels_yolo, original_width, original_height).unsqueeze(0).to(device) # B N 5
+
+        augmented_batch = geometric_pipeline(img_batch, params={"bbox": bboxes_xyxy})
+        geo_augmented_image = augmented_batch["input"]
+        geo_augmented_bboxes_xyxy = augmented_batch["bbox"]
+
+        intensity_augmented_image = intensity_pipeline(geo_augmented_image)
+
+        # Convert back to CPU and numpy
+        augmented_image_np = KU.tensor_to_image((intensity_augmented_image.squeeze(0).cpu() * 255.0).byte())
+        augmented_bboxes_xyxy_cpu = geo_augmented_bboxes_xyxy.squeeze(0).cpu()
+
+        augmented_bboxes_yolo = xyxy_to_yolo(augmented_bboxes_xyxy_cpu, original_width, original_height)
+        augmented_bboxes_yolo_corrected = correct_bboxes(augmented_bboxes_yolo)
+
+
+        processed_image_labels.append(ImageLabel(
+            image=augmented_image_np,
+            labels=augmented_bboxes_yolo_corrected,
+            image_path=os.path.join(processed_images_dir, f'{Path(image_filename).stem}_{i + 1}{Path(image_filename).suffix}'),
+            labels_path=os.path.join(processed_labels_dir, f'{Path(image_filename).stem}_{i + 1}.txt')
+        ))
+    return processed_image_labels
+
+def write_result(img_ann: ImageLabel):
+    cv2.imwrite(img_ann.image_path, cv2.cvtColor(img_ann.image, cv2.COLOR_RGB2BGR)) # Save as BGR
+    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()
+    print(f'{img_ann.labels_path} written')
+
+
+def process_batch_wrapper(batch_data, geometric_pipeline, intensity_pipeline, device):
+    processed_batch_image_labels = []
+    for image, labels_yolo, image_filename in batch_data:
+        results = process_image_and_labels(image, labels_yolo, image_filename, geometric_pipeline, intensity_pipeline, device)
+        processed_batch_image_labels.extend(results)
+    return processed_batch_image_labels
+
+def save_batch_results(batch_image_labels):
+    global total_files_processed
+    for img_ann in batch_image_labels:
+        write_result(img_ann)
+        total_files_processed += 1
+        print(f"Total processed images: {total_files_processed}")
+
+
+def main():
+    global total_files_processed
+    total_files_processed = 0
+
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    print(f"Using device: {device}")
+
+    geometric_pipeline = K.AugmentationSequential(
+        K.RandomHorizontalFlip(p=0.5),
+        K.RandomAffine(degrees=25, translate=(0.1, 0.1), scale=(0.8, 1.2), p=0.5),
+        data_keys=["input", "bbox"],
+        same_on_batch=False
+    ).to(device)
+
+    intensity_pipeline = K.AugmentationSequential(
+        K.ColorJitter(brightness=0.1, contrast=0.07, saturation=0.1, hue=0.1, p=0.5),
+        RandomFog(p=0.2),
+        RandomShadow(p=0.3),
+        K.RandomMotionBlur(kernel_size=3, angle=35., direction=0.5, p=0.3),
+        data_keys=["input"],
+        same_on_batch=False
+    ).to(device)
+
+
+    dataset = ImageDataset(data_images_dir, data_labels_dir)
+    dataloader = DataLoader(dataset, batch_size=32, shuffle=False, num_workers=os.cpu_count()) # Adjust batch_size as needed
+
+    processed_images_set = set(os.listdir(processed_images_dir))
+    images_to_process_indices = [i for i, filename in enumerate(dataset.image_filenames) if filename not in processed_images_set]
+    dataloader_filtered = torch.utils.data.Subset(dataset, images_to_process_indices)
+    filtered_dataloader = DataLoader(dataloader_filtered, batch_size=32, shuffle=False, num_workers=os.cpu_count())
+
+
+    start_time = time.time()
+
+    try:
+        for batch_data in filtered_dataloader:
+            batch_image = batch_data[0]
+            batch_labels = batch_data[1]
+            batch_filenames = batch_data[2]
+
+            batch_processed_image_labels = process_batch_wrapper(list(zip(batch_image, batch_labels, batch_filenames)), geometric_pipeline, intensity_pipeline, device)
+            save_batch_results(batch_processed_image_labels)
+    except Exception as e:
+        print(e)
+
+    end_time = time.time()
+    elapsed_time = end_time - start_time
+    print(f"Total processing time: {elapsed_time:.2f} seconds")
+    images_per_hour = (total_files_processed / elapsed_time) * 3600
+    print(f"Processed images per hour: {images_per_hour:.2f}")
+    print("Augmentation process completed.")
+
+
+if __name__ == '__main__':
+    main()

preprocessing.py

@@ -1,15 +1,32 @@
+import concurrent.futures
 import os.path
 import time
-from datetime import datetime, timedelta
 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 constants import (data_images_dir, data_labels_dir, processed_images_dir, processed_labels_dir)
-                       annotation_classes, checkpoint_file, checkpoint_date_format)
 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
@@ -37,31 +54,18 @@ def correct_bboxes(labels):
 
 
 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.SafeRotate(limit=90, 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')),
-        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'))
-    ]
-
     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):
+    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)
@@ -87,7 +91,8 @@ def write_result(img_ann: ImageLabel):
                  img_ann.labels]
         f.writelines(lines)
         f.close()
-    print(f'{img_ann.labels_path} written')
+    global total_files_processed
+    print(f'{total_files_processed}. {img_ann.labels_path} written')
 
 
 def read_labels(labels_path) -> [[]]:
@@ -104,19 +109,10 @@ def read_labels(labels_path) -> [[]]:
         return arr
 
 
-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,
-        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)
-    ))
-
-
 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)
 
@@ -126,27 +122,38 @@ def preprocess_annotations():
         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)
 
-    for image_file in images:
+def process_image_file(image_file): # this function will be executed in thread
+    try:
+        image_path = os.path.join(data_images_dir, image_file.name)
+        labels_path = os.path.join(data_labels_dir, f'{Path(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:
-            image_path = os.path.join(data_images_dir, image_file.name)
+            results = image_processing(img_ann)
-            labels_path = os.path.join(data_labels_dir, f'{Path(image_path).stem}.txt')
+            for res_ann in results:
-            image = cv2.imdecode(np.fromfile(image_path, dtype=np.uint8), cv2.IMREAD_UNCHANGED)
+                write_result(res_ann)
-            process_image(ImageLabel(
-                image_path=image_path,
-                image=image,
-                labels_path=labels_path,
-                labels=read_labels(labels_path)
-            ))
         except Exception as e:
-            print(f'Error appeared {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(120)
+        time.sleep(300)
 
 
 if __name__ == '__main__':

requirements.txt

@@ -3,14 +3,15 @@ torch
 torchvision
 torchaudio
 ultralytics
-albumentations~=2.0.4
+albumentations
 
-opencv-python~=4.11.0.86
+opencv-python
-matplotlib~=3.10.0
+matplotlib
-PyYAML~=6.0.2
+PyYAML
-cryptography~=44.0.1
+cryptography
-numpy~=2.1.1
+numpy
-requests~=2.32.3
+requests
 pyyaml
+boto3
 msgpack
 rstream

train.py

@@ -1,4 +1,4 @@
-import io
+import concurrent.futures
 import os
 import random
 import shutil
@@ -7,7 +7,6 @@ from datetime import datetime
 from os import path, replace, listdir, makedirs, scandir
 from os.path import abspath
 from pathlib import Path
-from utils import Dotdict
 
 import yaml
 from ultralytics import YOLO
@@ -15,13 +14,14 @@ from ultralytics import YOLO
 import constants
 from azaion_api import ApiCredentials, Api
 from cdn_manager import CDNCredentials, CDNManager
-from security import Security
 from constants import (processed_images_dir,
                        processed_labels_dir,
                        annotation_classes,
                        prefix, date_format,
                        datasets_dir, models_dir,
                        corrupted_images_dir, corrupted_labels_dir, sample_dir)
+from security import Security
+from utils import Dotdict
 
 today_folder = f'{prefix}{datetime.now():{date_format}}'
 today_dataset = path.join(datasets_dir, today_folder)
@@ -31,7 +31,7 @@ test_set = 10
 old_images_percentage = 75
 
 DEFAULT_CLASS_NUM = 80
-
+total_files_copied = 0
 
 def form_dataset(from_date: datetime):
     makedirs(today_dataset, exist_ok=True)
@@ -67,6 +67,25 @@ def form_dataset(from_date: datetime):
 
 
 def copy_annotations(images, folder):
+    global total_files_copied
+    total_files_copied = 0
+
+    def copy_image(image):
+        global total_files_copied
+        total_files_copied += 1
+        label_name = f'{Path(image.path).stem}.txt'
+        label_path = path.join(processed_labels_dir, label_name)
+        if check_label(label_path):
+            shutil.copy(image.path, path.join(destination_images, image.name))
+            shutil.copy(label_path, path.join(destination_labels, label_name))
+        else:
+            shutil.copy(image.path, path.join(corrupted_images_dir, image.name))
+            shutil.copy(label_path, path.join(corrupted_labels_dir, label_name))
+            print(f'Label {label_path} is corrupted! Copy with its image to the corrupted directory ({corrupted_labels_dir})')
+
+        if total_files_copied % 1000 == 0:
+            print(f'{total_files_copied} copied...')
+
     destination_images = path.join(today_dataset, folder, 'images')
     makedirs(destination_images, exist_ok=True)
 
@@ -78,19 +97,10 @@ def copy_annotations(images, folder):
 
     copied = 0
     print(f'Copying annotations to {destination_images} and {destination_labels} folders:')
-    for image in images:
+    with concurrent.futures.ThreadPoolExecutor() as executor:
-        label_name = f'{Path(image.path).stem}.txt'
+        executor.map(copy_image, images)
-        label_path = path.join(processed_labels_dir, label_name)
+
-        if check_label(label_path):
+
-            shutil.copy(image.path, path.join(destination_images, image.name))
-            shutil.copy(label_path, path.join(destination_labels, label_name))
-        else:
-            shutil.copy(image.path, path.join(corrupted_images_dir, image.name))
-            shutil.copy(label_path, path.join(corrupted_labels_dir, label_name))
-            print(f'Label {label_path} is corrupted! Copy with its image to the corrupted directory ({corrupted_labels_dir})')
-        copied = copied + 1
-        if copied % 1000 == 0:
-            print(f'{copied} copied...')
     print(f'Copied all {copied} annotations to {destination_images} and {destination_labels} folders')
 
 
@@ -143,12 +153,15 @@ def revert_to_processed_data(date):
 
 def get_latest_model():
     def convert(d: str):
+        if not d.startswith(prefix):
+            return None
         dir_date = datetime.strptime(d.replace(prefix, ''), '%Y-%m-%d')
         dir_model_path = path.join(models_dir, d, 'weights', 'best.pt')
         return {'date': dir_date, 'path': dir_model_path}
 
     dates = [convert(d) for d in next(os.walk(models_dir))[1]]
-    sorted_dates = list(sorted(dates, key=lambda x: x['date']))
+    dates = list(filter(lambda x : x is not None, dates))
+    sorted_dates = list(sorted(dates, key=lambda x: x['date'] ))
     if len(sorted_dates) == 0:
         return None, None
     last_model = sorted_dates[-1]
@@ -223,9 +236,8 @@ def validate(model_path):
 
 
 def upload_model(model_path: str):
-
+    model = YOLO(model_path)
-    # model = YOLO(model_path)
+    model.export(format="onnx", imgsz=1280, nms=True, batch=4)
-    # model.export(format="onnx", imgsz=1280, nms=True, batch=4)
     onnx_model = path.dirname(model_path) + Path(model_path).stem + '.onnx'
 
     with open(onnx_model, 'rb') as f_in:
@@ -250,9 +262,8 @@ def upload_model(model_path: str):
     api.upload_file('azaion.onnx.small', onnx_part_small)
 
 if __name__ == '__main__':
-    # model_path = train_dataset('2024-10-26', from_scratch=True)
+    model_path = train_dataset(from_scratch=True)
-    # validate(path.join('runs', 'detect', 'train7', 'weights', 'best.pt'))
+    validate(path.join('runs', 'detect', 'train7', 'weights', 'best.pt'))
-    # form_data_sample(500)
+    form_data_sample(500)
-    # convert2rknn()
+    convert2rknn()
-    model_path = 'azaion.pt'
+    upload_model('azaion.pt')
-    upload_model(model_path)