update cdn manager

download big engine file

cdn_manager.py

@@ -42,29 +42,3 @@ class CDNManager:
         except Exception as e:
             print(e)
             return False
-
-
-config_file = os.path.join('downloader_config.yaml')
-with open(config_file, 'r') as f:
-    config = yaml.safe_load(f)
-
-cdn_manager = CDNManager(CDNCredentials(
-    config['host'],
-    config['downloader_access_key'],
-    config['downloader_access_secret'],
-    config['uploader_access_key'],
-    config['uploader_access_secret']
-))
-
-input_action = sys.argv[1]
-input_bucket = sys.argv[2]
-input_filename = sys.argv[3]
-if len(sys.argv) > 4: # 0 is this script's path, hence 5 args is max
-    input_path = sys.argv[4]
-
-if input_action == 'upload':
-    with open(input_path, 'rb') as f:
-        cdn_manager.upload(input_bucket, input_filename, bytearray(f.read()))
-elif input_action == 'download':
-    cdn_manager.download(input_bucket, input_filename)
-

preprocessing-cuda.py

@@ -1,269 +0,0 @@
-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()

start_inference.py

@@ -3,6 +3,7 @@ import re
 import yaml
 import constants
 from azaion_api import Api, ApiCredentials
+from cdn_manager import CDNManager, CDNCredentials
 from inference.inference import Inference
 from inference.onnx_engine import OnnxEngine
 from inference.tensorrt_engine import TensorRTEngine
@@ -22,13 +23,18 @@ if __name__ == "__main__":
     # Inference(TensorRTEngine('azaion-2025-03-10_batch8.engine', batch_size=8),
     #           confidence_threshold=0.5, iou_threshold=0.3).process('ForAI_test.mp4')
 
-    CDNManager()
-
     with open(constants.CONFIG_FILE, "r") as f:
         config_dict = yaml.safe_load(f)
         d_config = Dotdict(config_dict)
         cdn_c = Dotdict(d_config.cdn)
         api_c = Dotdict(d_config.api)
+
+    cdn_manager = CDNManager(CDNCredentials(cdn_c.host,
+                                            cdn_c.downloader_access_key, cdn_c.downloader_access_secret,
+                                            cdn_c.uploader_access_key, cdn_c.uploader_access_secret))
+    cdn_manager.download(cdn_c.bucket, constants.AI_TENSOR_MODEL_FILE_BIG)
+    cdn_manager.download(cdn_c.bucket, constants.AI_ONNX_MODEL_FILE_BIG)
+
     api_client = Api(ApiCredentials(api_c.url, api_c.user, api_c.pw, api_c.folder))
 
     tensor_model_bytes = api_client.load_resource(constants.AI_TENSOR_MODEL_FILE_BIG, constants.AI_TENSOR_MODEL_FILE_SMALL)