multiple images intake

api.py

@@ -1,12 +1,17 @@
 """
-EASI Severity Prediction REST API
-==================================
-
+EASI Severity Prediction REST API with Batch Processing
+========================================================
 FastAPI-based REST API for predicting EASI scores from dermatological images.
-Optimized for Hugging Face Spaces deployment.
+Now supports both single and batch image processing!
+
+New Features:
+- POST /predict/batch - Process multiple images in one request
+- Configurable max batch size and timeout
+- Parallel processing for faster batch predictions
 
 Endpoints:
-- POST /predict - Upload image and get EASI predictions
+- POST /predict - Upload single image and get EASI predictions
+- POST /predict/batch - Upload multiple images (up to 10 at once)
 - GET /health - Health check endpoint
 - GET /conditions - Get list of available conditions
 - GET /docs - Interactive API documentation
@@ -16,12 +21,6 @@ pip install fastapi uvicorn python-multipart pillow tensorflow numpy pandas hugg
 
 Run locally:
 uvicorn api:app --host 0.0.0.0 --port 8000 --reload
-
-Deploy to HF Spaces:
-1. Create Space with Docker SDK
-2. Upload this file + Dockerfile + requirements.txt + trained_model/
-3. Accept terms for google/derm-foundation
-4. Space auto-builds!
 """
 
 import os
@@ -30,6 +29,8 @@ import logging
 from typing import List, Dict, Any, Optional
 from io import BytesIO
 from pathlib import Path
+import asyncio
+from concurrent.futures import ThreadPoolExecutor
 
 # Suppress warnings
 os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
@@ -52,31 +53,34 @@ import pickle
 import pandas as pd
 from huggingface_hub import hf_hub_download
 
+# Configuration
+MAX_BATCH_SIZE = 10  # Maximum images per batch request
+BATCH_TIMEOUT = 300  # Timeout in seconds for batch processing
+HF_REPO_ID = "google/derm-foundation"
+DERM_FOUNDATION_PATH = "./derm_foundation/"
+EASI_MODEL_PATH = './trained_model/easi_severity_model_derm_foundation_individual.pkl'
+HF_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_TOKEN")
+
 # Initialize FastAPI app
 app = FastAPI(
     title="EASI Severity Prediction API",
-    description="REST API for predicting EASI scores from skin images. Deployed on Hugging Face Spaces.",
-    version="2.0.0",
+    description="REST API for predicting EASI scores from skin images. Supports single and batch processing.",
+    version="2.1.0",
     docs_url="/docs",
     redoc_url="/redoc"
 )
 
-# CORS middleware for Flutter web/mobile
+# CORS middleware
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  # In production, specify your Flutter app domain
+    allow_origins=["*"],
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
-# Configuration
-HF_REPO_ID = "google/derm-foundation"
-DERM_FOUNDATION_PATH = "./derm_foundation/"
-EASI_MODEL_PATH = './trained_model/easi_severity_model_derm_foundation_individual.pkl'
-
-# HF Spaces automatically injects HF_TOKEN for authenticated users
-HF_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_TOKEN")
+# Thread pool for parallel processing
+executor = ThreadPoolExecutor(max_workers=4)
 
 # Response Models
 class ConditionPrediction(BaseModel):
@@ -101,12 +105,23 @@ class PredictionResponse(BaseModel):
     summary_statistics: Dict[str, float]
     image_info: Dict[str, Any]
 
+class BatchPredictionResponse(BaseModel):
+    success: bool
+    total_images_processed: int
+    successful_predictions: int
+    failed_predictions: int
+    results: List[Optional[PredictionResponse]]
+    errors: List[Optional[str]]
+    processing_time_seconds: float
+
 class HealthResponse(BaseModel):
     status: str
     models_loaded: Dict[str, bool]
     available_conditions: int
     hf_token_configured: bool
     deployment_platform: str
+    batch_processing_enabled: bool
+    max_batch_size: int
     space_info: Optional[Dict[str, str]] = None
 
 class ErrorResponse(BaseModel):
