boc/iom/ai_pipeline/data_collection.py

"""
Data Collection Pipeline
Collect and annotate infrastructure images for training
"""

from typing import Dict, List, Optional
from dataclasses import dataclass
from pathlib import Path
import json
import os


@dataclass
class ImageAnnotation:
    """Image annotation"""
    image_id: str
    filename: str
    width: int
    height: int
    objects: List[Dict]
    scene_type: str
    location: Optional[Dict] = None


class DataCollector:
    """
    Collect and prepare training data
    
    Sources:
    - quiXzoom app submissions
    - Public datasets
    - Synthetic data generation
    - Manual annotation
    """
    
    def __init__(self, output_path: str):
        self.output_path = Path(output_path)
        self.annotations: List[ImageAnnotation] = []
        
        # Create directories
        (self.output_path / "images" / "train").mkdir(parents=True, exist_ok=True)
        (self.output_path / "images" / "val").mkdir(parents=True, exist_ok=True)
        (self.output_path / "images" / "test").mkdir(parents=True, exist_ok=True)
        (self.output_path / "labels" / "train").mkdir(parents=True, exist_ok=True)
        (self.output_path / "labels" / "val").mkdir(parents=True, exist_ok=True)
        (self.output_path / "labels" / "test").mkdir(parents=True, exist_ok=True)
    
    def add_annotation(
        self,
        image_id: str,
        filename: str,
        width: int,
        height: int,
        objects: List[Dict],
        scene_type: str,
        split: str = "train"
    ):
        """Add image annotation"""
        annotation = ImageAnnotation(
            image_id=image_id,
            filename=filename,
            width=width,
            height=height,
            objects=objects,
            scene_type=scene_type
        )
        
        self.annotations.append(annotation)
        
        # Save label file
        label_path = self.output_path / "labels" / split / f"{image_id}.txt"
        with open(label_path, "w") as f:
            for obj in objects:
                # YOLO format: class x_center y_center width height
                f.write(f"{obj['class_id']} {obj['x_center']} {obj['y_center']} {obj['width']} {obj['height']}\n")
    
    def create_data_yaml(self):
        """Create data.yaml for YOLO training"""
        data = {
            "path": str(self.output_path.absolute()),
            "train": "images/train",
            "val": "images/val",
            "test": "images/test",
            "nc": 10,
            "names": [
                "street_light",
                "traffic_sign",
                "bench",
                "trash_can",
                "sidewalk",
                "road",
                "building",
                "bridge",
                "tree",
                "graffiti"
            ]
        }
        
        with open(self.output_path / "data.yaml", "w") as f:
            json.dump(data, f, indent=2)
    
    def get_stats(self) -> Dict:
        """Get dataset statistics"""
        return {
            "total_images": len(self.annotations),
            "total_objects": sum(len(a.objects) for a in self.annotations),
            "scene_types": list(set(a.scene_type for a in self.annotations)),
            "splits": {
                "train": len(list((self.output_path / "labels" / "train").glob("*"))),
                "val": len(list((self.output_path / "labels" / "val").glob("*"))),
                "test": len(list((self.output_path / "labels" / "test").glob("*")))
            }
        }


class SyntheticDataGenerator:
    """Generate synthetic training data"""
    
    def __init__(self):
        self.scenes = [
            "street_view", "building_facade", "bridge", "road",
            "sidewalk", "park", "industrial_area", "residential_area"
        ]
        
        self.defects = [
            "crack", "corrosion", "graffiti", "damage",
            "wear", "dirt", "missing", "broken"
        ]
    
    def generate_synthetic_image(
        self,
        scene_type: str,
        num_defects: int = 3
    ) -> Dict:
        """Generate synthetic image annotation"""
        import random
        
        width = 640
        height = 480
        
        objects = []
        for _ in range(num_defects):
            # Random defect
            defect = random.choice(self.defects)
            
            # Random bounding box
            x_center = random.uniform(0.1, 0.9)
            y_center = random.uniform(0.1, 0.9)
            w = random.uniform(0.05, 0.3)
            h = random.uniform(0.05, 0.3)
            
            objects.append({
                "class_id": random.randint(0, 9),
                "x_center": x_center,
                "y_center": y_center,
                "width": w,
                "height": h,
                "defect_type": defect
            })
        
        return {
            "width": width,
            "height": height,
            "scene_type": scene_type,
            "objects": objects
        }


# Example usage
def example_data_collection():
    """Example: Data collection"""
    print("=== Data Collection ===\n")
    
    collector = DataCollector("/tmp/iom_dataset")
    
    # Add sample annotations
    collector.add_annotation(
        image_id="img_001",
        filename="img_001.jpg",
        width=640,
        height=480,
        objects=[
            {"class_id": 0, "x_center": 0.5, "y_center": 0.3, "width": 0.1, "height": 0.2},
            {"class_id": 5, "x_center": 0.7, "y_center": 0.6, "width": 0.15, "height": 0.1}
        ],
        scene_type="street_view",
        split="train"
    )
    
    collector.add_annotation(
        image_id="img_002",
        filename="img_002.jpg",
        width=640,
        height=480,
        objects=[
            {"class_id": 2, "x_center": 0.4, "y_center": 0.5, "width": 0.2, "height": 0.15}
        ],
        scene_type="park",
        split="train"
    )
    
    # Create data.yaml
    collector.create_data_yaml()
    
    # Get stats
    stats = collector.get_stats()
    print("Dataset statistics:")
    print(f"  Total images: {stats['total_images']}")
    print(f"  Total objects: {stats['total_objects']}")
    print(f"  Scene types: {stats['scene_types']}")
    
    # Generate synthetic data
    print("\nGenerating synthetic data...")
    generator = SyntheticDataGenerator()
    
    for i in range(5):
        synthetic = generator.generate_synthetic_image("street_view")
        print(f"  Synthetic image {i+1}: {len(synthetic['objects'])} objects")
    
    return collector


if __name__ == '__main__':
    example_data_collection()