Update app/main.py

app/main.py

@@ -1,140 +1,223 @@
-import uuid
-from pathlib import Path
-from datetime import datetime
-from fastapi import FastAPI, UploadFile, File, BackgroundTasks, HTTPException
-from typing import List
-
-from . import configs as config
-from .models import PredictionResponse, HistoryItem
-from .services.predictor import predict_image
-from .services.segmenter import run_segmentation
+from fastapi import FastAPI, File, UploadFile, HTTPException
+from fastapi.responses import JSONResponse
+from typing import Dict, Any
+import numpy as np
+from PIL import Image
+import io
+import tensorflow as tf
+
+from app.configs import (
+    get_classification_model,
+    get_segmentation_model,
+    model_classes,
+    request_history
+)
+from app.core.config import settings
 
 app = FastAPI(
-    title="AI Image Classification API",
-    version="2.0.0"
+    title=settings.TITLE,
+    description=settings.DESCRIPTION,
+    version=settings.VERSION
 )
 
-# =========================
-# Models Status
-# =========================
-@app.get("/models/status")
-async def models_status():
-    classification_model = config.get_classification_model()
-    segmentation_model = config.get_segmentation_model()
-
-    return {
-        "classification_loaded": classification_model is not None,
-        "segmentation_loaded": segmentation_model is not None,
-        "storage_paths": {
-            "images": str(config.IMAGES_DIR),
-            "segments": str(config.SEGMENTS_DIR)
-        }
-    }
-
-# =========================
-# Prediction Endpoint
-# =========================
-@app.post("/predict", response_model=PredictionResponse)
-async def predict(
-    background_tasks: BackgroundTasks,
-    file: UploadFile = File(...)
-):
-    file_bytes = await file.read()
-
-    #  validation
-    if len(file_bytes) > 10 * 1024 * 1024:
-        raise HTTPException(status_code=400, detail="File too large (10MB max)")
-
+# ======================
+# Helper Functions
+# ======================
+def preprocess_image(file: UploadFile, target_size=(224, 224)) -> np.ndarray:
+    """Preprocess image for model input"""
     try:
-        prediction, confidence, image_path = predict_image(
-            file_bytes,
-            file.filename or "image.jpg"
-        )
+        # Read image
+        contents = file.file.read()
+        image = Image.open(io.BytesIO(contents))
+        
+        # Convert to RGB if necessary
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+        
+        # Resize and normalize
+        image = image.resize(target_size)
+        img_array = np.array(image, dtype=np.float32) / 255.0
+        
+        # Add batch dimension
+        img_array = np.expand_dims(img_array, axis=0)
+        
+        return img_array
     except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+        raise HTTPException(status_code=400, detail=f"Error processing image: {str(e)}")
 
-    # get model 
-    model = config.get_classification_model()
-    model_version = getattr(model, "name", "mlflow-production")
-
-    #create history
-    request_id = str(uuid.uuid4())
-
-    history_item = {
-        "request_id": request_id,
-        "filename": file.filename or "image.jpg",
-        "image_path": image_path,
-        "prediction": prediction,
-        "confidence": confidence,
-        "model_version": model_version,
-        "timestamp": datetime.now().isoformat(),
-        "status": "classified",
-        "segmentation_path": None
-    }
-
-    config.request_history.append(history_item)
-
-    # run segmentation in background if needed
-    if prediction == "malicious":
-        background_tasks.add_task(
-            run_segmentation,
-            request_id,
-            Path(image_path).name,
-            image_path
-        )
-
-    return PredictionResponse(**history_item)
-
-# =========================
-# History
-# =========================
-@app.get("/history", response_model=List[HistoryItem])
-async def get_history():
-    return [HistoryItem(**item) for item in config.request_history]
-
-@app.get("/history/{request_id}", response_model=HistoryItem)
-async def get_prediction(request_id: str):
-    for item in config.request_history:
-        if item["request_id"] == request_id:
-            return HistoryItem(**item)
-
-    raise HTTPException(status_code=404, detail="Prediction not found")
-
-# =========================
-# Root
-# =========================
+# ======================
+# Endpoints
+# ======================
 @app.get("/")
 async def root():
-    class_ready = "READY" if config.get_classification_model() else "REQUIRED !!!!"
-    seg_ready = " READY" if config.get_segmentation_model() else "PTIONAL !"
+    """Health check endpoint"""
+    try:
+        class_model = get_classification_model()
+        seg_model = get_segmentation_model()
+        
+        return {
+            "message": "Skin Cancer API is running!",
+            "classification_model": "READY" if class_model else "REQUIRED !!!!",
+            "segmentation_model": "READY" if seg_model else "REQUIRED !!!!",
+            "endpoints": {
+                "/classify": "POST - Classify skin cancer image",
+                "/segment": "POST - Segment skin cancer image",
+                "/docs": "GET - API documentation"
+            }
+        }
+    except Exception as e:
+        return {
+            "message": "Skin Cancer API is running!",
+            "error": str(e),
+            "classification_model": "ERROR",
+            "segmentation_model": "ERROR"
+        }
 
