cli.py

"""
Command Line Interface for Disaster Victim Identification App
"""

import argparse
import sys
import os
import cv2
import json
from pathlib import Path
from datetime import datetime

from detector import VictimDetector, VideoProcessor
from harmonizer import ImageHarmonizer
from utils import ResultsExporter, ConfigManager


def process_image(args):
    """Process a single image"""
    print(f"Processing image: {args.input}")
    
    # Load image
    image = cv2.imread(args.input)
    if image is None:
        print(f"Error: Could not load image {args.input}")
        return False
    
    # Initialize detector
    detector = VictimDetector(model_path=args.model)
    
    # Initialize harmonizer if enabled
    harmonizer = None
    if args.harmonize:
        harmonizer = ImageHarmonizer()
    
    # Process image
    if harmonizer:
        processed_image = harmonizer.process(image)
    else:
        processed_image = image
    
    # Run detection
    detections = detector.detect(processed_image, confidence=args.confidence)
    
    # Print results
    print(f"Detected {len(detections)} victims")
    for i, detection in enumerate(detections, 1):
        x1, y1, x2, y2 = detection['bbox']
        confidence = detection['confidence']
        print(f"Victim {i}: ({x1:.0f}, {y1:.0f}) to ({x2:.0f}, {y2:.0f}) - Confidence: {confidence:.3f}")
    
    # Save results
    if args.output:
        exporter = ResultsExporter()
        
        # Export detection results
        if args.output.endswith('.json'):
            exporter.export_json(detections, args.output)
        elif args.output.endswith('.csv'):
            exporter.export_csv(detections, args.output)
        else:
            # Default to JSON
            exporter.export_json(detections, args.output + '.json')
        
        print(f"Results saved to {args.output}")
    
    # Save annotated image
    if args.save_image:
        output_image = image.copy()
        for detection in detections:
            x1, y1, x2, y2 = detection['bbox']
            confidence = detection['confidence']
            
            # Draw bounding box
            cv2.rectangle(output_image, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), 2)
            
            # Draw label
            label = f"Victim: {confidence:.2f}"
            cv2.putText(output_image, label, (int(x1), int(y1)-10),
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
        
        image_path = args.save_image if args.save_image else f"result_{Path(args.input).stem}.jpg"
        cv2.imwrite(image_path, output_image)
        print(f"Annotated image saved to {image_path}")
    
    return True


def process_video(args):
    """Process a video file"""
    print(f"Processing video: {args.input}")
    
    # Initialize detector
    detector = VictimDetector(model_path=args.model)
    processor = VideoProcessor(detector)
    
    # Process video
    try:
        processor.process_video(
            args.input,
            output_path=args.output,
            confidence=args.confidence,
            show_preview=not args.no_preview
        )
        print("Video processing completed")
        return True
    except Exception as e:
        print(f"Error processing video: {e}")
        return False


def process_camera(args):
    """Process live camera feed"""
    print("Starting camera processing...")
    
    # Initialize detector
    detector = VictimDetector(model_path=args.model)
    processor = VideoProcessor(detector)
    
    # Process camera feed
    try:
        processor.process_camera(camera_id=args.camera_id, confidence=args.confidence)
        return True
    except Exception as e:
        print(f"Error processing camera: {e}")
        return False


def process_batch(args):
    """Process multiple images"""
    print(f"Processing batch from directory: {args.input}")
    
    # Get all image files
    image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff'}
    image_files = []
    
    for file_path in Path(args.input).iterdir():
        if file_path.suffix.lower() in image_extensions:
            image_files.append(str(file_path))
    
    if not image_files:
        print("No image files found in directory")
        return False
    
    print(f"Found {len(image_files)} images")
    
    # Initialize detector
    detector = VictimDetector(model_path=args.model)
    
    # Initialize harmonizer if enabled
    harmonizer = None
    if args.harmonize:
        harmonizer = ImageHarmonizer()
    
    # Process images
    all_results = {}
    exporter = ResultsExporter()
    
    for i, image_path in enumerate(image_files):
        print(f"Processing {i+1}/{len(image_files)}: {Path(image_path).name}")
        
        # Load image
        image = cv2.imread(image_path)
        if image is None:
            print(f"Warning: Could not load {image_path}")
            continue
        
        # Process image
        if harmonizer:
            processed_image = harmonizer.process(image)
        else:
            processed_image = image
        
        # Run detection
        detections = detector.detect(processed_image, confidence=args.confidence)
        all_results[image_path] = detections
        
        print(f"  Detected {len(detections)} victims")
    
    # Save batch results
    if args.output:
        batch_data = {
            'timestamp': datetime.now().isoformat(),
            'total_images': len(image_files),
            'total_detections': sum(len(detections) for detections in all_results.values()),
            'results': all_results
        }
        
        with open(args.output, 'w') as f:
            json.dump(batch_data, f, indent=2)
        
        print(f"Batch results saved to {args.output}")
    
    return True


def main():
    """Main CLI function"""
    parser = argparse.ArgumentParser(
        description="Disaster Victim Identification App - Command Line Interface",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Process a single image
  python cli.py image --input disaster_scene.jpg --output results.json

  # Process a video
  python cli.py video --input disaster_video.mp4 --output processed_video.mp4

  # Process live camera feed
  python cli.py camera --camera-id 0

  # Process batch of images
  python cli.py batch --input ./images/ --output batch_results.json

  # Use custom model
  python cli.py image --input scene.jpg --model path/to/model.pt
        """
    )
    
    subparsers = parser.add_subparsers(dest='command', help='Available commands')
    
    # Image processing command
    image_parser = subparsers.add_parser('image', help='Process a single image')
    image_parser.add_argument('--input', '-i', required=True, help='Input image path')
    image_parser.add_argument('--output', '-o', help='Output results file')
    image_parser.add_argument('--model', '-m', help='Custom model path')
    image_parser.add_argument('--confidence', '-c', type=float, default=0.5, help='Confidence threshold')
    image_parser.add_argument('--harmonize', action='store_true', help='Enable image harmonization')
    image_parser.add_argument('--save-image', help='Save annotated image to specified path')
    
    # Video processing command
    video_parser = subparsers.add_parser('video', help='Process a video file')
    video_parser.add_argument('--input', '-i', required=True, help='Input video path')
    video_parser.add_argument('--output', '-o', help='Output video path')
    video_parser.add_argument('--model', '-m', help='Custom model path')
    video_parser.add_argument('--confidence', '-c', type=float, default=0.5, help='Confidence threshold')
    video_parser.add_argument('--no-preview', action='store_true', help='Disable preview window')
    
    # Camera processing command
    camera_parser = subparsers.add_parser('camera', help='Process live camera feed')
    camera_parser.add_argument('--camera-id', type=int, default=0, help='Camera ID')
    camera_parser.add_argument('--model', '-m', help='Custom model path')
    camera_parser.add_argument('--confidence', '-c', type=float, default=0.5, help='Confidence threshold')
    
    # Batch processing command
    batch_parser = subparsers.add_parser('batch', help='Process multiple images')
    batch_parser.add_argument('--input', '-i', required=True, help='Input directory path')
    batch_parser.add_argument('--output', '-o', help='Output results file')
    batch_parser.add_argument('--model', '-m', help='Custom model path')
    batch_parser.add_argument('--confidence', '-c', type=float, default=0.5, help='Confidence threshold')
    batch_parser.add_argument('--harmonize', action='store_true', help='Enable image harmonization')
    
    # Global arguments
    parser.add_argument('--config', help='Configuration file path')
    parser.add_argument('--verbose', '-v', action='store_true', help='Verbose output')
    
    args = parser.parse_args()
    
    if not args.command:
        parser.print_help()
        return 1
    
    # Load configuration if provided
    if args.config:
        config_manager = ConfigManager(args.config)
        print(f"Loaded configuration from {args.config}")
    
    # Set logging level
    if args.verbose:
        import logging
        logging.basicConfig(level=logging.DEBUG)
    
    # Execute command
    try:
        if args.command == 'image':
            success = process_image(args)
        elif args.command == 'video':
            success = process_video(args)
        elif args.command == 'camera':
            success = process_camera(args)
        elif args.command == 'batch':
            success = process_batch(args)
        else:
            print(f"Unknown command: {args.command}")
            return 1
        
        return 0 if success else 1
        
    except KeyboardInterrupt:
        print("\nOperation cancelled by user")
        return 1
    except Exception as e:
        print(f"Error: {e}")
        if args.verbose:
            import traceback
            traceback.print_exc()
        return 1


if __name__ == "__main__":
    sys.exit(main())