@@ -134,24 +149,14 @@ class DermFoundationNeuralNetwork:
             self.confidence_scaler = model_data['confidence_scaler']
             self.weighted_scaler = model_data['weighted_scaler']
             
-            # Get the original keras model path from pickle
             keras_model_path = model_data['keras_model_path']
             
-            # If the path doesn't exist, try looking in the same directory as the pickle file
             if not os.path.exists(keras_model_path):
                 print(f"Original keras path not found: {keras_model_path}")
-                
-                # Get the directory where the pickle file is located
                 pickle_dir = os.path.dirname(os.path.abspath(filepath))
-                
-                # Extract just the filename, handling both Windows and Unix paths
-                # Replace backslashes with forward slashes first
                 normalized_path = keras_model_path.replace('\\', '/')
                 keras_filename = normalized_path.split('/')[-1]
-                
                 print(f"Extracted filename: {keras_filename}")
-                
-                # Try looking for it in the same directory as the pickle
                 alternative_path = os.path.join(pickle_dir, keras_filename)
                 print(f"Trying alternative path: {alternative_path}")
                 
@@ -160,16 +165,10 @@ class DermFoundationNeuralNetwork:
                     print(f"✓ Found keras model at: {keras_model_path}")
                 else:
                     print(f"✗ Keras model not found at alternative path either")
-                    print(f"Files in {pickle_dir}:")
-                    try:
-                        print(os.listdir(pickle_dir))
-                    except:
-                        pass
                     return False
             else:
                 print(f"✓ Found keras model at original path: {keras_model_path}")
             
-            # Load the keras model
             self.model = tf.keras.models.load_model(keras_model_path)
             print(f"✓ Keras model loaded successfully")
             return True
@@ -253,7 +252,6 @@ class DermFoundationNeuralNetwork:
 def download_derm_foundation_from_hf(output_dir):
     """Download Derm Foundation model from Hugging Face Hub"""
     try:
-        # Get token - on HF Spaces it's auto-injected
         hf_token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_TOKEN")
         
         print("=" * 80)
@@ -264,13 +262,9 @@ def download_derm_foundation_from_hf(output_dir):
             print(f"✓ HF Token found (length: {len(hf_token)})")
         else:
             print("⚠ No HF Token found - attempting anonymous download")
-            print("  Note: If this fails, you need to:")
-            print("  1. Accept terms at https://huggingface.co/google/derm-foundation")
-            print("  2. Add HF_TOKEN to Space secrets")
         
         os.makedirs(output_dir, exist_ok=True)
         
-        # Files to download
         files_to_download = [
             "saved_model.pb",
             "variables/variables.data-00000-of-00001",
@@ -290,7 +284,6 @@ def download_derm_foundation_from_hf(output_dir):
                     resume_download=True
                 )
                 
-                # Verify file exists and get size
                 if os.path.exists(downloaded_path):
                     file_size_mb = os.path.getsize(downloaded_path) / (1024 * 1024)
                     print(f"✓ Downloaded successfully ({file_size_mb:.2f} MB)")
@@ -313,11 +306,6 @@ def download_derm_foundation_from_hf(output_dir):
         print("✗ ERROR DOWNLOADING MODEL")
         print("=" * 80)
         print(f"Error: {str(e)}")
-        print("\nTroubleshooting steps:")
-        print("1. Ensure you've accepted the model terms at:")
-        print("   https://huggingface.co/google/derm-foundation")
-        print("2. Add HF_TOKEN to your Space secrets (Settings → Repository secrets)")
-        print("3. Make sure your token has 'Read access to gated repos' permission")
         import traceback
         traceback.print_exc()
         return False
