Update app/main.py

app/main.py

@@ -1,183 +1,139 @@
-# app/main.py
-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
+import uuid
+from pathlib import Path
+from datetime import datetime
+from fastapi import FastAPI, UploadFile, File, BackgroundTasks, HTTPException
+from typing import List
+
+from app import configs as config
+from app.models import PredictionResponse, HistoryItem
+from app.services.predictor import predict_image
+from app.services.segmenter import run_segmentation
 
 app = FastAPI(
-    title=settings.TITLE,
-    description=settings.DESCRIPTION,
-    version=settings.VERSION
+    title="AI Image Classification API",
+    version="2.0.0"
 )
 
-# ======================
-# Helper Functions
-# ======================
-def preprocess_image(file: UploadFile, target_size=(224, 224)) -> np.ndarray:
-    """Preprocess image for model input"""
-    try:
-        contents = file.file.read()
-        image = Image.open(io.BytesIO(contents))
-        
-        if image.mode != 'RGB':
-            image = image.convert('RGB')
-        
-        image = image.resize(target_size)
-        img_array = np.array(image, dtype=np.float32) / 255.0
-        img_array = np.expand_dims(img_array, axis=0)
-        
-        return img_array
-    except Exception as e:
-        raise HTTPException(status_code=400, detail=f"Error processing image: {str(e)}")
+# =========================
+# Models Status
+# =========================
+@app.get("/models/status")
+async def models_status():
+    classification_model = config.get_classification_model()
+    segmentation_model = config.get_segmentation_model()
 
-# ======================
-# Endpoints
-# ======================
-@app.get("/")
-async def root():
-    """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 {
+        "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)
         }
+    }
 
-@app.post("/classify", response_model=Dict[str, Any])
-async def classify(file: UploadFile = File(...)):
-    """Classify skin cancer image"""
-    try:
-        model = get_classification_model()
-        
-        if model is None:
-            raise HTTPException(status_code=503, detail="Classification model not loaded")
-        
-        img_array = preprocess_image(file, target_size=(224, 224))
-        
-        # 🔥 Call SavedModel signature
-        result = model(tf.constant(img_array))
-        
-        # Get prediction
-        if isinstance(result, dict):
-            predictions = list(result.values())[0].numpy()
-        else:
-            predictions = result.numpy()
-        
-        predicted_class = int(np.argmax(predictions[0]))
-        confidence = float(np.max(predictions[0]))
-        class_name = model_classes.get(predicted_class, "unknown")
-        
-        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)}")
+# =========================
+# 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)")
 
-@app.post("/segment")
-async def segment(file: UploadFile = File(...)):
-    """Segment skin cancer image"""
     try:
-        model = get_segmentation_model()
-        
-        if model is None:
-            raise HTTPException(status_code=503, detail="Segmentation model not loaded")
-        
-        contents = file.file.read()
-        original_image = Image.open(io.BytesIO(contents))
-        
-        if original_image.mode != 'RGB':
-            original_image = original_image.convert('RGB')
-        
-        original_size = original_image.size
-        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
-        result = model(tf.constant(img_array))
-        
-        if isinstance(result, dict):
-            mask = list(result.values())[0].numpy()
-        else:
-            mask = result.numpy()
-        
-        mask = mask[0]
-        
-        if mask.ndim == 3:
-            mask = mask[:, :, 0]
-        
-        mask = (mask * 255).astype(np.uint8)
-        mask_image = Image.fromarray(mask).resize(original_size)
-        
-        img_byte_arr = io.BytesIO()
-        mask_image.save(img_byte_arr, format='PNG')
-        img_byte_arr.seek(0)
-        
-        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
+        prediction, confidence, image_path = predict_image(
+            file_bytes,
+            file.filename or "image.jpg"
+        )
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Segmentation error: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
 
-@app.get("/history")
+    # Get model version
+    model_version = "savedmodel-v1"
+
+    # 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 malignant
+    if prediction.lower() == "malignant":
+        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():
-    """Get request 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
+# =========================
+@app.get("/")
+async def root():
+    class_ready = "READY" if config.get_classification_model() else "REQUIRED !!!!"
+    seg_ready = "READY" if config.get_segmentation_model() else "OPTIONAL !"
+
+    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"
+    }
+
+# =========================
+# Health Check
+# =========================
+@app.get("/health")
+async def health():
     return {
-        "total_requests": len(request_history),
-        "history": request_history[-10:]
+        "status": "healthy",
+        "predict_ready": config.get_classification_model() is not None,
+        "total_predictions": len(config.request_history)
     }
 
+# =========================
+# Run Server
+# =========================
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(
+        "app.main:app",
+        host="0.0.0.0",
+        port=7860
+    )