Time profiler

models/create_yaml.py

@@ -0,0 +1,51 @@
+import argparse
+import yaml
+import os
+
+def main(args):
+    # Load the YAML content
+    yaml_path = args.yaml_path
+    dataset_path = args.dataset_path
+    augment = args.augment
+
+    yaml_content = {
+        'path': dataset_path,
+        'train': 'images/train',
+        'val': 'images/val',
+        'test': 'images/test',
+        'nc': 1,
+        'names': {
+            0: 'seep',
+        },
+        'patience': 20
+    }
+
+
+    if augment:
+        yaml_content.update({
+            'hsv_h': 0.015,  # Hue augmentation
+            'hsv_s': 0.7,    # Saturation augmentation
+            'hsv_v': 0.4,    # Value (brightness) augmentation
+            'degrees': 90,    # Image rotation (+/- deg)
+            'translate': 0.1,  # Image translation (+/- fraction)
+            'scale': 0.5,    # Image scale (+/- gain)
+            'shear': 0.0,    # Image shear (+/- deg)
+            'perspective': 0.0,  # Image perspective (+/- fraction), range 0-0.001
+            'flipud': 0.5,   # Image flip up-down (probability)
+            'fliplr': 0.5,   # Image flip left-right (probability)
+            'mosaic': 1.0,   # Image mosaic (probability)
+            'mixup': 0.5,    # Image mixup (probability)
+            'copy_paste': 0.5,  # Segment copy-paste (probability)
+        })
+
+    # Save the YAML content
+    with open(yaml_path, 'w') as file:
+        yaml.dump(yaml_content, file, default_flow_style=False)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Generate YAML file for YOLO training')
+    parser.add_argument("--yaml-path", required=True, help="Path to save the YAML file")
+    parser.add_argument("--dataset-path", required=True, help="Path to the dataset directory")
+    parser.add_argument("--augment", action='store_true', help="Enable data augmentation")
+    args = parser.parse_args()
+    main(args)

models/readme.md

@@ -0,0 +1,25 @@
+# Overview
+=====================================
+
+The purpose of the 'create_yaml.py' file here is to create a YOLOv8 training YAML file in a specified file path.
+
+## Usage
+-------------
+
+### Arguments
+
+The `create_yaml.py` script takes three mandatory arguments:
+
+1. **Path to Save the YAML File**:
+   - The directory path where the generated YAML file will be saved.
+2. **Path to the Dataset Directory**:
+   - The directory path containing the dataset used for training.
+3. **Data Augmentation Flag**:
+   - A boolean flag indicating whether data augmentation should be enabled (`True`) or disabled (`False`).
+
+### Example Command
+
+To run the script, use the following command format:
+
+```bash
+python create_yaml.py /path/to/save/yaml /path/to/dataset True

models/utils/readme.md

@@ -0,0 +1,47 @@
+# Overview
+=====================================
+This script is designed to profile the execution of a YOLO (You Only Look Once) model using PyTorch and the Ultralytics YOLO library. It allows users to measure the performance of the model on various input sizes and configurations, including the option to use CUDA for accelerated profiling.
+
+## Usage
+-------------
+
+To use this script, you need to have the following dependencies installed:
+
+- **Python**: Ensure you have Python installed on your system.
+- **PyTorch**: Install PyTorch using `pip install torch`.
+- **Ultralytics YOLO**: Install the Ultralytics YOLO library using `pip install ultralytics`.
+
+### Command-Line Arguments
+
+The script takes the following command-line arguments:
+
+#### `--yolo-path`
+- **Description**: The path to the YOLO model file.
+- **Required**: Yes
+
+#### `--batch-size`
+- **Description**: The batch size of the input tensor.
+- **Default**: 1
+- **Type**: Integer
+
+#### `--height`
+- **Description**: The height of the input tensor.
+- **Default**: 1024
+- **Type**: Integer
+
+#### `--width`
+- **Description**: The width of the input tensor.
+- **Default**: 1280
+- **Type**: Integer
+
+#### `--cuda`
+- **Description**: Use CUDA for profiling if available.
+- **Default**: False (CPU)
+- **Type**: Boolean (action='store_true')
+
+### Example Command
+
+To run the script, use the following command format:
+
+```bash
+python profile_yolo.py --yolo-path /path/to/yolo/model --batch-size 4 --height 640 --width 640 --cuda

models/utils/time_profiler.py