@@ -477,7 +465,115 @@ def generate_derm_foundation_embedding(model, image):
         
         return embedding_vector
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error generating embedding: {str(e)}")
+        raise Exception(f"Error generating embedding: {str(e)}")
+
+
+def process_single_image_sync(image_bytes: bytes, filename: str) -> Dict[str, Any]:
+    """
+    Synchronous function to process a single image.
+    Returns dict with 'success', 'result', and 'error' keys.
+    """
+    try:
+        # Read and process image
+        original_image = Image.open(BytesIO(image_bytes)).convert('RGB')
+        original_size = original_image.size
+        
+        # Process to 448x448
+        cropped_img = smart_crop_to_square(original_image)
+        processed_img = cropped_img.resize((448, 448), Image.Resampling.LANCZOS)
+        
+        # Generate embedding
+        embedding = generate_derm_foundation_embedding(derm_model, processed_img)
+        
+        # Make prediction
+        predictions = easi_model.predict(embedding)
+        
+        if predictions is None:
+            return {
+                'success': False,
+                'result': None,
+                'error': "Prediction failed - model returned None"
+            }
+        
+        # Calculate EASI scores
+        easi_results, total_easi = calculate_easi_scores(predictions)
+        severity = get_severity_interpretation(total_easi)
+        
+        # Format predicted conditions
+        predicted_conditions = []
+        for i, condition in enumerate(predictions['dermatologist_skin_condition_on_label_name']):
+            prob = predictions['all_condition_probabilities'][condition]
+            conf = predictions['dermatologist_skin_condition_confidence'][i]
+            weight = predictions['weighted_skin_condition_label'][condition]
+            
+            # Find EASI category
+            easi_category = None
+            easi_contribution = 0
+            for cat_key, cat_info in easi_results.items():
+                for contrib in cat_info['contributing_conditions']:
+                    if contrib['condition'] == condition:
+                        easi_category = cat_info['name']
+                        easi_contribution = contrib['individual_score']
+                        break
+            
+            predicted_conditions.append(ConditionPrediction(
+                condition=condition,
+                probability=float(prob),
+                confidence=float(conf),
+                weight=float(weight),
+                easi_category=easi_category,
+                easi_contribution=easi_contribution
+            ))
+        
+        # Summary statistics
+        summary_stats = {
+            "total_conditions": len(predicted_conditions),
+            "average_confidence": float(np.mean(predictions['dermatologist_skin_condition_confidence'])) if predicted_conditions else 0.0,
+            "average_weight": float(np.mean(list(predictions['weighted_skin_condition_label'].values()))) if predicted_conditions else 0.0,
+            "total_weight": float(sum(predictions['weighted_skin_condition_label'].values()))
+        }
+        
+        # Format EASI components
+        easi_components_formatted = {
+            component: EASIComponent(
+                name=result['name'],
+                score=result['score'],
+                contributing_conditions=result['contributing_conditions']
+            )
+            for component, result in easi_results.items()
+        }
+        
+        result = PredictionResponse(
+            success=True,
+            total_easi_score=total_easi,
+            severity_interpretation=severity,
+            easi_components=easi_components_formatted,
+            predicted_conditions=predicted_conditions,
+            summary_statistics=summary_stats,
+            image_info={
+                "original_size": f"{original_size[0]}x{original_size[1]}",
+                "processed_size": "448x448",
+                "filename": filename
+            }
+        )
+        
+        return {
+            'success': True,
+            'result': result,
+            'error': None
+        }
+        
+    except Exception as e:
+        import traceback
+        error_traceback = traceback.format_exc()
+        print(f"Error processing image {filename}: {str(e)}")
+        print(error_traceback)
+        
+        return {
+            'success': False,
+            'result': None,
+            'error': str(e)
+        }
 
 
 # Global model instances
@@ -491,12 +587,11 @@ async def load_models():
     """Load models on startup"""
     global derm_model, easi_model, deployment_platform
     
-    # Force garbage collection before starting
     import gc
     gc.collect()
     
     print("\n" + "=" * 80)
-    print("🚀 STARTING EASI API ON HUGGING FACE SPACES")
+    print("🚀 STARTING EASI API WITH BATCH PROCESSING")
     print("=" * 80)
     
     # Detect if running on HF Spaces
@@ -513,6 +608,8 @@ async def load_models():
         deployment_platform = "local"
         print("📍 Running locally")
     
