bplusplus 1.2.2__py3-none-any.whl → 1.2.4__py3-none-any.whl

bplusplus/prepare.py

@@ -2,17 +2,12 @@ import os
 import random
 import shutil
 import tempfile
-from collections import defaultdict
 from pathlib import Path
-from typing import Any, Optional
+from typing import Optional
 
-import matplotlib.pyplot as plt
-import numpy as np
 import requests
 import torch
-import yaml
-from PIL import Image, ImageDraw, ImageFont
-from prettytable import PrettyTable
+from PIL import Image
 from torch import serialization
 from torch.nn import Module, ModuleDict, ModuleList
 from torch.nn.modules.activation import LeakyReLU, ReLU, SiLU
@@ -31,65 +26,155 @@ from ultralytics.nn.modules.block import DFL
 from ultralytics.nn.modules.conv import Conv
 from ultralytics.nn.tasks import DetectionModel
 
-from .collect import Group, collect
-
-
-def prepare(input_directory: str, output_directory: str, one_stage: bool = False, with_background: bool = False, size_filter: bool = False, sizes: list = None):
 
+def prepare(input_directory: str, output_directory: str, img_size: int = 40):
     """
-    Prepares the dataset for training by performing the following steps:
-    1. Copies images from the input directory to a temporary directory.
-    2. Deletes corrupted images.
-    3. Downloads YOLOv5 weights if not already present.
-    4. Runs YOLOv5 inference to generate labels for the images.
-    5. Deletes orphaned images and inferences.
-    6. Updates labels based on class mapping.
-    7. Splits the data into train, test, and validation sets.
-    8. Counts the total number of images across all splits.
-    9. Makes a YAML configuration file for YOLOv8.
+    Prepares a YOLO classification dataset by performing the following steps:
+    1. Copies images from input directory to temporary directory and creates class mapping.
+    2. Deletes corrupted images and downloads YOLO model weights if not present.
+    3. Runs YOLO inference to generate detection labels (bounding boxes) for the images.
+    4. Cleans up orphaned images, invalid labels, and updates labels with class indices.
+    5. Crops detected objects from images based on bounding boxes and resizes them.
+    6. Splits data into train/valid sets with classification folder structure (train/class_name/image.jpg).
 
     Args:
         input_directory (str): The path to the input directory containing the images.
-        output_directory (str): The path to the output directory where the prepared dataset will be saved.
+        output_directory (str): The path to the output directory where the prepared classification dataset will be saved.
+        img_size (int, optional): The target size for the smallest dimension of cropped images. Defaults to 40.
     """
-
     input_directory = Path(input_directory)
     output_directory = Path(output_directory)
 
-    class_mapping={}
+    print("="*60)
+    print("STARTING BPLUSPLUS DATASET PREPARATION")
+    print("="*60)
+    print(f"Input directory: {input_directory}")
+    print(f"Output directory: {output_directory}")
+    print(f"Target image size: {img_size}px (smallest dimension)")
+    print()
 
     with tempfile.TemporaryDirectory() as temp_dir:
-
         temp_dir_path = Path(temp_dir)
-        images_path = temp_dir_path / "images"
-
-        images_path.mkdir(parents=True, exist_ok=True)
-
-        for folder_directory in input_directory.iterdir():
-            images_names = []
-            if folder_directory.is_dir():
-                folder_name = folder_directory.name
-                for image_file in folder_directory.glob("*.jpg"):
-                    shutil.copy(image_file, images_path)
-                    image_name = image_file.name
-                    images_names.append(image_name)
-
-                class_mapping[folder_name] = images_names   
-
-        original_image_count = len(list(images_path.glob("*.jpg"))) + len(list(images_path.glob("*.jpeg")))
-
-        __delete_corrupted_images(images_path)
-
-        current_dir = Path(__file__).resolve().parent
-
-        weights_path = current_dir / 'small-generic.pt'
-
-        github_release_url = 'https://github.com/orlandocloss/TwoStageInsectDetection/releases/download/models/small-generic.pt'
+        print(f"Using temporary directory: {temp_dir_path}")
+        print()
+        
+        # Step 1: Setup directories and copy images
+        print("STEP 1: Setting up directories and copying images...")
+        print("-" * 50)
+        class_mapping, original_image_count = _setup_directories_and_copy_images(
+            input_directory, temp_dir_path
+        )
+        print(f"✓ Step 1 completed: {original_image_count} images copied from {len(class_mapping)} classes")
+        print()
+        
+        # Step 2-3: Clean images and setup model
+        print("STEP 2: Cleaning images and setting up YOLO model...")
+        print("-" * 50)
+        weights_path = _prepare_model_and_clean_images(temp_dir_path)
+        print(f"✓ Step 2 completed: Model ready at {weights_path}")
+        print()
+        
+        # Step 4: Run YOLO inference
+        print("STEP 3: Running YOLO inference to detect objects...")
+        print("-" * 50)
+        labels_path = _run_yolo_inference(temp_dir_path, weights_path)
+        print(f"✓ Step 3 completed: Labels generated at {labels_path}")
+        print()
+        
+        # Step 5-6: Clean up labels and update class mapping
+        print("STEP 4: Cleaning up orphaned files and processing labels...")
+        print("-" * 50)
+        class_idxs = _cleanup_and_process_labels(
+            temp_dir_path, labels_path, class_mapping
+        )
+        print(f"✓ Step 4 completed: Processed {len(class_idxs)} classes")
+        print()
+        
+        # Step 7-9: Finalize dataset
+        print("STEP 5: Creating classification dataset with cropped images...")
+        print("-" * 50)
+        _finalize_dataset(
+            class_mapping, temp_dir_path, output_directory, 
+            class_idxs, original_image_count, img_size
+        )
+        print("✓ Step 5 completed: Classification dataset ready!")
+        print()
+        
+    print("="*60)
+    print("BPLUSPLUS DATASET PREPARATION COMPLETED SUCCESSFULLY!")
+    print("="*60)
 
