diff --git a/orangepi5/00 install-os.md b/orangepi5/00 install-os.md deleted file mode 100644 index 1ee04f9..0000000 --- a/orangepi5/00 install-os.md +++ /dev/null @@ -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) diff --git a/orangepi5/convert-install.sh b/orangepi5/convert-install.sh deleted file mode 100644 index 3684ad8..0000000 --- a/orangepi5/convert-install.sh +++ /dev/null @@ -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 diff --git a/orangepi5/convert.sh b/orangepi5/convert.sh deleted file mode 100755 index 7f0b9ab..0000000 --- a/orangepi5/convert.sh +++ /dev/null @@ -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 \ No newline at end of file diff --git a/preprocessing-cuda.py b/preprocessing-cuda.py new file mode 100644 index 0000000..8b8eb46 --- /dev/null +++ b/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() \ No newline at end of file diff --git a/preprocessing.py b/preprocessing.py index 170547a..d4c30d6 100644 --- a/preprocessing.py +++ b/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, - annotation_classes, checkpoint_file, checkpoint_date_format) +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 @@ -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) - 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) - process_image(ImageLabel( - image_path=image_path, - image=image, - labels_path=labels_path, - labels=read_labels(labels_path) - )) + results = image_processing(img_ann) + for res_ann in results: + write_result(res_ann) 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__': diff --git a/rabbit_receiver.py b/rabbit_receiver.py new file mode 100644 index 0000000..e69de29 diff --git a/requirements.txt b/requirements.txt index 61cbea4..33a0976 100644 --- a/requirements.txt +++ b/requirements.txt @@ -3,14 +3,15 @@ torch torchvision torchaudio ultralytics -albumentations~=2.0.4 +albumentations -opencv-python~=4.11.0.86 -matplotlib~=3.10.0 -PyYAML~=6.0.2 -cryptography~=44.0.1 -numpy~=2.1.1 -requests~=2.32.3 +opencv-python +matplotlib +PyYAML +cryptography +numpy +requests pyyaml +boto3 msgpack rstream \ No newline at end of file diff --git a/train.py b/train.py index 371ae5b..05cd5f3 100644 --- a/train.py +++ b/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: - 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})') - copied = copied + 1 - if copied % 1000 == 0: - print(f'{copied} copied...') + with concurrent.futures.ThreadPoolExecutor() as executor: + executor.map(copy_image, images) + + 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.export(format="onnx", imgsz=1280, nms=True, batch=4) + model = YOLO(model_path) + 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) - # validate(path.join('runs', 'detect', 'train7', 'weights', 'best.pt')) - # form_data_sample(500) - # convert2rknn() - model_path = 'azaion.pt' - upload_model(model_path) + model_path = train_dataset(from_scratch=True) + validate(path.join('runs', 'detect', 'train7', 'weights', 'best.pt')) + form_data_sample(500) + convert2rknn() + upload_model('azaion.pt')