Initial commit: Wood knot detection model and GUI

split_coco_dataset.py

@@ -0,0 +1,248 @@
+#!/usr/bin/env python3
+"""
+Split the Kaggle wood defects COCO dataset into train/valid/test splits.
+Creates both COCO format and YOLO format annotations.
+
+Usage:
+    python split_coco_dataset.py --input bbox_coco_dataset.json --images IMAGE/ --output dataset_split
+"""
+
+import argparse
+import json
+import random
+import shutil
+from pathlib import Path
+from typing import Dict, List
+
+
+def coco_to_yolo_bbox(bbox: List[float], img_width: int, img_height: int) -> List[float]:
+    """
+    Convert COCO bbox [x, y, width, height] to YOLO format [x_center, y_center, width, height].
+    All values normalized to [0, 1].
+    
+    Args:
+        bbox: COCO format [x_min, y_min, width, height]
+        img_width: Image width in pixels
+        img_height: Image height in pixels
+    
+    Returns:
+        YOLO format [x_center, y_center, width, height] normalized
+    """
+    x_min, y_min, width, height = bbox
+    
+    # Calculate center coordinates
+    x_center = (x_min + width / 2) / img_width
+    y_center = (y_min + height / 2) / img_height
+    
+    # Normalize width and height
+    norm_width = width / img_width
+    norm_height = height / img_height
+    
+    return [x_center, y_center, norm_width, norm_height]
+
+
+def split_coco_dataset(
+    input_json: Path,
+    images_dir: Path,
+    output_dir: Path,
+    train_split: float = 0.8,
+    valid_split: float = 0.1,
+    seed: int = 42
+):
+    """
+    Split COCO dataset into train/valid/test splits.
+
+    Args:
+        input_json: Path to input COCO JSON file
+        images_dir: Directory containing all images
+        output_dir: Output directory for splits
+        train_split: Fraction for training (default 0.8)
+        valid_split: Fraction for validation (default 0.1)
+        seed: Random seed for reproducibility
+    """
+    # Load COCO data
+    with input_json.open('r') as f:
+        data = json.load(f)
+
+    images = data['images']
+    annotations = data['annotations']
+    categories = data['categories']
+
+    # Set random seed for reproducibility
+    random.seed(seed)
+
+    # Shuffle images
+    random.shuffle(images)
+
+    # Calculate split sizes
+    n_images = len(images)
+    n_train = int(n_images * train_split)
+    n_valid = int(n_images * valid_split)
+    n_test = n_images - n_train - n_valid
+
+    print(f"Total images: {n_images}")
+    print(f"Train: {n_train}, Valid: {n_valid}, Test: {n_test}")
+
+    # Create splits
+    splits = {
+        'train': images[:n_train],
+        'valid': images[n_train:n_train + n_valid],
+        'test': images[n_train + n_valid:]
+    }
+
+    # Create output directories
+    output_dir.mkdir(parents=True, exist_ok=True)
+    
+    # Create category ID to index mapping (YOLO uses 0-indexed categories)
+    category_id_to_idx = {cat['id']: idx for idx, cat in enumerate(categories)}
+    
+    # Create image_id to image info mapping
+    image_info = {img['id']: img for img in images}
+
+    for split_name, split_images in splits.items():
+        split_dir = output_dir / split_name
+        split_dir.mkdir(exist_ok=True)
+        
+        # Create labels directory for YOLO format
+        labels_dir = split_dir / 'labels'
+        labels_dir.mkdir(exist_ok=True)
+
+        # Get image IDs for this split
+        split_image_ids = {img['id'] for img in split_images}
+
+        # Filter annotations for this split
+        split_annotations = [
+            ann for ann in annotations
+            if ann['image_id'] in split_image_ids
+        ]
+
+        # Create COCO data for this split
+        split_data = {
+            'images': split_images,
+            'annotations': split_annotations,
+            'categories': categories
+        }
+
+        # Save COCO JSON
+        json_path = split_dir / '_annotations.coco.json'
+        with json_path.open('w') as f:
+            json.dump(split_data, f, indent=2)
+        
+        # Group annotations by image_id for YOLO format
+        annotations_by_image: Dict[int, List] = {}
+        for ann in split_annotations:
+            img_id = ann['image_id']
+            if img_id not in annotations_by_image:
+                annotations_by_image[img_id] = []
+            annotations_by_image[img_id].append(ann)
+
+        # Copy images and create YOLO labels
+        copied = 0
+        for img in split_images:
+            src_path = images_dir / img['file_name']
+            dst_path = split_dir / img['file_name']
+
+            if src_path.exists():
+                shutil.copy2(src_path, dst_path)
+                copied += 1
+                
+                # Create YOLO format label file
+                img_id = img['id']
+                label_file = labels_dir / f"{Path(img['file_name']).stem}.txt"
+                
+                with label_file.open('w') as f:
+                    if img_id in annotations_by_image:
+                        for ann in annotations_by_image[img_id]:
+                            # Convert COCO bbox to YOLO format
+                            yolo_bbox = coco_to_yolo_bbox(
+                                ann['bbox'],
+                                img['width'],
+                                img['height']
+                            )
+                            
+                            # Get category index
+                            cat_idx = category_id_to_idx[ann['category_id']]
+                            
+                            # Write YOLO format: class x_center y_center width height
+                            f.write(f"{cat_idx} {yolo_bbox[0]:.6f} {yolo_bbox[1]:.6f} "
+                                   f"{yolo_bbox[2]:.6f} {yolo_bbox[3]:.6f}\n")
+            else:
+                print(f"Warning: {src_path} not found")
+
+        print(f"{split_name}: {len(split_images)} images, {len(split_annotations)} annotations, {copied} copied")
+    
+    # Create data.yaml for YOLO training
+    data_yaml_path = output_dir / 'data.yaml'
+    data_yaml_content = f"""# YOLO dataset configuration
+path: {output_dir.absolute()}  # dataset root dir
+train: train  # train images (relative to 'path')
+val: valid  # val images (relative to 'path')
+test: test  # test images (relative to 'path')
+
+# Classes
+names:
+"""
+    for idx, cat in enumerate(categories):
+        data_yaml_content += f"  {idx}: {cat['name']}\n"
+    
+    with data_yaml_path.open('w') as f:
+        f.write(data_yaml_content)
+
+    print(f"\nDataset split complete! Saved to: {output_dir}")
+    print(f"Created YOLO format labels in {output_dir}/{{train,valid,test}}/labels/")
+    print(f"Created data.yaml at {data_yaml_path}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Split COCO dataset into train/valid/test")
+    parser.add_argument(
+        "--input",
+        type=Path,
+        default="bbox_coco_dataset.json",
+        help="Input COCO JSON file"
+    )
+    parser.add_argument(
+        "--images",
+        type=Path,
+        default="IMAGE",
+        help="Directory containing images"
+    )
+    parser.add_argument(
+        "--output",
+        type=Path,
+        default="dataset_split",
+        help="Output directory for splits"
+    )
+    parser.add_argument(
+        "--train-split",
+        type=float,
+        default=0.8,
+        help="Training split fraction"
+    )
+    parser.add_argument(
+        "--valid-split",
+        type=float,
+        default=0.1,
+        help="Validation split fraction"
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=42,
+        help="Random seed"
+    )
+
+    args = parser.parse_args()
+
+    split_coco_dataset(
+        input_json=args.input,
+        images_dir=args.images,
+        output_dir=args.output,
+        train_split=args.train_split,
+        valid_split=args.valid_split,
+        seed=args.seed
+    )
+
+
+if __name__ == "__main__":
+    main()