+    print(f"🔢 Max batch size: {MAX_BATCH_SIZE}")
+    print(f"⏱️  Batch timeout: {BATCH_TIMEOUT}s")
     print("=" * 80)
     
     # Check HF Token
@@ -521,7 +618,6 @@ async def load_models():
         print(f"✓ HF Token configured (length: {len(hf_token)})")
     else:
         print("⚠ No HF Token found")
-        print("  If model download fails, add HF_TOKEN to Space secrets")
     
     print("=" * 80)
     
@@ -537,7 +633,6 @@ async def load_models():
         
         if not success:
             print("\n❌ CRITICAL: Failed to download Derm Foundation model!")
-            print("   API will not function correctly.")
             return
     else:
         print("\n✓ Derm Foundation model found locally (using cache)")
@@ -577,7 +672,6 @@ async def load_models():
             easi_model = None
     else:
         print(f"✗ EASI model not found at: {EASI_MODEL_PATH}")
-        print("  Make sure trained_model/ folder is included in your Space")
     
     # Final status
     print("\n" + "=" * 80)
@@ -585,13 +679,14 @@ async def load_models():
     print("=" * 80)
     print(f"Derm Foundation Model: {'✓ Loaded' if derm_model else '✗ Failed'}")
     print(f"EASI Prediction Model: {'✓ Loaded' if easi_model else '✗ Failed'}")
+    print(f"Batch Processing: ✓ Enabled (max {MAX_BATCH_SIZE} images)")
     print(f"Platform: {deployment_platform}")
     print("=" * 80)
     
     if derm_model and easi_model:
         print("✅ All systems ready! API is operational.")
     else:
-        print("⚠️  WARNING: Some models failed to load. API may not work correctly.")
+        print("⚠️  WARNING: Some models failed to load.")
     
     print("=" * 80 + "\n")
 