-        if not weights_path.exists():
-            __download_file_from_github_release(github_release_url, weights_path)
+def _setup_directories_and_copy_images(input_directory: Path, temp_dir_path: Path):
+    """
+    Sets up temporary directories and copies images from input directory.
+    
+    Returns:
+        tuple: (class_mapping dict, original_image_count int)
+    """
+    images_path = temp_dir_path / "images"
+    images_path.mkdir(parents=True, exist_ok=True)
+    print(f"  Created temporary images directory: {images_path}")
+    
+    class_mapping = {}
+    total_copied = 0
+    
+    print("  Scanning input directory for class folders...")
+    class_folders = [d for d in input_directory.iterdir() if d.is_dir()]
+    print(f"  Found {len(class_folders)} class folders")
+    
+    for folder_directory in class_folders:
+        images_names = []
+        if folder_directory.is_dir():
+            folder_name = folder_directory.name
+            image_files = list(folder_directory.glob("*.jpg"))
+            print(f"  Copying {len(image_files)} images from class '{folder_name}'...")
+            
+            for image_file in image_files:
+                shutil.copy(image_file, images_path)
+                image_name = image_file.name
+                images_names.append(image_name)
+                total_copied += 1
+            
+            class_mapping[folder_name] = images_names
+            print(f"    ✓ {len(images_names)} images copied for class '{folder_name}'")
+    
+    original_image_count = len(list(images_path.glob("*.jpg"))) + len(list(images_path.glob("*.jpeg")))
+    print(f"  Total images in temporary directory: {original_image_count}")
+    
+    return class_mapping, original_image_count
 