-    return {
-        "service": "AI Image Classification API",
-        "classification": class_ready,
-        "segmentation": seg_ready,
-        "endpoint": "POST /predict",
-        "history": f"GET /history ({len(config.request_history)} records)",
-        "docs": "/docs"
-    }
+@app.post("/classify", response_model=Dict[str, Any])
+async def classify(file: UploadFile = File(...)):
+    """
+    Classify skin cancer image
+    
+    - **file**: Image file (JPG, PNG)
+    """
+    try:
+        # Get model
+        model = get_classification_model()
+        
+        if model is None:
+            raise HTTPException(status_code=503, detail="Classification model not loaded")
+        
+        # Preprocess image
+        img_array = preprocess_image(file, target_size=(224, 224))
+        
+        # 🔥 Call SavedModel signature directly
+        # The signature expects 'input_1' or 'serving_default' input
+        if hasattr(model, 'structured_input_signature'):
+            # Get input tensor name from signature
+            input_tensor = tf.constant(img_array)
+            result = model(input_1=input_tensor)
+        else:
+            # Fallback: call directly
+            result = model(tf.constant(img_array))
+        
+        # Get prediction
+        # SavedModel returns a dict with output keys
+        if isinstance(result, dict):
+            # Get the first output value
+            predictions = list(result.values())[0].numpy()
+        else:
+            predictions = result.numpy()
+        
+        # Get class and confidence
+        predicted_class = int(np.argmax(predictions[0]))
+        confidence = float(np.max(predictions[0]))
+        
+        # Get class name
+        class_name = model_classes.get(predicted_class, "unknown")
+        
+        # Log request
+        request_history.append({
+            "endpoint": "/classify",
+            "prediction": class_name,
+            "confidence": confidence
+        })
+        
+        return {
+            "prediction": class_name,
+            "class_id": predicted_class,
+            "confidence": confidence,
+            "class_probabilities": {
+                model_classes.get(i, "unknown"): float(prob)
+                for i, prob in enumerate(predictions[0])
+            }
+        }
+        
+    except HTTPException:
+        raise
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Classification error: {str(e)}")
+
+@app.post("/segment")
+async def segment(file: UploadFile = File(...)):
+    """
+    Segment skin cancer image
+    
+    - **file**: Image file (JPG, PNG)
+    """
+    try:
+        # Get model
+        model = get_segmentation_model()
+        
+        if model is None:
+            raise HTTPException(status_code=503, detail="Segmentation model not loaded")
+        
+        # Read and preprocess image
+        contents = file.file.read()
+        original_image = Image.open(io.BytesIO(contents))
+        
+        if original_image.mode != 'RGB':
+            original_image = original_image.convert('RGB')
+        
+        # Keep original size for output
+        original_size = original_image.size
+        
+        # Resize for model input
+        img_array = np.array(original_image.resize((256, 256)), dtype=np.float32) / 255.0
+        img_array = np.expand_dims(img_array, axis=0)
+        
+        # 🔥 Call SavedModel signature directly
+        if hasattr(model, 'structured_input_signature'):
+            result = model(input_1=tf.constant(img_array))
+        else:
+            result = model(tf.constant(img_array))
+        
+        # Get segmentation mask
+        if isinstance(result, dict):
+            mask = list(result.values())[0].numpy()
+        else:
+            mask = result.numpy()
+        
+        # Post-process mask
+        mask = mask[0]  # Remove batch dimension
+        
+        # Resize mask to original size
+        if mask.ndim == 3:
+            mask = mask[:, :, 0]  # Take first channel if needed
+        
+        mask = (mask * 255).astype(np.uint8)
+        mask_image = Image.fromarray(mask).resize(original_size)
+        
+        # Convert to bytes
+        img_byte_arr = io.BytesIO()
+        mask_image.save(img_byte_arr, format='PNG')
+        img_byte_arr.seek(0)
+        
+        # Log request
+        request_history.append({
+            "endpoint": "/segment",
+            "status": "success"
+        })
+        
+        return {
+            "message": "Segmentation completed",
+            "output_format": "PNG",
+            "original_size": original_size,
+            "mask_size": mask_image.size
+        }
+        
+    except HTTPException:
+        raise
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"Segmentation error: {str(e)}")
 
-# =========================
-# Health Check
-# =========================
-@app.get("/health")
-async def health():
+@app.get("/history")
+async def get_history():
+    """Get request history"""
     return {
-        "status": "healthy",
-        "predict_ready": config.get_classification_model() is not None,
-        "total_predictions": len(config.request_history)
+        "total_requests": len(request_history),
+        "history": request_history[-10:]  # Last 10 requests
     }
 
-# =========================
-# Run Server
-# =========================
-if __name__ == "main":
+if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(
-    "app.main:app",
-    host="0.0.0.0",
-    port=7860
-)
+    uvicorn.run(app, host="0.0.0.0", port=7860)