@@ -0,0 +1,33 @@
+import argparse
+import torch
+from ultralytics import YOLO
+
+def main(args):
+    # Load the YOLO model
+    model = YOLO(args.yolo_path)
+
+    # Create a random input tensor with the specified shape
+    x = torch.randn((args.batch_size, 3, args.height, args.width), requires_grad=True)
+
+    # Move the model and input to CUDA if specified
+    if args.cuda:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        model.to(device)
+        x = x.to(device)
+
+    # Profile the model execution
+    with torch.autograd.profiler.profile(use_cuda=args.cuda) as prof:
+        model(x)
+
+    # Print the profiling results
+    print(prof)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Profile YOLO model execution')
+    parser.add_argument("--yolo-path", required=True, help="Path to the YOLO model file")
+    parser.add_argument("--batch-size", type=int, default=1, help="Batch size of the input tensor")
+    parser.add_argument("--height", type=int, default=1024, help="Height of the input tensor")
+    parser.add_argument("--width", type=int, default=1280, help="Width of the input tensor")
+    parser.add_argument("--cuda", action='store_true', help="Use CUDA for profiling")
+    args = parser.parse_args()
+    main(args)

preprocess_data/utils.py

@@ -0,0 +1,214 @@
+import cv2
+import os
+import numpy as np
+from shapely.geometry import Polygon
+from PIL import Image, ImageDraw
+import shutil
+from sklearn.model_selection import train_test_split
+import pandas as pd
+from pathlib import Path
+
+
+colors = [
+        (255, 0, 0),   # Red
+        (0, 255, 0),   # Green
+        (0, 0, 255),   # Blue
+        (255, 255, 0), # Yellow
+        (255, 0, 255), # Magenta
+        (0, 255, 255), # Cyan
+        (128, 128, 128) # Gray
+    ]
+
+def mask_to_polygons_multi(mask):
+    h, w = mask.shape
+    polygons = []
+    labels = []
+    annotations = []
+
+    for label in np.unique(mask):
+        if label == 0:
+            continue  # Skip background
+
+        binary_mask = (mask == label).astype(np.uint8)
+        contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+
+        for contour in contours:
+            epsilon = 0.005 * cv2.arcLength(contour, True)
+            approx = cv2.approxPolyDP(contour, epsilon, True)
+
+            # Check if the contour has at least 4 points
+            if len(approx) < 4:
+                continue  # Skip this contour if it has fewer than 4 points
+
+            try:
+                poly = Polygon(approx.reshape(-1, 2))
+                poly = poly.simplify(1, preserve_topology=True)
+
+                if poly.is_valid and poly.area > 0:
+                    min_x, min_y, max_x, max_y = poly.bounds
+
+                    coords = np.array(poly.exterior.coords)
+                    coords[:, 0] /= w
+                    coords[:, 1] /= h
+
+                    polygons.append(coords.flatten().tolist())
+                    labels.append(int(label))
+
+                    bbox = [min_x, min_y, max_x - min_x, max_y - min_y]
+
+                    annotations.append({
+                        "segmentation": [coords.flatten().tolist()],
+                        "category_id": int(label),
+                        "bbox": bbox,
+                        "area": poly.area
+                    })
+            except Exception as e:
+                print(f"Error processing contour: {e}")
+                continue  # Skip this contour if there's any error
+
+    return polygons, labels, annotations
+
+def annotations_to_mask(annotations, image_shape=(256,256)):
+    """Convert COCO-style annotations back to a mask."""
+    mask = np.zeros(image_shape, dtype=np.uint8)
+    for ann in annotations:
+        polygon = np.array(ann['segmentation'][0]).reshape(-1, 2)
+        polygon = (polygon * np.array([image_shape[1], image_shape[0]])).astype(int)
+        cv2.fillPoly(mask, [polygon], int(ann['category_id'] + 1))
+    return mask
+
+
+def overlay_mask_on_image_with_boxes(image, mask, annotations, alpha=0.5):
+    """Overlay a colored mask on an image with bounding boxes and labels."""
+    # Convert grayscale image to RGB
+    if len(image.shape) == 2 or image.shape[2] == 1:
+        image_rgb = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
+    else:
+        image_rgb = image
+
+    # Define a color map with distinct colors for each class
+
+
+    # Create a colored mask
+    colored_mask = np.zeros((*mask.shape, 3), dtype=np.uint8)
+    for i in range(1, mask.max() + 1):
+        colored_mask[mask == i] = colors[i % len(colors)]
+
+    # Overlay the mask on the image
+    overlay = (image_rgb * (1 - alpha) + colored_mask * alpha).astype(np.uint8)
+    overlay_image = Image.fromarray(overlay)
+    draw = ImageDraw.Draw(overlay_image)
+
+    # Draw bounding boxes
+    for ann in annotations:
+        bbox = ann['bbox']
+        x, y, w, h = bbox
+        color = colors[(ann['category_id']) % len(colors)]
+        draw.rectangle([x, y, x+w, y+h], outline=color, width=2)
+
+    return np.array(overlay_image)
+
+def extract_instance(image, mask, instance_id):
+    """Extract a single instance from the image using the mask."""
+    instance_mask = (mask == instance_id).astype(np.uint8) * instance_id
+    roi = cv2.bitwise_and(image, image, mask=(instance_mask > 0).astype(np.uint8))
+    return roi, instance_mask
+
+def paste_instance(base_image, roi, instance_mask):
+    """Paste the extracted ROI onto a copy of the base image."""
+    result_image = base_image.copy()
+    result_image[instance_mask > 0] = roi[instance_mask > 0]
+    return result_image
+
+def process_mask_images_to_poygon_txt(input_dir, output_dir):
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+
+    mask_files = [f for f in os.listdir(input_dir) if f.endswith(('.png', '.jpg', '.tif'))]
+
+    for mask_file in mask_files:
+        mask_path = os.path.join(input_dir, mask_file)
+        mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
+
+        polygons, labels, _ = mask_to_polygons_multi(mask)
+
+        txt_filename = os.path.splitext(mask_file)[0] + '.txt'
+        txt_path = os.path.join(output_dir, txt_filename)
+
+        with open(txt_path, 'w') as f:
+            for polygon, label in zip(polygons, labels):
+                # Subtract 1 from label to start class indexing from 0 for YOLO format
+                line = f"{label-1} " + " ".join(map(str, polygon))
+                f.write(line + '\n')
+
+def process_tif_to_png_and_verify(tif_path):
+    # Load the TIFF image using cv2
+    original_image = cv2.imread(tif_path)
+
+    if original_image is None:
+        raise ValueError(f"Failed to load image: {tif_path}")
+
+    # Construct the output PNG path
+    png_path = tif_path.replace('.tif', '.png')
+
+    # Save the image as PNG
+    cv2.imwrite(png_path, original_image)
+
+    # Read back the PNG image
+    loaded_image = cv2.imread(png_path)
+
+    # Compare the arrays
+    are_equal = np.array_equal(original_image, loaded_image)
+
+    # Delete the PNG file
+    os.remove(png_path)
+
+    return are_equal, original_image, loaded_image
+
+
+def create_dataset_df(base_dir):
+    data = []
+    label_dir = os.path.join(base_dir, 'labels')
+    image_dir = os.path.join(base_dir, 'preprocessed_images')
+
+    # Iterate through all txt files in the label directory
+    for txt_file in Path(label_dir).glob('*.txt'):
+        txt_path = str(txt_file.relative_to(base_dir))
+        png_name = txt_file.stem + '.png'
+        png_path = str(Path('preprocessed_images') / png_name)
+
+        # Check if the corresponding image file exists
+        if not os.path.exists(os.path.join(base_dir, png_path)):
+            print(f"Warning: Image file not found for {txt_path}")
+            continue
+
+        # Read labels from the txt file
+        with open(txt_file, 'r') as f:
+            for line in f:
+                label = line.strip().split()[0]  # Get the first value of each line
+                data.append([txt_path, png_path, label])
+
+    # Create DataFrame
+    df = pd.DataFrame(data, columns=['label_file', 'image_file', 'label'])
+    return df
+
+def split_dataset(df):
+    # Create unique image-label pairs
+    unique_pairs = df[['label_file', 'image_file', 'label']].drop_duplicates().reset_index(drop=True)
+
+    # Split into train and test (85:15)
+    train_val, test = train_test_split(unique_pairs, test_size=0.05, stratify=unique_pairs['label'], random_state=42)
+
+    # Further split train into train and validation (80:20 of the original train set)
+    train, val = train_test_split(train_val, test_size=0.15, stratify=train_val['label'], random_state=42)
+
+    return train, val, test
+
+def move_files(df, source_dir, dest_dir):
+    os.makedirs(dest_dir, exist_ok=True)
+    for _, row in df.iterrows():
+        image_file = os.path.join(source_dir, row['image_file'])
+        label_file = os.path.join(source_dir, row['label_file'])
+
+        shutil.copy(image_file, dest_dir)
+        shutil.copy(label_file, dest_dir)