@@ -608,14 +703,20 @@ async def root():
     }
     
     return {
-        "message": "EASI Severity Prediction API",
-        "version": "2.0.0",
+        "message": "EASI Severity Prediction API with Batch Processing",
+        "version": "2.1.0",
         "platform": deployment_platform,
         "space_info": space_info,
         "status": "operational" if (derm_model and easi_model) else "degraded",
+        "batch_processing": {
+            "enabled": True,
+            "max_batch_size": MAX_BATCH_SIZE,
+            "timeout_seconds": BATCH_TIMEOUT
+        },
         "endpoints": {
             "health": "/health",
-            "predict": "/predict",
+            "predict": "/predict (single image)",
+            "predict_batch": "/predict/batch (multiple images)",
             "conditions": "/conditions",
             "docs": "/docs",
             "redoc": "/redoc"
@@ -644,6 +745,8 @@ async def health_check():
         "available_conditions": len(easi_model.mlb.classes_) if easi_model else 0,
         "hf_token_configured": (os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_TOKEN")) is not None,
         "deployment_platform": deployment_platform,
+        "batch_processing_enabled": True,
+        "max_batch_size": MAX_BATCH_SIZE,
         "space_info": space_info
     }
 
@@ -695,85 +798,19 @@ async def predict_easi(
         )
     
     try:
-        # Read and process image
+        # Read image bytes
         image_bytes = await file.read()
-        original_image = Image.open(BytesIO(image_bytes)).convert('RGB')
-        original_size = original_image.size
         
-        # Process to 448x448
-        cropped_img = smart_crop_to_square(original_image)
-        processed_img = cropped_img.resize((448, 448), Image.Resampling.LANCZOS)
+        # Process image synchronously
+        result = process_single_image_sync(image_bytes, file.filename)
         
-        # Generate embedding
-        embedding = generate_derm_foundation_embedding(derm_model, processed_img)
-        
-        # Make prediction
-        predictions = easi_model.predict(embedding)
-        
-        if predictions is None:
-            raise HTTPException(status_code=500, detail="Prediction failed - model returned None")
-        
-        # Calculate EASI scores
-        easi_results, total_easi = calculate_easi_scores(predictions)
-        severity = get_severity_interpretation(total_easi)
-        
-        # Format predicted conditions
-        predicted_conditions = []
-        for i, condition in enumerate(predictions['dermatologist_skin_condition_on_label_name']):
-            prob = predictions['all_condition_probabilities'][condition]
-            conf = predictions['dermatologist_skin_condition_confidence'][i]
-            weight = predictions['weighted_skin_condition_label'][condition]
-            
-            # Find EASI category
-            easi_category = None
-            easi_contribution = 0
-            for cat_key, cat_info in easi_results.items():
-                for contrib in cat_info['contributing_conditions']:
-                    if contrib['condition'] == condition:
-                        easi_category = cat_info['name']
-                        easi_contribution = contrib['individual_score']
-                        break
-            
-            predicted_conditions.append(ConditionPrediction(
-                condition=condition,
-                probability=float(prob),
-                confidence=float(conf),
-                weight=float(weight),
-                easi_category=easi_category,
-                easi_contribution=easi_contribution
-            ))
-        
-        # Summary statistics
-        summary_stats = {
-            "total_conditions": len(predicted_conditions),
-            "average_confidence": float(np.mean(predictions['dermatologist_skin_condition_confidence'])) if predicted_conditions else 0.0,
-            "average_weight": float(np.mean(list(predictions['weighted_skin_condition_label'].values()))) if predicted_conditions else 0.0,
-            "total_weight": float(sum(predictions['weighted_skin_condition_label'].values()))
-        }
-        
-        # Format EASI components
-        easi_components_formatted = {
-            component: EASIComponent(
-                name=result['name'],
-                score=result['score'],
-                contributing_conditions=result['contributing_conditions']
+        if not result['success']:
+            raise HTTPException(
+                status_code=500,
+                detail=f"Error processing image: {result['error']}"
             )
-            for component, result in easi_results.items()
-        }
         
-        return PredictionResponse(
-            success=True,
-            total_easi_score=total_easi,
-            severity_interpretation=severity,
-            easi_components=easi_components_formatted,
-            predicted_conditions=predicted_conditions,
-            summary_statistics=summary_stats,
-            image_info={
-                "original_size": f"{original_size[0]}x{original_size[1]}",
-                "processed_size": "448x448",
-                "filename": file.filename
-            }
-        )
+        return result['result']
         
     except HTTPException:
         raise
@@ -789,6 +826,166 @@ async def predict_easi(
         )
 
 
+@app.post("/predict/batch", response_model=BatchPredictionResponse)
+async def predict_easi_batch(
+    files: List[UploadFile] = File(..., description=f"Multiple skin image files (max {MAX_BATCH_SIZE})")
+):
+    """
+    Predict EASI scores from multiple uploaded skin images in parallel.
+    
+    - **files**: List of image files (JPG, JPEG, PNG) - max 10 images per request
+    - Returns: Batch results with individual predictions and errors
+    
+    **Example Usage (Python):**
+    ```python
+    import requests
+    
+    files = [
+        ('files', open('image1.jpg', 'rb')),
+        ('files', open('image2.jpg', 'rb')),
+        ('files', open('image3.jpg', 'rb'))
+    ]
+    
+    response = requests.post('http://localhost:8000/predict/batch', files=files)
+    results = response.json()
+    ```
+    
+    **Example Usage (cURL):**
+    ```bash
+    curl -X POST "http://localhost:8000/predict/batch" \
+      -F "[email protected]" \
+      -F "[email protected]" \
+      -F "[email protected]"
+    ```
+    """
+    
+    import time
+    start_time = time.time()
+    
+    # Validate models loaded
+    if derm_model is None or easi_model is None:
+        error_detail = []
+        if derm_model is None:
+            error_detail.append("Derm Foundation model not loaded")
+        if easi_model is None:
+            error_detail.append("EASI model not loaded")
+        
+        raise HTTPException(
+            status_code=503,
+            detail=f"Models not available: {', '.join(error_detail)}. Check /health endpoint."
+        )
+    
+    # Validate batch size
+    num_files = len(files)
+    if num_files == 0:
+        raise HTTPException(
+            status_code=400,
+            detail="No files provided. Please upload at least one image."
+        )
+    
+    if num_files > MAX_BATCH_SIZE:
+        raise HTTPException(
+            status_code=400,
+            detail=f"Too many files. Maximum batch size is {MAX_BATCH_SIZE}, received {num_files}."
+        )
+    
+    print(f"\n🔄 Processing batch of {num_files} images...")
+    
+    # Validate file types and read all files
+    image_data = []
+    for idx, file in enumerate(files):
+        if not file.content_type or not file.content_type.startswith('image/'):
+            raise HTTPException(
+                status_code=400,
+                detail=f"File {idx+1} ('{file.filename}') is not an image. Received: {file.content_type}"
+            )
+        
+        try:
+            image_bytes = await file.read()
+            image_data.append({
+                'bytes': image_bytes,
+                'filename': file.filename,
+                'index': idx
+            })
+        except Exception as e:
+            raise HTTPException(
+                status_code=400,
+                detail=f"Error reading file {idx+1} ('{file.filename}'): {str(e)}"
+            )
+    
+    # Process images in parallel using thread pool
+    try:
+        loop = asyncio.get_event_loop()
+        
+        # Create tasks for parallel processing
+        tasks = [
+            loop.run_in_executor(
+                executor,
+                process_single_image_sync,
+                img['bytes'],
+                img['filename']
+            )
+            for img in image_data
+        ]
+        
+        # Wait for all tasks with timeout
+        results = await asyncio.wait_for(
+            asyncio.gather(*tasks, return_exceptions=True),
+            timeout=BATCH_TIMEOUT
+        )
+        
+    except asyncio.TimeoutError:
+        raise HTTPException(
+            status_code=504,
+            detail=f"Batch processing timeout after {BATCH_TIMEOUT} seconds. Try reducing batch size."
+        )
+    except Exception as e:
+        import traceback
+        traceback.print_exc()
+        raise HTTPException(
+            status_code=500,
+            detail=f"Error during batch processing: {str(e)}"
+        )
+    
+    # Collect results and errors
+    prediction_results = []
+    error_messages = []
+    successful_count = 0
+    failed_count = 0
+    
+    for idx, result in enumerate(results):
+        if isinstance(result, Exception):
+            # Handle exception during processing
+            prediction_results.append(None)
+            error_messages.append(f"Exception: {str(result)}")
+            failed_count += 1
+            print(f"  ✗ Image {idx+1} failed: {str(result)}")
+        elif result['success']:
+            prediction_results.append(result['result'])
+            error_messages.append(None)
+            successful_count += 1
+            print(f"  ✓ Image {idx+1} processed successfully")
+        else:
+            prediction_results.append(None)
+            error_messages.append(result['error'])
+            failed_count += 1
+            print(f"  ✗ Image {idx+1} failed: {result['error']}")
+    
+    processing_time = time.time() - start_time
+    
+    print(f"✅ Batch complete: {successful_count} successful, {failed_count} failed in {processing_time:.2f}s\n")
+    
+    return BatchPredictionResponse(
+        success=True,
+        total_images_processed=num_files,
+        successful_predictions=successful_count,
+        failed_predictions=failed_count,
+        results=prediction_results,
+        errors=error_messages,
+        processing_time_seconds=round(processing_time, 2)
+    )
+
+
 @app.exception_handler(HTTPException)
 async def http_exception_handler(request, exc):
     """Custom HTTP exception handler"""
@@ -822,7 +1019,10 @@ if __name__ == "__main__":
     import uvicorn
     
     print("=" * 80)
-    print("🚀 Starting EASI API Server")
+    print("🚀 Starting EASI API Server with Batch Processing")
+    print("=" * 80)
+    print(f"Max batch size: {MAX_BATCH_SIZE} images")
+    print(f"Batch timeout: {BATCH_TIMEOUT} seconds")
     print("=" * 80)
     print("Access the API at: http://localhost:8000")
     print("Interactive docs: http://localhost:8000/docs")