-        # Add all required classes to safe globals
+def _prepare_model_and_clean_images(temp_dir_path: Path):
+    """
+    Cleans corrupted images and downloads/prepares the YOLO model.
+    
+    Returns:
+        Path: weights_path for the YOLO model
+    """
+    images_path = temp_dir_path / "images"
+    
+    # Clean corrupted images
+    print("  Checking for corrupted images...")
+    images_before = len(list(images_path.glob("*.jpg")))
+    __delete_corrupted_images(images_path)
+    images_after = len(list(images_path.glob("*.jpg")))
+    deleted_count = images_before - images_after
+    print(f"  ✓ Cleaned {deleted_count} corrupted images ({images_after} images remain)")
+    
+    # Setup model weights
+    current_dir = Path(__file__).resolve().parent
+    weights_path = current_dir / 'v11small-generic.pt'
+    github_release_url = 'https://github.com/Tvenver/Bplusplus/releases/download/v1.2.3/v11small-generic.pt'
+    
+    print(f"  Checking for YOLO model weights at: {weights_path}")
+    if not weights_path.exists():
+        print("  Model weights not found, downloading from GitHub...")
+        __download_file_from_github_release(github_release_url, weights_path)
+        print(f"  ✓ Model weights downloaded successfully")
+    else:
+        print("  ✓ Model weights already exist")
+    
+    # Add all required classes to safe globals
+    if hasattr(serialization, 'add_safe_globals'):
         serialization.add_safe_globals([
             DetectionModel, Sequential, Conv, Conv2d, BatchNorm2d, 
             SiLU, ReLU, LeakyReLU, MaxPool2d, Linear, Dropout, Upsample,
@@ -101,324 +186,131 @@ def prepare(input_directory: str, output_directory: str, one_stage: bool = False
             torch._utils._rebuild_tensor_v2,
             torch._utils._rebuild_parameter
         ])
-
-        labels_path = temp_dir_path / "predict" / "labels"
-        
-        try:
-            print(f"Loading YOLO model from {weights_path}")
-            model = YOLO(weights_path)
-            
-            # Get list of all image files
-            image_files = list(images_path.glob('*.jpg'))
-            print(f"Found {len(image_files)} images to process")
-            
-            # Ensure predict directory exists
-            predict_dir = temp_dir_path / "predict"
-            predict_dir.mkdir(exist_ok=True)
-            labels_path.mkdir(parents=True, exist_ok=True)
-            
-            result_count = 0
-            error_count = 0
-            
-            for img_path in image_files:
-                try:
-                    results = model.predict(
-                        source=str(img_path),
-                        conf=0.5,
-                        save=True,
-                        save_txt=True,
-                        project=temp_dir_path,
-                        name="predict",
-                        exist_ok=True,
-                        verbose=True
-                    )
-                    
-                    result_count += 1
-                    
-                except Exception as e:
-                    error_count += 1
-                    print(f"Error processing {img_path.name}: {e}")
-                    continue
-            
-            print(f"Model prediction completed: {result_count} successful, {error_count} failed")
-            print(f"Checking for labels in {labels_path}")
-            
-            # Verify labels were created
-            label_files = list(labels_path.glob("*.txt"))
-            print(f"Found {len(label_files)} label files")
-            
-            if len(label_files) == 0:
-                print("WARNING: No label files were created by the model prediction!")
-                
-        except Exception as e:
-            print(f"Error during model prediction setup: {e}")
-            import traceback
-            traceback.print_exc()
-
-        if one_stage:
-
-            if size_filter and len(sizes) <= 2:
-                __filter_by_size(images_path, labels_path, sizes)
-                print(f"\nFiltered {len(list(images_path.glob('*.jpg')))} images by size out of {original_image_count} input images.\n NOTE: Some images may be filtered due to corruption or inaccurate labels.")
-
-            __delete_orphaned_images_and_inferences(images_path, labels_path)
-            __delete_invalid_txt_files(images_path, labels_path)
-            class_idxs = update_labels(class_mapping, labels_path)
-            __split_data(class_mapping, temp_dir_path, output_directory)
-
-            # __save_class_idx_to_file(class_idxs, output_directory)
-            final_image_count = count_images_across_splits(output_directory)
-            print(f"\nOut of {original_image_count} input images, {final_image_count} are eligible for detection. \nThese are saved across train, test and valid split in {output_directory}.")
-            __generate_sample_images_with_detections(output_directory, class_idxs)
-
-            if with_background:
-                print("\nCollecting and splitting background images.")
-
-                bg_images=int(final_image_count*0.06)
-
-                search: dict[str, Any] = {
-                    "scientificName": ["Plantae"]
-                }
-
-                collect(
-                    group_by_key=Group.scientificName,
-                    search_parameters=search, 
-                    images_per_group=bg_images,
-                    output_directory=temp_dir_path,
-                    num_threads=3
-                )
-
-                __delete_corrupted_images(temp_dir_path / "Plantae")
-
-                __split_background_images(temp_dir_path / "Plantae", output_directory)
-
-            __count_classes_and_output_table(output_directory, class_idxs)
-
-            __make_yaml_file(output_directory, class_idxs)
-        else:
-            # try:
-            #     sized_dir = temp_dir_path / "sized"
-            #     sized_dir.mkdir(parents=True, exist_ok=True)
-            #     __two_stage_update(class_mapping, filtered, sized_dir, images_path)
-            #     __classification_split(sized_dir, output_directory)
-            #     __count_classification_split(output_directory, class_mapping)
-            # except:
-            __delete_orphaned_images_and_inferences(images_path, labels_path)
-            __delete_invalid_txt_files(images_path, labels_path)
-            __classification_split(images_path, labels_path, output_directory, class_mapping)
-            __count_classification_split(output_directory, class_mapping)
-            
-def __count_classification_split(output_directory: str, class_mapping: dict):
-    """
-    Counts the number of images in the train and valid splits for each class.
-
-    Args:
-        output_directory (str): Path to the output directory containing train and valid splits.
-        class_mapping (dict): Dictionary mapping class names to image file names.
-    """
-    class_counts = {}
-    train_counts = {}
-    valid_counts = {}
     
-    for class_name in class_mapping.keys():
-        train_dir = output_directory / 'train' / class_name
-        valid_dir = output_directory / 'valid' / class_name
-        
-        train_count = len(list(train_dir.glob("*.jpg"))) if train_dir.exists() else 0
-        valid_count = len(list(valid_dir.glob("*.jpg"))) if valid_dir.exists() else 0
-        total_count = train_count + valid_count
-        
-        class_counts[class_name] = total_count
-        train_counts[class_name] = train_count
-        valid_counts[class_name] = valid_count
-    
-    table = PrettyTable()
-    table.field_names = ["Class", "Train", "Valid", "Total"]
-    for class_name in class_mapping.keys():
-        table.add_row([
-            class_name, 
-            train_counts[class_name], 
-            valid_counts[class_name], 
-            class_counts[class_name]
-        ])
-    print(table)
-    # print(f"Saved in {output_directory}")
+    return weights_path
 
-def __classification_split(input_directory: str, labels_directory: str, output_directory: str, class_mapping: dict):
+def _run_yolo_inference(temp_dir_path: Path, weights_path: Path):
     """
-    Splits the data into train and validation sets for classification tasks,
-    cropping images according to their YOLO labels but preserving original class structure.
+    Runs YOLO inference on all images to generate labels.
     
-    Args:
-        input_directory (str): Path to the input directory containing images.
-        labels_directory (str): Path to the directory containing YOLO label files.
-        output_directory (str): Path to the output directory where train and valid splits will be created.
-        class_mapping (dict): Dictionary mapping class names to image file names.
+    Returns:
+        Path: labels_path where the generated labels are stored
     """
-    input_directory = Path(input_directory)
-    labels_directory = Path(labels_directory)
-    output_directory = Path(output_directory)
-    
-    # Create train and valid directories
-    train_dir = output_directory / 'train'
-    valid_dir = output_directory / 'valid'
+    images_path = temp_dir_path / "images"
+    labels_path = temp_dir_path / "predict" / "labels"
     
-    train_dir.mkdir(parents=True, exist_ok=True)
-    valid_dir.mkdir(parents=True, exist_ok=True)
-    
-    # Create class directories based on class_mapping
-    for class_name in class_mapping:
-        (train_dir / class_name).mkdir(exist_ok=True)
-        (valid_dir / class_name).mkdir(exist_ok=True)
-        print(f"Created directory for class: {class_name}")
-    
-    # Process each class folder and its images
-    valid_images = []
-    
-    # First, collect all valid label files
-    valid_label_stems = {label_file.stem for label_file in labels_directory.glob("*.txt") 
-                        if label_file.exists() and os.path.getsize(label_file) > 0}
-    
-    print(f"Found {len(valid_label_stems)} valid label files")
-    
-    for class_name, image_names in class_mapping.items():
-        print(f"Processing class: {class_name} with {len(image_names)} images")
+    try:
+        print(f"  Loading YOLO model from: {weights_path}")
+        model = YOLO(weights_path)
+        print("  ✓ YOLO model loaded successfully")
         
-        for image_name in image_names:
-            # Check if the image exists directly in the input directory
-            image_path = input_directory / image_name
-            
-            if not image_path.exists():
-                continue
-                
-            # Skip images that don't have a valid label
-            if image_path.stem not in valid_label_stems:
-                continue
-                
-            label_file = labels_directory / (image_path.stem + '.txt')
-            
+        # Get list of all image files
+        image_files = list(images_path.glob('*.jpg'))
+        print(f"  Found {len(image_files)} images to process with YOLO")
+        
+        # Ensure predict directory exists
+        predict_dir = temp_dir_path / "predict"
+        predict_dir.mkdir(exist_ok=True)
+        labels_path.mkdir(parents=True, exist_ok=True)
+        print(f"  Created prediction output directory: {predict_dir}")
+        
+        result_count = 0
+        error_count = 0
+        
+        print("  Starting YOLO inference...")
+        print(f"  Progress: 0/{len(image_files)} images processed", end="", flush=True)
+        
+        for i, img_path in enumerate(image_files, 1):
             try:
-                img = Image.open(image_path)
-                
-                if label_file.exists():
-                    # If label exists, crop the image
-                    with open(label_file, 'r') as f:
-                        lines = f.readlines()
-                        if lines:
-                            parts = lines[0].strip().split()
-                            if len(parts) >= 5:
-                                x_center, y_center, width, height = map(float, parts[1:5])
-                                
-                                img_width, img_height = img.size
-                                x_min = int((x_center - width/2) * img_width)
-                                y_min = int((y_center - height/2) * img_height)
-                                x_max = int((x_center + width/2) * img_width)
-                                y_max = int((y_center + height/2) * img_height)
-                                
-                                x_min = max(0, x_min)
-                                y_min = max(0, y_min)
-                                x_max = min(img_width, x_max)
-                                y_max = min(img_height, y_max)
-                                
-                                img = img.crop((x_min, y_min, x_max, y_max))
+                results = model.predict(
+                    source=str(img_path),
+                    conf=0.35,
+                    save=True,
+                    save_txt=True,
+                    project=temp_dir_path,
+                    name="predict",
+                    exist_ok=True,
+                    verbose=False  # Set to False to reduce YOLO's own output
+                )
                 
-                img_width, img_height = img.size
-                if img_width < img_height:
-                    # Width is smaller, make it 40
-                    new_width = 40
-                    new_height = int((img_height / img_width) * 40)
-                else:
-                    # Height is smaller, make it 40
-                    new_height = 40
-                    new_width = int((img_width / img_height) * 40)
+                result_count += 1
                 
-                #blur the image
-                img = img.resize((new_width, new_height), Image.LANCZOS)
+                # Update progress every 10% or every 100 images, whichever is smaller
+                update_interval = max(1, min(100, len(image_files) // 10))
+                if i % update_interval == 0 or i == len(image_files):
+                    print(f"\r  Progress: {i}/{len(image_files)} images processed", end="", flush=True)
                 
-                valid_images.append((image_path, img, class_name))
             except Exception as e:
-                print(f"Error processing {image_path}: {e}")
-    
-    print(f"Successfully processed {len(valid_images)} valid images for classification")
-    
-    # Shuffle and split images
-    random.shuffle(valid_images)
-    split_idx = int(len(valid_images) * 0.9)
-    train_images = valid_images[:split_idx]
-    valid_images = valid_images[split_idx:]
-    
-    print(f"Split into {len(train_images)} training images and {len(valid_images)} validation images")
-    
-    # Save images to train/valid directories
-    for image_set, dest_dir in [(train_images, train_dir), (valid_images, valid_dir)]:
-        for orig_file, img, class_name in image_set:
-            output_path = dest_dir / class_name / (orig_file.stem + '.jpg')
+                error_count += 1
+                print(f"\n  Error processing {img_path.name}: {e}")
+                continue
+        
+        print()  # New line after progress
+        print(f"  ✓ YOLO inference completed: {result_count} successful, {error_count} failed")
+        
+        # Verify labels were created
+        label_files = list(labels_path.glob("*.txt"))
+        print(f"  Generated {len(label_files)} label files")
+        
+        if len(label_files) == 0:
+            print("WARNING: No label files were created by the model prediction!")
             
-            # Convert any non-RGB mode to RGB before saving
-            if img.mode != 'RGB':
-                img = img.convert('RGB')
-                
-            img.save(output_path, format='JPEG', quality=95)
+    except Exception as e:
+        print(f"Error during model prediction setup: {e}")
+        import traceback
+        traceback.print_exc()
     
-    # Print summary
-    print(f"\nData split complete. Images saved to train and validation sets in {output_directory}")
-    for class_name in class_mapping:
-        train_count = len(list((train_dir / class_name).glob('*.*')))
-        valid_count = len(list((valid_dir / class_name).glob('*.*')))
-        print(f"  - {class_name}: {train_count} images in train, {valid_count} images in valid")
+    return labels_path
 
-def __filter_by_size(images_path: Path, labels_path: Path, sizes: list):
+def _cleanup_and_process_labels(temp_dir_path: Path, labels_path: Path, class_mapping: dict):
     """
-    Filters images by size and updates labels accordingly.
-
-    Args:
-        images_path (Path): The path to the directory containing images.
-        labels_path (Path): The path to the directory containing labels.
-        sizes (list): A list of sizes to filter by.
+    Cleans up orphaned images and invalid labels, then creates class index mapping.
+    
+    Returns:
+        dict: class_idxs mapping class indices to class names
     """
-    size_map={
-    "small": [0, 0.15],
-    "medium": [0.15, 0.3],
-    "large": [0.3, 1],
-    }
-        
-    filtered_images = []
-    for image_file in images_path.glob("*.jpg"):
-        label_file = labels_path / (image_file.stem + ".txt")
-        image_name = image_file.name
-
-        if label_file.exists():
-            with open(label_file, 'r') as file:
-                lines = file.readlines()
-                if len(lines) != 1:
-                    continue
-                else:
-                    parts = lines[0].split()
-                    _, _, width, height = map(float, parts[1:])
-                    for size in sizes:
-                        if width < size_map[size][1] and width >= size_map[size][0] and height < size_map[size][1] and height >= size_map[size][0]:
-                            filtered_images.append(image_name)
+    images_path = temp_dir_path / "images"
     
-    for image_file in images_path.glob("*.jpg"):
-        label_file = labels_path / (image_file.stem + ".txt")
-        image_name = image_file.name
-        if image_name not in filtered_images:
-            image_file.unlink()
-            try:
-                label_file.unlink()
-            except FileNotFoundError:
-                pass
+    print("  Cleaning up orphaned images and labels...")
+    images_before = len(list(images_path.glob("*.jpg")))
+    labels_before = len(list(labels_path.glob("*.txt")))
+    
+    __delete_orphaned_images_and_inferences(images_path, labels_path)
+    __delete_invalid_txt_files(images_path, labels_path)
+    
+    images_after = len(list(images_path.glob("*.jpg")))
+    labels_after = len(list(labels_path.glob("*.txt")))
+    
+    deleted_images = images_before - images_after
+    deleted_labels = labels_before - labels_after
+    print(f"  ✓ Cleaned up {deleted_images} orphaned images and {deleted_labels} invalid labels")
+    print(f"  Final counts: {images_after} images, {labels_after} valid labels")
+    
+    # Create class index mapping for classification
+    class_idxs = {}
+    for idx, class_name in enumerate(class_mapping.keys()):
+        class_idxs[idx] = class_name
+    
+    print(f"  Created class mapping for {len(class_idxs)} classes: {list(class_idxs.values())}")
+    
+    return class_idxs
 
-def __two_stage_update(class_mapping: dict, filtered_images: Path, output_directory: Path, images_path: Path):
+def _finalize_dataset(class_mapping: dict, temp_dir_path: Path, output_directory: Path, 
+                     class_idxs: dict, original_image_count: int, img_size: int):
     """
-    Prepares folders with class name containing filtered images.
+    Finalizes the dataset by creating cropped classification images and splitting into train/valid sets.
     """
-
-    for class_name, images in class_mapping.items():
-        for image_name in images:
-            if image_name in filtered_images:
-                (output_directory / class_name).mkdir(parents=True, exist_ok=True)
-                shutil.copy(images_path / image_name, output_directory /  class_name / image_name)
+    # Split data into train/valid with cropped classification images
+    __classification_split(class_mapping, temp_dir_path, output_directory, img_size)
+    
+    # Generate final report
+    print("  Generating final statistics...")
+    final_image_count = count_images_across_splits(output_directory)
+    print(f"  Dataset Statistics:")
+    print(f"    - Original images: {original_image_count}")
+    print(f"    - Final cropped images: {final_image_count}")
+    print(f"    - Success rate: {final_image_count/original_image_count*100:.1f}%")
+    print(f"    - Output directory: {output_directory}")
 
 def __delete_corrupted_images(images_path: Path):
      
@@ -500,7 +392,7 @@ def __delete_orphaned_images_and_inferences(images_path: Path, labels_path: Path
             # print(f"Deleting orphaned image: {image_file.name}")
             image_file.unlink()
 
-    print("Orphaned images files without corresponding labels have been deleted.")
+
 
 def __delete_invalid_txt_files(images_path: Path, labels_path: Path):
 
@@ -535,291 +427,188 @@ def __delete_invalid_txt_files(images_path: Path, labels_path: Path):
                 image_file_jpeg.unlink()
                 # print(f"Deleted corresponding image file: {image_file_jpeg.name}")
 
-    print("Invalid text files and their corresponding images files have been deleted.")
 
 
-def __split_data(class_mapping: dict, temp_dir_path: Path, output_directory: Path):
-    """
-    Splits the data into train, test, and validation sets.
 
+def __classification_split(class_mapping: dict, temp_dir_path: Path, output_directory: Path, img_size: int):
+    """
+    Splits the data into train and validation sets for classification tasks,
+    cropping images according to their YOLO labels but preserving original class structure.
+    
     Args:
         class_mapping (dict): A dictionary mapping class names to image file names.
         temp_dir_path (Path): The path to the temporary directory containing the images.
-        output_directory (Path): The path to the output directory where the split data will be saved.
+        output_directory (Path): The path to the output directory where train and valid splits will be created.
+        img_size (int): The target size for the smallest dimension of cropped images.
     """
     images_dir = temp_dir_path / "images"
     labels_dir = temp_dir_path / "predict" / "labels"
-
-    def create_dirs(split):
-        (output_directory / split).mkdir(parents=True, exist_ok=True)
-        (output_directory / split / "images").mkdir(parents=True, exist_ok=True)
-        (output_directory / split / "labels").mkdir(parents=True, exist_ok=True)
-
-    def copy_files(file_list, split):
-        for image_file in file_list:
-            image_file_path = images_dir / image_file
-
-            if not image_file_path.exists():
-                continue  
-
-            shutil.copy(image_file_path, output_directory / split / "images" / image_file_path.name)
-
-            label_file = labels_dir / (image_file_path.stem + ".txt")
-            if label_file.exists():
-                shutil.copy(label_file, output_directory / split / "labels" / label_file.name)
-
-    for split in ["train", "test", "valid"]:
-        create_dirs(split)
-
-    for _, files in class_mapping.items():
-        random.shuffle(files) 
-        num_files = len(files)
-
-        train_count = int(0.8 * num_files)
-        test_count = int(0.1 * num_files)
-        valid_count = num_files - train_count - test_count
-
-        train_files = files[:train_count]
-        test_files = files[train_count:train_count + test_count]
-        valid_files = files[train_count + test_count:]
-
-        copy_files(train_files, "train")
-        copy_files(test_files, "test")
-        copy_files(valid_files, "valid")
-
-    print("Data has been split into train, test, and valid.")
-
-def __save_class_idx_to_file(class_idxs: dict, output_directory: Path):
-    """
-    Saves the class indices to a file.
-
-    Args:
-        class_idxs (dict): A dictionary mapping class names to class indices.
-        output_directory (Path): The path to the output directory where the class index file will be saved.
-    """
-    class_idx_file = output_directory / "class_idx.txt"
-    with open(class_idx_file, 'w') as f:
-        for class_name, idx in class_idxs.items():
-            f.write(f"{class_name}: {idx}\n")
-    print(f"Class indices have been saved to {class_idx_file}")
-
-def __generate_sample_images_with_detections(main_dir: Path, class_idxs: dict):
-
-    """
-    Generates one sample image with multiple detections for each of train, test, valid, combining up to 6 images in one output.
-
-    Args:
-        main_dir (str): The main directory containing the train, test, and valid splits.
-    """
-
-    def resize_and_contain(image, target_size):
-        image.thumbnail(target_size, Image.LANCZOS)
-        new_image = Image.new("RGB", target_size, (0, 0, 0))
-        new_image.paste(image, ((target_size[0] - image.width) // 2, (target_size[1] - image.height) // 2))
-        return new_image
-
-    def draw_bounding_boxes(image, labels_path, class_mapping, color_map):
-        draw = ImageDraw.Draw(image)
-        img_width, img_height = image.size
-        try:
-            font = ImageFont.truetype("DejaVuSans-Bold.ttf", 20)
-        except IOError:
-            font = ImageFont.load_default()
-
-        if labels_path.exists():
-            with open(labels_path, 'r') as label_file:
-                for line in label_file.readlines():
-                    parts = line.strip().split()
-                    class_idx = int(parts[0])
-                    center_x, center_y, width, height = map(float, parts[1:])
-                    x_min = int((center_x - width / 2) * img_width)
-                    y_min = int((center_y - height / 2) * img_height)
-                    x_max = int((center_x + width / 2) * img_width)
-                    y_max = int((center_y + height / 2) * img_height)
-                    class_name = class_mapping.get(class_idx, str(class_idx))
-                    color = color_map[class_idx]
-                    draw.rectangle([x_min, y_min, x_max, y_max], outline=color, width=3)
-                    draw.text((x_min, y_min - 20), class_name, fill=color, font=font)
-        return image
-
-    def combine_images(images, grid_size=(3, 2), target_size=(416, 416)):
-        resized_images = [resize_and_contain(img, target_size) for img in images]
-        width, height = target_size
-        combined_image = Image.new('RGB', (width * grid_size[0], height * grid_size[1]))
-
-        for i, img in enumerate(resized_images):
-            row = i // grid_size[0]
-            col = i % grid_size[0]
-            combined_image.paste(img, (col * width, row * height))
+    
+    # Create train and valid directories
+    train_dir = output_directory / 'train'
+    valid_dir = output_directory / 'valid'
+    
+    train_dir.mkdir(parents=True, exist_ok=True)
+    valid_dir.mkdir(parents=True, exist_ok=True)
+    
+    # Create class directories based on class_mapping
+    print(f"  Creating train and validation directories for {len(class_mapping)} classes...")
+    for class_name in class_mapping:
+        (train_dir / class_name).mkdir(exist_ok=True)
+        (valid_dir / class_name).mkdir(exist_ok=True)
+        print(f"    ✓ Created directories for class: {class_name}")
+    
+    # Process each class folder and its images
+    valid_images = []
+    
+    # First, collect all valid label files
+    valid_label_stems = {label_file.stem for label_file in labels_dir.glob("*.txt") 
+                        if label_file.exists() and os.path.getsize(label_file) > 0}
+    
+    print(f"  Found {len(valid_label_stems)} valid label files for cropping")
+    
+    print("  Starting image cropping and resizing...")
+    total_processed = 0
+    total_valid = 0
+    
+    for class_name, image_names in class_mapping.items():
+        print(f"  Processing class '{class_name}' ({len(image_names)} images)...")
+        class_processed = 0
+        class_valid = 0
         
-        return combined_image
-
-    def generate_color_map(class_mapping):
-        colors = ['red', 'blue', 'green', 'purple', 'orange', 'yellow', 'pink', 'cyan', 'magenta']
-        color_map = {idx: random.choice(colors) for idx in class_mapping.keys()}
-        return color_map
-
-    splits = ['train', 'test', 'valid']
-    class_mapping = class_idxs
-    color_map = generate_color_map(class_mapping)
-
-    for split in splits:
-        images_dir = Path(main_dir) / split / 'images'
-        labels_dir = Path(main_dir) / split / 'labels'
-        image_files = list(images_dir.glob("*.jpg"))
-        if not image_files:
-            continue
+        for image_name in image_names:
+            # Check if the image exists in the images directory
+            image_path = images_dir / image_name
+            class_processed += 1
+            total_processed += 1
+            
+            if not image_path.exists():
+                continue
+                
+            # Skip images that don't have a valid label
+            if image_path.stem not in valid_label_stems:
+                continue
+                
+            label_file = labels_dir / (image_path.stem + '.txt')
+            
+            try:
+                img = Image.open(image_path)
+                
+                if label_file.exists():
+                    # If label exists, crop the image
+                    with open(label_file, 'r') as f:
+                        lines = f.readlines()
+                        if lines:
+                            parts = lines[0].strip().split()
+                            if len(parts) >= 5:
+                                x_center, y_center, width, height = map(float, parts[1:5])
+                                
+                                img_width, img_height = img.size
+                                x_min = int((x_center - width/2) * img_width)
+                                y_min = int((y_center - height/2) * img_height)
+                                x_max = int((x_center + width/2) * img_width)
+                                y_max = int((y_center + height/2) * img_height)
+                                
+                                x_min = max(0, x_min)
+                                y_min = max(0, y_min)
+                                x_max = min(img_width, x_max)
+                                y_max = min(img_height, y_max)
+                                
+                                img = img.crop((x_min, y_min, x_max, y_max))
+                
+                img_width, img_height = img.size
+                if img_width < img_height:
+                    # Width is smaller, set to img_size
+                    new_width = img_size
+                    new_height = int((img_height / img_width) * img_size)
+                else:
+                    # Height is smaller, set to img_size
+                    new_height = img_size
+                    new_width = int((img_width / img_height) * img_size)
+                
+                # Resize the image
+                img = img.resize((new_width, new_height), Image.LANCZOS)
+                
+                valid_images.append((image_path, img, class_name))
+                class_valid += 1
+                total_valid += 1
+            except Exception as e:
+                print(f"    Error processing {image_path}: {e}")
+        
+        print(f"    ✓ Class '{class_name}': {class_valid} valid images from {class_processed} processed")
+    
+    print(f"  ✓ Successfully processed {total_valid} valid images from {total_processed} total images")
+    
+    # Shuffle and split images
+    print("  Shuffling and splitting images into train/validation sets...")
+    random.shuffle(valid_images)
+    split_idx = int(len(valid_images) * 0.9)
+    train_images = valid_images[:split_idx]
+    valid_images_split = valid_images[split_idx:]
+    
+    print(f"  Split: {len(train_images)} training images, {len(valid_images_split)} validation images")
+    
+    # Save images to train/valid directories
+    print("  Saving cropped and resized images...")
+    saved_train = 0
+    saved_valid = 0
+    
+    for image_set, dest_dir, split_name in [(train_images, train_dir, "train"), (valid_images_split, valid_dir, "valid")]:
+        print(f"    Saving {len(image_set)} images to {split_name} set...")
+        for orig_file, img, class_name in image_set:
+            output_path = dest_dir / class_name / (orig_file.stem + '.jpg')
+            
+            # Convert any non-RGB mode to RGB before saving
+            if img.mode != 'RGB':
+                img = img.convert('RGB')
+                
+            img.save(output_path, format='JPEG', quality=95)
+            
+            if split_name == "train":
+                saved_train += 1
+            else:
+                saved_valid += 1
+    
+    print(f"    ✓ Saved {saved_train} train images and {saved_valid} validation images")
+    
+    # Print detailed summary table
+    print(f"  Final dataset summary:")
+    print()
+    
+    # Calculate column widths for proper alignment
+    max_class_name_length = max(len(class_name) for class_name in class_mapping.keys())
+    class_col_width = max(max_class_name_length, len("Class"))
+    
+    # Print table header
+    print(f"  {'Class':<{class_col_width}} | {'Train':<7} | {'Valid':<7} | {'Total':<7}")
+    print(f"  {'-' * class_col_width}-+-{'-' * 7}-+-{'-' * 7}-+-{'-' * 7}")
+    
+    # Print data for each class and calculate totals
+    total_train = 0
+    total_valid = 0
+    total_overall = 0
+    
+    for class_name in sorted(class_mapping.keys()):  # Sort for consistent output
+        train_count = len(list((train_dir / class_name).glob('*.*')))
+        valid_count = len(list((valid_dir / class_name).glob('*.*')))
+        class_total = train_count + valid_count
         
-        sample_images = []
-        for image_file in image_files[:6]:
-            label_file = labels_dir / (image_file.stem + '.txt')
-            image = Image.open(image_file)
-            image_with_boxes = draw_bounding_boxes(image, label_file, class_mapping, color_map)
-            sample_images.append(image_with_boxes)
+        print(f"  {class_name:<{class_col_width}} | {train_count:<7} | {valid_count:<7} | {class_total:<7}")
         
-        if sample_images:
-            combined_image = combine_images(sample_images, grid_size=(3, 2), target_size=(416, 416))
-            combined_image_path = Path(main_dir) / split / f"{split}_sample_with_detections.jpg"
-            combined_image.save(combined_image_path)
+        total_train += train_count
+        total_valid += valid_count
+        total_overall += class_total
     
-
-def __split_background_images(background_dir: Path, output_directory: Path):
-    """
-    Splits the background images into train, test, and validation sets.
-
-    Args:
-        temp_dir_path (Path): The path to the temporary directory containing the background images.
-        output_directory (Path): The path to the output directory where the split background images will be saved.
-    """
-
-    image_files = list(Path(background_dir).glob("*.jpg"))
-    random.shuffle(image_files)
-
-    num_images = len(image_files)
-    train_split = int(0.8 * num_images)
-    valid_split = int(0.1 * num_images)
-
-    train_files = image_files[:train_split]
-    valid_files = image_files[train_split:train_split + valid_split]
-    test_files = image_files[train_split + valid_split:]
-
-    def copy_files(image_list, split):
-        for image_file in image_list:
-            shutil.copy(image_file, Path(output_directory) / split / 'images' / image_file.name)
-
-            label_file = Path(output_directory) / split / 'labels' / (image_file.stem + ".txt")
-            label_file.touch()  
-
-    copy_files(train_files, 'train')
-    copy_files(valid_files, 'valid')
-    copy_files(test_files, 'test')
-
-    print(f"Background data has been split: {len(train_files)} train, {len(valid_files)} valid, {len(test_files)} test")
+    # Print totals row
+    print(f"  {'-' * class_col_width}-+-{'-' * 7}-+-{'-' * 7}-+-{'-' * 7}")
+    print(f"  {'TOTAL':<{class_col_width}} | {total_train:<7} | {total_valid:<7} | {total_overall:<7}")
+    print()
     
-
-def __count_classes_and_output_table(output_directory: Path, class_idxs: dict):
-    """
-    Counts the number of images per class and outputs a table.
-
-    Args:
-        output_directory (Path): The path to the output directory containing the split data.
-        class_idxs (dict): A dictionary mapping class indices to class names.
-    """
-
-    def count_classes_in_split(labels_dir):
-        class_counts = defaultdict(int)
-        for label_file in os.listdir(labels_dir):
-            if label_file.endswith(".txt"):
-                label_path = os.path.join(labels_dir, label_file)
-                with open(label_path, 'r') as f:
-                    lines = f.readlines()
-                    if not lines:
-                        # Count empty files as 'null' class (background images)
-                        class_counts['null'] += 1
-                    else:
-                        for line in lines:
-                            class_index = int(line.split()[0])
-                            class_counts[class_index] += 1
-        return class_counts
-
-    splits = ['train', 'test', 'valid']
-    total_counts = defaultdict(int)
-
-    table = PrettyTable()
-    table.field_names = ["Class", "Class Index", "Train Count", "Test Count", "Valid Count", "Total"]
-
-    split_counts = {split: defaultdict(int) for split in splits}
-
-    for split in splits:
-        labels_dir = output_directory / split / 'labels'
-        if not os.path.exists(labels_dir):
-            print(f"Warning: {labels_dir} does not exist, skipping {split}.")
-            continue
-
-        class_counts = count_classes_in_split(labels_dir)
-        for class_index, count in class_counts.items():
-            split_counts[split][class_index] = count
-            total_counts[class_index] += count
-
-    for class_index, total in total_counts.items():
-        class_name = class_idxs.get(class_index, "Background" if class_index == 'null' else f"Class {class_index}")
-        train_count = split_counts['train'].get(class_index, 0)
-        test_count = split_counts['test'].get(class_index, 0)
-        valid_count = split_counts['valid'].get(class_index, 0)
-        table.add_row([class_name, class_index, train_count, test_count, valid_count, total])
-
-    print(table)
-
-def update_labels(class_mapping: dict, labels_path: Path) -> dict:
-    """
-    Updates the labels based on the class mapping.
-
-    Args:
-        class_mapping (dict): A dictionary mapping class names to image file names.
-        labels_path (Path): The path to the directory containing the label files.
-
-    Returns:
-        dict: A dictionary mapping class names to class indices.
-    """
-    class_index_mapping = {}
-    class_index_definition = {}
-
-    for idx, (class_name, images) in enumerate(class_mapping.items()):
-        class_index_definition[idx] = class_name
-        for image_name in images:
-            class_index_mapping[image_name] = idx
-
-    for txt_file in labels_path.glob("*.txt"):
-        image_name_jpg = txt_file.stem + ".jpg"
-        image_name_jpeg = txt_file.stem + ".jpeg"
-
-        if image_name_jpg in class_index_mapping:
-            class_index = class_index_mapping[image_name_jpg]
-        elif image_name_jpeg in class_index_mapping:
-            class_index = class_index_mapping[image_name_jpeg]
-        else:
-            print(f"Warning: No corresponding image found for {txt_file.name}")
-            continue
-
-        with open(txt_file, 'r') as file:
-            lines = file.readlines()
-
-        updated_lines = []
-        for line in lines:
-            parts = line.split()
-            if len(parts) > 0:
-                parts[0] = str(class_index)
-                updated_lines.append(" ".join(parts))
-
-        with open(txt_file, 'w') as file:
-            file.write("\n".join(updated_lines))
-
-    print(f"Labels updated successfully")
-    return class_index_definition
+    print(f"  ✓ Classification dataset created successfully at: {output_directory}")
 
 def count_images_across_splits(output_directory: Path) -> int:
     """
-    Counts the total number of images across train, test, and validation splits.
+    Counts the total number of images across train and validation splits for classification dataset.
 
     Args:
         output_directory (Path): The path to the output directory containing the split data.
@@ -828,33 +617,12 @@ def count_images_across_splits(output_directory: Path) -> int:
         int: The total number of images across all splits.
     """
     total_images = 0
-    for split in ['train', 'test', 'valid']:
-        split_dir = output_directory / split / 'images'
-        total_images += len(list(split_dir.glob("*.jpg"))) + len(list(split_dir.glob("*.jpeg")))
-
-    return total_images
-
-def __make_yaml_file(output_directory: Path, class_idxs: dict):
-    """
-    Creates a YAML configuration file for YOLOv8.
-
-    Args:
-        output_directory (Path): The path to the output directory where the YAML file will be saved.
-        class_idxs (dict): A dictionary mapping class indices to class names.
-    """
-
-    # Define the structure of the YAML file
-    yaml_content = {
-        'path': str(output_directory.resolve()),
-        'train': 'train/images',
-        'val': 'valid/images',
-        'test': 'test/images',
-        'names': {idx: name for idx, name in class_idxs.items()}
-    }
-
-    # Write the YAML content to a file
-    yaml_file_path = output_directory / 'dataset.yaml'
-    with open(yaml_file_path, 'w') as yaml_file:
-        yaml.dump(yaml_content, yaml_file, default_flow_style=False, sort_keys=False)
-
-    print(f"YOLOv8 YAML file created at {yaml_file_path}")
+    for split in ['train', 'valid']:
+        split_dir = output_directory / split
+        if split_dir.exists():
+            # Count all images in all class subdirectories
+            for class_dir in split_dir.iterdir():
+                if class_dir.is_dir():
+                    total_images += len(list(class_dir.glob("*.jpg"))) + len(list(class_dir.glob("*.jpeg")))
+
+    return total_images