bplusplus 2.0.1__tar.gz → 2.0.4__tar.gz

{bplusplus-2.0.1 → bplusplus-2.0.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: bplusplus
-Version: 2.0.1
+Version: 2.0.4
 Summary: A simple method to create AI models for biodiversity, with collect and prepare pipeline
 License: MIT
 Author: Titus Venverloo
@@ -128,10 +128,13 @@ bplusplus.prepare(
     output_directory=PREPARED_DATA_DIR,
     img_size=640,        # Target image size for training
     conf=0.6,            # Detection confidence threshold (0-1)
-    valid=0.1,  # Validation split ratio (0-1), set to 0 for no validation
+    valid=0.1,           # Validation split ratio (0-1), set to 0 for no validation
+    blur=None,           # Gaussian blur as fraction of image size (0-1), None = disabled
 )
 ```
 
+**Note:** The `blur` parameter applies Gaussian blur before resizing, which can help reduce noise. Values are relative to image size (e.g., `blur=0.01` means 1% of the smallest dimension). Supported image formats: JPG, JPEG, and PNG.
+
 #### Step 3: Train Model
 Train the hierarchical classification model on your prepared data. The model learns to identify family, genus, and species.
 
@@ -191,11 +194,15 @@ results = bplusplus.inference(
     output_dir=OUTPUT_DIR,
     fps=None,               # None = process all frames
     backbone="resnet50",    # Must match training
+    save_video=True,        # Set to False to skip video rendering (only CSV output)
+    img_size=60,            # Must match training
 )
 
 print(f"Detected {results['tracks']} tracks ({results['confirmed_tracks']} confirmed)")
 ```
 
+**Note:** Set `save_video=False` to skip generating the annotated and debug videos, which speeds up processing when you only need the CSV detection data.
+
 **Custom Detection Configuration:**
 
 For advanced control over detection parameters, provide a YAML config file:

{bplusplus-2.0.1 → bplusplus-2.0.4}/README.md

@@ -86,10 +86,13 @@ bplusplus.prepare(
     output_directory=PREPARED_DATA_DIR,
     img_size=640,        # Target image size for training
     conf=0.6,            # Detection confidence threshold (0-1)
-    valid=0.1,  # Validation split ratio (0-1), set to 0 for no validation
+    valid=0.1,           # Validation split ratio (0-1), set to 0 for no validation
+    blur=None,           # Gaussian blur as fraction of image size (0-1), None = disabled
 )
 ```
 
+**Note:** The `blur` parameter applies Gaussian blur before resizing, which can help reduce noise. Values are relative to image size (e.g., `blur=0.01` means 1% of the smallest dimension). Supported image formats: JPG, JPEG, and PNG.
+
 #### Step 3: Train Model
 Train the hierarchical classification model on your prepared data. The model learns to identify family, genus, and species.
 
@@ -149,11 +152,15 @@ results = bplusplus.inference(
     output_dir=OUTPUT_DIR,
     fps=None,               # None = process all frames
     backbone="resnet50",    # Must match training
+    save_video=True,        # Set to False to skip video rendering (only CSV output)
+    img_size=60,            # Must match training
 )
 
 print(f"Detected {results['tracks']} tracks ({results['confirmed_tracks']} confirmed)")
 ```
 
+**Note:** Set `save_video=False` to skip generating the annotated and debug videos, which speeds up processing when you only need the CSV detection data.
+
 **Custom Detection Configuration:**
 
 For advanced control over detection parameters, provide a YAML config file:

{bplusplus-2.0.1 → bplusplus-2.0.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "bplusplus"
-version = "2.0.1"
+version = "2.0.4"
 description = "A simple method to create AI models for biodiversity, with collect and prepare pipeline"
 authors = ["Titus Venverloo <tvenver@mit.edu>", "Deniz Aydemir <deniz@aydemir.us>", "Orlando Closs <orlandocloss@pm.me>", "Ase Hatveit <aase@mit.edu>"]
 license = "MIT"

bplusplus-2.0.1/src/bplusplus/insect_detector.py → bplusplus-2.0.4/src/bplusplus/detector.py

@@ -1,6 +1,6 @@
 """
-Insect Detection Backend Module
-===============================
+Detection Backend Module
+========================
 
 This module provides motion-based insect detection utilities used by the inference pipeline.
 It is NOT meant to be run directly - use inference.py instead.

{bplusplus-2.0.1 → bplusplus-2.0.4}/src/bplusplus/inference.py

@@ -18,7 +18,7 @@ import requests
 import logging
 
 from .tracker import InsectTracker
-from .insect_detector import (
+from .detector import (
     DEFAULT_DETECTION_CONFIG,
     get_default_config,
     build_detection_params,
@@ -210,7 +210,7 @@ class HierarchicalInsectClassifier(nn.Module):
     def forward(self, x):
         features = self.backbone(x)
         return [branch(features) for branch in self.branches]
-        
+
 
 # ============================================================================
 # VISUALIZATION
@@ -257,8 +257,8 @@ class FrameVisualizer:
         lines = [track_display]
         
         for level, conf_key in [("family", "family_confidence"), 
-                          ("genus", "genus_confidence"), 
-                          ("species", "species_confidence")]:
+                                ("genus", "genus_confidence"), 
+                                ("species", "species_confidence")]:
             if detection_data.get(level):
                 name = detection_data[level]
                 conf = detection_data.get(conf_key, 0)
@@ -307,7 +307,7 @@ class VideoInferenceProcessor:
     and track-based prediction aggregation.
     """
     
-    def __init__(self, species_list, hierarchical_model_path, params, backbone="resnet50"):
+    def __init__(self, species_list, hierarchical_model_path, params, backbone="resnet50", img_size=60):
         """
         Initialize the processor.
         
@@ -316,7 +316,9 @@ class VideoInferenceProcessor:
             hierarchical_model_path: Path to trained model weights
             params: Detection parameters dict
             backbone: ResNet backbone ('resnet18', 'resnet50', 'resnet101')
+            img_size: Image size for classification (should match training)
         """
+        self.img_size = img_size
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         self.species_list = species_list
         self.params = params
@@ -354,8 +356,7 @@ class VideoInferenceProcessor:
         self.model.eval()
         
         self.transform = transforms.Compose([
-            transforms.Resize((768, 768)),
-            transforms.CenterCrop(640),
+            transforms.Resize((self.img_size, self.img_size)),
             transforms.ToTensor(),
             transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
         ])
@@ -463,13 +464,14 @@ class VideoInferenceProcessor:
         
         return fg_mask, frame_detections
     
-    def classify_confirmed_tracks(self, video_path, confirmed_track_ids):
+    def classify_confirmed_tracks(self, video_path, confirmed_track_ids, crops_dir=None):
         """
         Classify only the confirmed tracks by re-reading relevant frames.
         
         Args:
             video_path: Path to original video
             confirmed_track_ids: Set of track IDs that passed topology analysis
+            crops_dir: Optional directory to save cropped frames
             
         Returns:
             dict: track_id -> list of classifications
@@ -480,6 +482,15 @@ class VideoInferenceProcessor:
         
         print(f"\nClassifying {len(confirmed_track_ids)} confirmed tracks...")
         
+        # Setup crops directory if requested
+        if crops_dir:
+            os.makedirs(crops_dir, exist_ok=True)
+            # Create subdirectory for each track
+            for track_id in confirmed_track_ids:
+                track_dir = os.path.join(crops_dir, str(track_id)[:8])
+                os.makedirs(track_dir, exist_ok=True)
+            print(f"  Saving crops to: {crops_dir}")
+        
         # Group detections by frame for confirmed tracks
         frames_to_classify = defaultdict(list)
         for det in self.all_detections:
@@ -518,11 +529,22 @@ class VideoInferenceProcessor:
                 track_classifications[det['track_id']].append(classification)
                 classified_count += 1
                 
+                # Save crop if requested
+                if crops_dir:
+                    track_id = det['track_id']
+                    track_dir = os.path.join(crops_dir, str(track_id)[:8])
+                    crop = frame[int(y1):int(y2), int(x1):int(x2)]
+                    if crop.size > 0:
+                        crop_path = os.path.join(track_dir, f"frame_{target_frame:06d}.jpg")
+                        cv2.imwrite(crop_path, crop)
+                
                 if classified_count % 20 == 0:
                     print(f"  Classified {classified_count} detections...", end='\r')
         
         cap.release()
         print(f"\n✓ Classified {classified_count} detections from {len(confirmed_track_ids)} tracks")
+        if crops_dir:
+            print(f"✓ Saved {classified_count} crops to {crops_dir}")
         
         return track_classifications
     
@@ -699,7 +721,7 @@ class VideoInferenceProcessor:
         print("\n" + "="*60)
         print("🐛 FINAL SUMMARY")
         print("="*60)
-    
+        
         if results:
             print(f"\n✓ CONFIRMED INSECTS ({num_confirmed}):")
             for r in results:
@@ -736,7 +758,7 @@ class VideoInferenceProcessor:
 # VIDEO PROCESSING
 # ============================================================================
 
-def process_video(video_path, processor, output_paths, show_video=False, fps=None):
+def process_video(video_path, processor, output_paths, show_video=False, fps=None, crops_dir=None):
     """
     Process video file with efficient classification (confirmed tracks only).
     
@@ -752,6 +774,7 @@ def process_video(video_path, processor, output_paths, show_video=False, fps=Non
         output_paths: Dict with output file paths
         show_video: Display video while processing
         fps: Target FPS (skip frames if lower than input)
+        crops_dir: Optional directory to save cropped frames for each track
         
     Returns:
         list: Aggregated results
@@ -832,7 +855,7 @@ def process_video(video_path, processor, output_paths, show_video=False, fps=Non
     print("="*60)
     
     if confirmed_track_ids:
-        processor.classify_confirmed_tracks(video_path, confirmed_track_ids)
+        processor.classify_confirmed_tracks(video_path, confirmed_track_ids, crops_dir=crops_dir)
         results = processor.hierarchical_aggregation(confirmed_track_ids)
     else:
         results = []
@@ -840,23 +863,29 @@ def process_video(video_path, processor, output_paths, show_video=False, fps=Non
     # ==========================================================================
     # PHASE 4: Render Videos
     # ==========================================================================
-    print("\n" + "="*60)
-    print("PHASE 4: RENDERING VIDEOS")
-    print("="*60)
-    
-    # Render debug video (all detections, showing confirmed vs unconfirmed)
-    print(f"\nRendering debug video (all detections)...")
-    _render_debug_video(
-        video_path, output_paths["debug_video"],
-        processor, confirmed_track_ids, all_track_info, input_fps
-    )
-    
-    # Render annotated video (confirmed tracks with classifications)
-    print(f"\nRendering annotated video ({len(confirmed_track_ids)} confirmed tracks)...")
-    _render_annotated_video(
-        video_path, output_paths["annotated_video"],
-        processor, confirmed_track_ids, input_fps
-    )
+    # Render videos if requested
+    if "annotated_video" in output_paths or "debug_video" in output_paths:
+        print("\n" + "="*60)
+        print("PHASE 4: RENDERING VIDEOS")
+        print("="*60)
+        
+        # Render debug video (all detections, showing confirmed vs unconfirmed)
+        if "debug_video" in output_paths:
+            print(f"\nRendering debug video (all detections)...")
+            _render_debug_video(
+                video_path, output_paths["debug_video"],
+                processor, confirmed_track_ids, all_track_info, input_fps
+            )
+        
+        # Render annotated video (confirmed tracks with classifications)
+        if "annotated_video" in output_paths:
+            print(f"\nRendering annotated video ({len(confirmed_track_ids)} confirmed tracks)...")
+            _render_annotated_video(
+                video_path, output_paths["annotated_video"],
+                processor, confirmed_track_ids, input_fps
+            )
+    else:
+        print("\n(Video rendering skipped)")
     
     # Save results
     processor.save_results(results, output_paths)
@@ -1050,6 +1079,9 @@ def inference(
     fps=None,
     config=None,
     backbone="resnet50",
+    crops=False,
+    save_video=True,
+    img_size=60,
 ):
     """
     Run inference on a video file.
@@ -1066,20 +1098,24 @@ def inference(
             - dict: config parameters directly
         backbone: ResNet backbone ('resnet18', 'resnet50', 'resnet101').
                   If model checkpoint contains backbone info, it will be used instead.
+        crops: If True, save cropped frames for each classified track
+        save_video: If True, save annotated and debug videos. Defaults to True.
+        img_size: Image size for classification (should match training). Default: 60.
     
     Returns:
         dict: Processing results with output file paths
         
     Generated files in output_dir:
-        - {video_name}_annotated.mp4: Video with detection boxes and paths
-        - {video_name}_debug.mp4: Side-by-side with GMM motion mask
+        - {video_name}_annotated.mp4: Video with detection boxes and paths (if save_video=True)
+        - {video_name}_debug.mp4: Side-by-side with GMM motion mask (if save_video=True)
         - {video_name}_results.csv: Aggregated track results
         - {video_name}_detections.csv: Frame-by-frame detections
+        - {video_name}_crops/ (if crops=True): Directory with cropped frames per track
     """
     if not os.path.exists(video_path):
         print(f"Error: Video not found: {video_path}")
         return {"error": f"Video not found: {video_path}", "success": False}
-
+    
     # Build parameters from config
     if config is None:
         params = get_default_config()
@@ -1096,16 +1132,23 @@ def inference(
         raise ValueError("config must be None, a file path (str), or a dict")
     
     # Setup output directory and file paths
-        os.makedirs(output_dir, exist_ok=True)
+    os.makedirs(output_dir, exist_ok=True)
     video_name = os.path.splitext(os.path.basename(video_path))[0]
     
     output_paths = {
-        "annotated_video": os.path.join(output_dir, f"{video_name}_annotated.mp4"),
-        "debug_video": os.path.join(output_dir, f"{video_name}_debug.mp4"),
         "results_csv": os.path.join(output_dir, f"{video_name}_results.csv"),
         "detections_csv": os.path.join(output_dir, f"{video_name}_detections.csv"),
     }
     
+    if save_video:
+        output_paths["annotated_video"] = os.path.join(output_dir, f"{video_name}_annotated.mp4")
+        output_paths["debug_video"] = os.path.join(output_dir, f"{video_name}_debug.mp4")
+    
+    # Setup crops directory if requested
+    crops_dir = os.path.join(output_dir, f"{video_name}_crops") if crops else None
+    if crops_dir:
+        output_paths["crops_dir"] = crops_dir
+    
     print("\n" + "="*60)
     print("BPLUSPLUS INFERENCE")
     print("="*60)
@@ -1126,14 +1169,16 @@ def inference(
         hierarchical_model_path=hierarchical_model_path,
         params=params,
         backbone=backbone,
+        img_size=img_size,
     )
-
+    
     try:
         results = process_video(
             video_path=video_path,
             processor=processor,
             output_paths=output_paths,
-            fps=fps
+            fps=fps,
+            crops_dir=crops_dir
         )
         
         return {
@@ -1174,6 +1219,7 @@ Output files generated in output directory:
   - {video_name}_debug.mp4: Side-by-side view with GMM motion mask  
   - {video_name}_results.csv: Aggregated track results
   - {video_name}_detections.csv: Frame-by-frame detections
+  - {video_name}_crops/ (with --crops): Cropped frames for each track
         """
     )
     
@@ -1192,6 +1238,12 @@ Output files generated in output directory:
     parser.add_argument('--backbone', '-b', default='resnet50',
                        choices=['resnet18', 'resnet50', 'resnet101'],
                        help='ResNet backbone (default: resnet50, overridden by checkpoint if saved)')
+    parser.add_argument('--crops', action='store_true',
+                       help='Save cropped frames for each classified track')
+    parser.add_argument('--no-video', action='store_true',
+                       help='Skip saving annotated and debug videos')
+    parser.add_argument('--img-size', type=int, default=60,
+                       help='Image size for classification (should match training, default: 60)')
     
     # Detection parameters (override config)
     defaults = DEFAULT_DETECTION_CONFIG
@@ -1267,6 +1319,9 @@ Output files generated in output directory:
         fps=args.fps,
         config=config,
         backbone=args.backbone,
+        crops=args.crops,
+        save_video=not args.no_video,
+        img_size=args.img_size,
     )
     
     if result.get("success"):
@@ -1279,4 +1334,4 @@ Output files generated in output directory:
 
 
 if __name__ == "__main__":
-    main() 
+    main()

{bplusplus-2.0.1 → bplusplus-2.0.4}/src/bplusplus/prepare.py

@@ -7,7 +7,7 @@ from typing import Optional
 
 import requests
 import torch
-from PIL import Image
+from PIL import Image, ImageFilter
 from torch import serialization
 from torch.nn import Module, ModuleDict, ModuleList
 from torch.nn.modules.activation import LeakyReLU, ReLU, SiLU
@@ -27,7 +27,7 @@ from ultralytics.nn.modules.conv import Conv
 from ultralytics.nn.tasks import DetectionModel
 
 
-def prepare(input_directory: str, output_directory: str, img_size: int = 40, conf: float = 0.35, valid: float = 0.1):
+def prepare(input_directory: str, output_directory: str, img_size: int = 40, conf: float = 0.35, valid: float = 0.1, blur: Optional[float] = None):
     """
     Prepares a YOLO classification dataset by performing the following steps:
     1. Copies images from input directory to temporary directory and creates class mapping.
@@ -44,10 +44,16 @@ def prepare(input_directory: str, output_directory: str, img_size: int = 40, con
         conf (float, optional): YOLO detection confidence threshold. Defaults to 0.35.
         valid (float, optional): Fraction of data for validation (0.0 to 1.0). 
                                  0 = no validation split, 0.1 = 10% validation. Defaults to 0.1.
+        blur (float, optional): Gaussian blur as fraction of image size (0.0 to 1.0).
+                                Applied before resizing. 0.01 = 1% of smallest dimension.
+                                None or 0 means no blur. Defaults to None.
     """
     # Validate the valid parameter
     if not 0 <= valid <= 1:
         raise ValueError(f"valid must be between 0 and 1, got {valid}")
+    # Validate the blur parameter
+    if blur is not None and not 0 <= blur <= 1:
+        raise ValueError(f"blur must be between 0 and 1, got {blur}")
     input_directory = Path(input_directory)
     output_directory = Path(output_directory)
 
@@ -62,6 +68,10 @@ def prepare(input_directory: str, output_directory: str, img_size: int = 40, con
         print(f"Validation split: {valid*100:.0f}% validation, {(1-valid)*100:.0f}% training")
     else:
         print("Validation split: disabled (all images to training)")
+    if blur and blur > 0:
+        print(f"Gaussian blur: {blur*100:.1f}% of image size")
+    else:
+        print("Gaussian blur: disabled")
     print()
 
     with tempfile.TemporaryDirectory() as temp_dir:
@@ -106,7 +116,7 @@ def prepare(input_directory: str, output_directory: str, img_size: int = 40, con
         print("-" * 50)
         _finalize_dataset(
             class_mapping, temp_dir_path, output_directory, 
-            class_idxs, original_image_count, img_size, valid
+            class_idxs, original_image_count, img_size, valid, blur
         )
         print("✓ Step 5 completed: Classification dataset ready!")
         print()
@@ -137,7 +147,7 @@ def _setup_directories_and_copy_images(input_directory: Path, temp_dir_path: Pat
         images_names = []
         if folder_directory.is_dir():
             folder_name = folder_directory.name
-            image_files = list(folder_directory.glob("*.jpg"))
+            image_files = list(folder_directory.glob("*.jpg")) + list(folder_directory.glob("*.png"))
             print(f"  Copying {len(image_files)} images from class '{folder_name}'...")
             
             for image_file in image_files:
@@ -149,7 +159,7 @@ def _setup_directories_and_copy_images(input_directory: Path, temp_dir_path: Pat
             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")))
+    original_image_count = len(list(images_path.glob("*.jpg"))) + len(list(images_path.glob("*.jpeg"))) + len(list(images_path.glob("*.png")))
     print(f"  Total images in temporary directory: {original_image_count}")
     
     return class_mapping, original_image_count
@@ -165,9 +175,9 @@ def _prepare_model_and_clean_images(temp_dir_path: Path):
     
     # Clean corrupted images
     print("  Checking for corrupted images...")
-    images_before = len(list(images_path.glob("*.jpg")))
+    images_before = len(list(images_path.glob("*.jpg"))) + len(list(images_path.glob("*.png")))
     __delete_corrupted_images(images_path)
-    images_after = len(list(images_path.glob("*.jpg")))
+    images_after = len(list(images_path.glob("*.jpg"))) + len(list(images_path.glob("*.png")))
     deleted_count = images_before - images_after
     print(f"  ✓ Cleaned {deleted_count} corrupted images ({images_after} images remain)")
     
@@ -208,7 +218,7 @@ def _run_yolo_inference(temp_dir_path: Path, weights_path: Path, conf: float):
         temp_dir_path (Path): Path to the working temp directory.
         weights_path (Path): Path to YOLO weights.
         conf (float): YOLO detection confidence threshold.
-    
+
     Returns:
         Path: labels_path where the generated labels are stored
     """
@@ -221,7 +231,7 @@ def _run_yolo_inference(temp_dir_path: Path, weights_path: Path, conf: float):
         print("  ✓ YOLO model loaded successfully")
         
         # Get list of all image files
-        image_files = list(images_path.glob('*.jpg'))
+        image_files = list(images_path.glob('*.jpg')) + list(images_path.glob('*.png'))
         print(f"  Found {len(image_files)} images to process with YOLO")
         
         # Ensure predict directory exists
@@ -288,13 +298,13 @@ def _cleanup_and_process_labels(temp_dir_path: Path, labels_path: Path, class_ma
     images_path = temp_dir_path / "images"
     
     print("  Cleaning up orphaned images and labels...")
-    images_before = len(list(images_path.glob("*.jpg")))
+    images_before = len(list(images_path.glob("*.jpg"))) + len(list(images_path.glob("*.png")))
     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")))
+    images_after = len(list(images_path.glob("*.jpg"))) + len(list(images_path.glob("*.png")))
     labels_after = len(list(labels_path.glob("*.txt")))
     
     deleted_images = images_before - images_after
@@ -313,15 +323,16 @@ def _cleanup_and_process_labels(temp_dir_path: Path, labels_path: Path, class_ma
 
 def _finalize_dataset(class_mapping: dict, temp_dir_path: Path, output_directory: Path, 
                      class_idxs: dict, original_image_count: int, img_size: int, 
-                     valid_fraction: float = 0.1):
+                     valid_fraction: float = 0.1, blur: Optional[float] = None):
     """
     Finalizes the dataset by creating cropped classification images and splitting into train/valid sets.
     
     Args:
         valid_fraction: Fraction of data for validation (0.0 to 1.0). 0 = no validation split.
+        blur: Gaussian blur as fraction of image size (0-1). None or 0 means no blur.
     """
     # Split data into train/valid with cropped classification images
-    __classification_split(class_mapping, temp_dir_path, output_directory, img_size, valid_fraction)
+    __classification_split(class_mapping, temp_dir_path, output_directory, img_size, valid_fraction, blur)
     
     # Generate final report
     print("  Generating final statistics...")
@@ -345,11 +356,12 @@ def __delete_corrupted_images(images_path: Path):
     it cannot be opened), the function deletes the corrupted image file.
     """
 
-    for image_file in images_path.glob("*.jpg"):
-        try:
-            Image.open(image_file)
-        except IOError:
-            image_file.unlink()
+    for pattern in ["*.jpg", "*.png"]:
+        for image_file in images_path.glob(pattern):
+            try:
+                Image.open(image_file)
+            except IOError:
+                image_file.unlink()
 
 def __download_file_from_github_release(url, dest_path):
 
@@ -399,13 +411,14 @@ def __delete_orphaned_images_and_inferences(images_path: Path, labels_path: Path
     for txt_file in labels_path.glob("*.txt"):
         image_file_jpg = images_path / (txt_file.stem + ".jpg")
         image_file_jpeg = images_path / (txt_file.stem + ".jpeg")
+        image_file_png = images_path / (txt_file.stem + ".png")
 
-        if not (image_file_jpg.exists() or image_file_jpeg.exists()):
+        if not (image_file_jpg.exists() or image_file_jpeg.exists() or image_file_png.exists()):
             # print(f"Deleting {txt_file.name} - No corresponding image file")
             txt_file.unlink()
             
     label_stems = {txt_file.stem for txt_file in labels_path.glob("*.txt")}
-    image_files = list(images_path.glob("*.jpg")) + list(images_path.glob("*.jpeg"))
+    image_files = list(images_path.glob("*.jpg")) + list(images_path.glob("*.jpeg")) + list(images_path.glob("*.png"))
 
     for image_file in image_files:
         if image_file.stem not in label_stems:
@@ -439,6 +452,7 @@ def __delete_invalid_txt_files(images_path: Path, labels_path: Path):
 
             image_file_jpg = images_path / (txt_file.stem + ".jpg")
             image_file_jpeg = images_path / (txt_file.stem + ".jpeg")
+            image_file_png = images_path / (txt_file.stem + ".png")
 
             if image_file_jpg.exists():
                 image_file_jpg.unlink()
@@ -446,11 +460,14 @@ def __delete_invalid_txt_files(images_path: Path, labels_path: Path):
             elif image_file_jpeg.exists():
                 image_file_jpeg.unlink()
                 # print(f"Deleted corresponding image file: {image_file_jpeg.name}")
+            elif image_file_png.exists():
+                image_file_png.unlink()
+                # print(f"Deleted corresponding image file: {image_file_png.name}")
 
 
 
 
-def __classification_split(class_mapping: dict, temp_dir_path: Path, output_directory: Path, img_size: int, valid_fraction: float = 0.1):
+def __classification_split(class_mapping: dict, temp_dir_path: Path, output_directory: Path, img_size: int, valid_fraction: float = 0.1, blur: Optional[float] = None):
     """
     Splits the data into train and validation sets for classification tasks,
     cropping images according to their YOLO labels but preserving original class structure.
@@ -461,6 +478,7 @@ def __classification_split(class_mapping: dict, temp_dir_path: Path, output_dire
         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.
         valid_fraction (float): Fraction of data for validation (0.0 to 1.0). 0 = no validation split.
+        blur (float, optional): Gaussian blur as fraction of image size (0-1). None or 0 means no blur.
     """
     images_dir = temp_dir_path / "images"
     labels_dir = temp_dir_path / "predict" / "labels"
@@ -544,6 +562,12 @@ def __classification_split(class_mapping: dict, temp_dir_path: Path, output_dire
                                 
                                 img = img.crop((x_min, y_min, x_max, y_max))
                 
+                # Apply Gaussian blur if specified (blur is fraction of smallest dimension)
+                if blur and blur > 0:
+                    img_width, img_height = img.size
+                    blur_radius = blur * min(img_width, img_height)
+                    img = img.filter(ImageFilter.GaussianBlur(radius=blur_radius))
+                
                 img_width, img_height = img.size
                 if img_width < img_height:
                     # Width is smaller, set to img_size
@@ -677,6 +701,6 @@ def count_images_across_splits(output_directory: Path) -> int:
             # 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")))
+                    total_images += len(list(class_dir.glob("*.jpg"))) + len(list(class_dir.glob("*.jpeg"))) + len(list(class_dir.glob("*.png")))
 
     return total_images

{bplusplus-2.0.1 → bplusplus-2.0.4}/src/bplusplus/train.py

@@ -74,7 +74,7 @@ def train(batch_size=4, epochs=30, patience=3, img_size=640, data_dir='input', o
             if any(f.lower().endswith(('.jpg', '.jpeg', '.png')) for f in files):
                 return True
         return False
-    
+
     train_dataset = InsectDataset(
         root_dir=train_dir,
         transform=train_transforms or get_transforms(is_training=True, img_size=img_size),
@@ -107,7 +107,7 @@ def train(batch_size=4, epochs=30, patience=3, img_size=640, data_dir='input', o
             shuffle=False,
             num_workers=num_workers
         )
-    
+
     train_loader = DataLoader(
         train_dataset,
         batch_size=batch_size,
@@ -819,7 +819,7 @@ def train_model(model, train_loader, val_loader, criterion, optimizer, level_to_
             print("Validation skipped (no valid data found).")
             print('-' * 60)
             continue
-        
+
         model.eval()
         val_running_loss = 0.0
         val_correct_predictions = [0] * model.num_levels
@@ -899,7 +899,7 @@ def train_model(model, train_loader, val_loader, criterion, optimizer, level_to_
             'backbone': backbone
         }, best_model_path)
         logger.info(f"Saved model (validation skipped) at {best_model_path}")
-    
+
     logger.info("Training completed successfully")
     return model