save_annotated_images.py

#!/usr/bin/env python3
"""
Simple script to save annotated images from the Marine Species API.
No external dependencies required - uses only Python standard library.
"""

import requests
import base64
import json
import time
from pathlib import Path
import sys


def encode_image_to_base64(image_path: str) -> str:
    """Encode an image file to base64 string."""
    try:
        with open(image_path, "rb") as image_file:
            return base64.b64encode(image_file.read()).decode('utf-8')
    except Exception as e:
        print(f"Error encoding image {image_path}: {e}")
        return None


def save_base64_image(base64_string: str, output_path: str) -> bool:
    """Save a base64 encoded image to a file."""
    try:
        # Decode base64 to bytes
        image_bytes = base64.b64decode(base64_string)
        
        # Save directly to file
        with open(output_path, 'wb') as f:
            f.write(image_bytes)
        
        print(f"✅ Saved annotated image: {output_path}")
        return True
        
    except Exception as e:
        print(f"❌ Failed to save image: {e}")
        return False


def process_image(api_url: str, image_path: str, output_dir: str):
    """Process an image and save the annotated result."""
    
    image_name = Path(image_path).stem
    print(f"\n🐟 Processing {Path(image_path).name}")
    print("-" * 50)
    
    # Encode input image
    print("📷 Encoding image...")
    image_b64 = encode_image_to_base64(image_path)
    if not image_b64:
        return None
    
    # API request
    detection_request = {
        "image": image_b64,
        "confidence_threshold": 0.25,
        "iou_threshold": 0.45,
        "image_size": 640,
        "return_annotated_image": True
    }
    
    try:
        print("🔍 Sending detection request...")
        start_time = time.time()
        
        response = requests.post(
            f"{api_url}/api/v1/detect",
            json=detection_request,
            timeout=60
        )
        
        request_time = time.time() - start_time
        print(f"⏱️  Request completed in {request_time:.2f}s")
        
        if response.status_code == 200:
            result = response.json()
            
            detections = result.get('detections', [])
            processing_time = result.get('processing_time', 0)
            annotated_image_b64 = result.get('annotated_image')
            
            print(f"✅ SUCCESS!")
            print(f"   Processing Time: {processing_time:.3f}s")
            print(f"   Detections Found: {len(detections)}")
            
            # Show detections
            if detections:
                print(f"   🎯 Detected Species:")
                for i, detection in enumerate(detections[:5]):
                    species = detection.get('class_name', 'Unknown')
                    confidence = detection.get('confidence', 0)
                    print(f"      {i+1}. {species} ({confidence:.1%})")
            
            # Save annotated image
            if annotated_image_b64:
                output_path = Path(output_dir) / f"{image_name}_annotated.jpg"
                success = save_base64_image(annotated_image_b64, str(output_path))
                
                if success:
                    return str(output_path)
            else:
                print("   ❌ No annotated image returned")
                
        else:
            print(f"❌ Request failed: {response.status_code}")
            
    except Exception as e:
        print(f"❌ Request failed: {e}")
    
    return None


def main():
    """Main function."""
    api_url = "https://seamo-ai-fishapi.hf.space"
    
    print("🐟 Marine Species API - Save Annotated Images")
    print("=" * 60)
    
    # Create output directory
    output_dir = Path("annotated_results")
    output_dir.mkdir(exist_ok=True)
    print(f"📁 Output directory: {output_dir.absolute()}")
    
    # Find test images
    image_dir = Path("docs/gradio/images")
    if not image_dir.exists():
        print(f"\n❌ Image directory not found: {image_dir}")
        return
    
    # Get image files
    image_files = []
    for ext in ['*.png', '*.jpg', '*.jpeg']:
        image_files.extend(image_dir.glob(ext))
    
    if not image_files:
        print(f"\n❌ No images found in {image_dir}")
        return
    
    print(f"📷 Found {len(image_files)} test images")
    
    # Process each image
    saved_images = []
    for image_path in sorted(image_files):
        result_path = process_image(api_url, str(image_path), str(output_dir))
        if result_path:
            saved_images.append(result_path)
        time.sleep(1)
    
    # Summary
    print("\n" + "=" * 60)
    print(f"🎯 Results:")
    print(f"   📷 Processed: {len(image_files)} images")
    print(f"   💾 Saved: {len(saved_images)} annotated images")
    
    if saved_images:
        print(f"\n🖼️  Annotated images saved:")
        for img_path in saved_images:
            print(f"   📄 {img_path}")
        
        print(f"\n💡 To view the annotated images:")
        print(f"   🖥️  macOS: open {output_dir}")
        print(f"   🐧 Linux: xdg-open {output_dir}")
        print(f"   🪟 Windows: explorer {output_dir}")
        print(f"   📂 Or browse to: {output_dir.absolute()}")
        
        print(f"\n🔍 The images show:")
        print(f"   • Bounding boxes around detected marine species")
        print(f"   • Species names and confidence scores")
        print(f"   • Color-coded detection results")


if __name__ == "__main__":
    main()