Upload 5 files

app/__init__.py

@@ -1,20 +1,16 @@
-from .configs import (
-    get_classification_model, 
-    get_segmentation_model,
-    model_classes,
-    request_history,
-    IMAGES_DIR,
-    SEGMENTS_DIR
-)
-from .models import PredictionResponse, HistoryItem
+from .configs import get_classification_model, get_segmentation_model, request_history, STORAGE_DIR, IMAGES_DIR, SEGMENTS_DIR
+from .models import UploadResponse, ClassifyRequest, ClassifyResponse, SegmentRequest, SegmentResponse
 
 __all__ = [
     "get_classification_model",
     "get_segmentation_model",
-    "model_classes",
     "request_history",
+    "STORAGE_DIR",
     "IMAGES_DIR",
     "SEGMENTS_DIR",
-    "PredictionResponse",
-    "HistoryItem"
-]
+    "UploadResponse",
+    "ClassifyRequest",
+    "ClassifyResponse",
+    "SegmentRequest",
+    "SegmentResponse",
+]

app/configs.py

@@ -1,150 +1,94 @@
-# app/configs.py
 import os
-import uuid
-from pathlib import Path
-from typing import Dict, Any, List
-from datetime import datetime
-from huggingface_hub import hf_hub_download
 import zipfile
 import shutil
-import tensorflow as tf
-
-os.environ["PYTHONPATH"] = "/app:" + os.environ.get("PYTHONPATH", "")
-
-try:
-    from hf_artifact_repo import HuggingFaceArtifactRepository
-    print("✅ HuggingFace Artifact Repository loaded!")
-except ImportError as e:
-    print(f"⚠️ Could not load HF plugin: {e}")
+import threading
+import tempfile
+from pathlib import Path
+from typing import Dict, Any, List, Optional
+from huggingface_hub import hf_hub_download
 
 HF_TOKEN = os.getenv("HF_TOKEN")
-if not HF_TOKEN:
-    raise ValueError("❌ HF_TOKEN environment variable is not set!")
-os.environ["HF_TOKEN"] = HF_TOKEN
-
-HF_REPO = "omarelrayes/mlflow-artifacts"
-
-CLASSIFIER_ZIP_PATH = "models/classifier_savedmodel.zip"
-SEGMENTER_ZIP_PATH = "models/segmenter_savedmodel.zip"
-
-print("="*60)
-print("🔧 Configuration:")
-print(f"   - HF Repository: {HF_REPO}")
-print(f"   - TensorFlow version: {tf.__version__}")
-print("="*60)
+HF_MODEL_REPO = os.getenv("HF_MODEL_REPO", "omarelrayes/mlflow-artifacts")
+CLASSIFIER_MODEL_PATH = os.getenv("CLASSIFIER_MODEL_PATH", "models/classifier_savedmodel.zip")
+SEGMENTER_MODEL_PATH = os.getenv("SEGMENTER_MODEL_PATH", "models/segmenter_savedmodel.zip")
+MODEL_CACHE_DIR = Path(os.getenv("MODEL_CACHE_DIR", tempfile.gettempdir())) / "savedmodels"
 
 _classification_model = None
 _segmentation_model = None
+_model_lock = threading.Lock()
 
-# Storage directories
-STORAGE_DIR = Path("storage")
-IMAGES_DIR = STORAGE_DIR / "images"
-SEGMENTS_DIR = STORAGE_DIR / "segments"
+CLASSIFIER_THRESHOLD = float(os.getenv("CLASSIFIER_THRESHOLD", "0.5"))
 
-for d in [STORAGE_DIR, IMAGES_DIR, SEGMENTS_DIR]:
-    d.mkdir(parents=True, exist_ok=True)
+request_history: List[Dict[str, Any]] = []
 
-model_classes = {
-    0: "benign",
-    1: "malignant"
-}
 
-request_history: List[Dict[str, Any]] = []
+def sigmoid_to_class(confidence: float) -> str:
+    return "malicious" if confidence >= CLASSIFIER_THRESHOLD else "benign"
 
-def save_image(image_bytes: bytes, filename: str) -> str:
-    """Save uploaded image to storage"""
-    image_path = IMAGES_DIR / filename
-    with open(image_path, 'wb') as f:
-        f.write(image_bytes)
-    return str(image_path)
-
-def save_segmentation(mask_image, request_id: str) -> str:
-    """Save segmentation mask"""
-    segmentation_path = SEGMENTS_DIR / f"{request_id}_mask.png"
-    mask_image.save(segmentation_path, format='PNG')
-    return str(segmentation_path)
-
-def download_and_extract_savedmodel(zip_path_in_repo, extract_dir):
-    print(f"📥 Downloading {zip_path_in_repo}...")
-    
+
+def _download_and_extract(zip_path_in_repo: str, extract_subdir: str):
+    extract_dir = MODEL_CACHE_DIR / extract_subdir
+    if extract_dir.exists():
+        return str(extract_dir)
+
+    print(f"Downloading {zip_path_in_repo} from {HF_MODEL_REPO}...")
     zip_file = hf_hub_download(
-        repo_id=HF_REPO,
+        repo_id=HF_MODEL_REPO,
         filename=zip_path_in_repo,
         repo_type="model",
-        token=HF_TOKEN,
-        cache_dir="/tmp/hf_cache"
+        token=HF_TOKEN or None,
     )
-    
-    print(f"📦 Downloaded to: {zip_file}")
-    
-    if os.path.exists(extract_dir):
+
+    if extract_dir.exists():
         shutil.rmtree(extract_dir)
-    
-    os.makedirs(extract_dir, exist_ok=True)
-    print(f"📂 Extracting to: {extract_dir}")
-    
-    with zipfile.ZipFile(zip_file, 'r') as zipf:
-        zipf.extractall(extract_dir)
-    
-    print(f"✅ Extracted successfully!")
-    return extract_dir
+    extract_dir.mkdir(parents=True, exist_ok=True)
+
+    with zipfile.ZipFile(zip_file, "r") as zf:
+        zf.extractall(extract_dir)
+
+    print(f"Extracted model to {extract_dir}")
+    return str(extract_dir)
+
+
+def _load_tf_model(zip_path: str, subdir: str):
+    import tensorflow as tf
+
+    model_dir = _download_and_extract(zip_path, subdir)
+    loaded = tf.saved_model.load(model_dir)
+    return loaded.signatures["serving_default"]
+
 
 def get_classification_model():
     global _classification_model
-
     if _classification_model is None:
-        print("="*60)
-        print("🔄 Loading Classification Model...")
-        print("="*60)
-
-        try:
-            extract_dir = download_and_extract_savedmodel(
-                CLASSIFIER_ZIP_PATH,
-                "/tmp/savedmodels/classifier"
-            )
-            
-            print(f"🤖 Loading SavedModel from: {extract_dir}")
-            
-            loaded = tf.saved_model.load(extract_dir)
-            _classification_model = loaded.signatures['serving_default']
-            _classification_model._loaded_obj = loaded
-            
-            print("✅ Classification Model Loaded Successfully!")
-        except Exception as e:
-            print(f"❌ Error loading classification model: {e}")
-            import traceback
-            traceback.print_exc()
-            raise
-
+        with _model_lock:
+            if _classification_model is None:
+                print("Loading Classification Model...")
+                _classification_model = _load_tf_model(CLASSIFIER_MODEL_PATH, "classifier")
     return _classification_model
 
+
 def get_segmentation_model():
     global _segmentation_model
-
     if _segmentation_model is None:
-        print("="*60)
-        print("🔄 Loading Segmentation Model...")
-        print("="*60)
-
-        try:
-            extract_dir = download_and_extract_savedmodel(
-                SEGMENTER_ZIP_PATH,
-                "/tmp/savedmodels/segmenter"
-            )
-            
-            print(f"🤖 Loading SavedModel from: {extract_dir}")
-            
-            loaded = tf.saved_model.load(extract_dir)
-            _segmentation_model = loaded.signatures['serving_default']
-            _segmentation_model._loaded_obj = loaded
-            
-            print("✅ Segmentation Model Loaded Successfully!")
-        except Exception as e:
-            print(f"❌ Error loading segmentation model: {e}")
-            import traceback
-            traceback.print_exc()
-            raise
-
+        with _model_lock:
+            if _segmentation_model is None:
+                print("Loading Segmentation Model...")
+                _segmentation_model = _load_tf_model(SEGMENTER_MODEL_PATH, "segmenter")
     return _segmentation_model
 
-print("✅ Models module initialized successfully!")
+
+def get_loaded_versions() -> List[str]:
+    return ["hf_savedmodel"]
+
+
+CLASSIFIER_URI = os.getenv("CLASSIFIER_MODEL_URI", "hf_savedmodel")
+SEGMENTER_URI = os.getenv("SEGMENTER_MODEL_URI", "hf_savedmodel")
+
+STORAGE_DIR = Path(os.getenv("STORAGE_DIR", "storage"))
+IMAGES_DIR = STORAGE_DIR / "images"
+SEGMENTS_DIR = STORAGE_DIR / "segments"
+RESULTS_DIR = STORAGE_DIR / "results"
+
+for dir_path in [STORAGE_DIR, IMAGES_DIR, SEGMENTS_DIR, RESULTS_DIR]:
+    dir_path.mkdir(parents=True, exist_ok=True)

app/main.py

@@ -1,22 +1,41 @@
 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
+from fastapi import FastAPI, UploadFile, File, HTTPException, Request
+from fastapi.responses import FileResponse
+
+from . import configs as config
+from .models import UploadResponse, ClassifyRequest, ClassifyResponse, SegmentRequest, SegmentResponse
+from .services.predictor import classify_image
+from .services.segmenter import segment_image
+from .services.routing import router as model_router
+from .storage import (
+    save_image,
+    get_image_path,
+    get_image_bytes,
+    save_classification_result,
+    get_classification_result,
+    save_segmentation_result,
+    get_segmentation_result,
+)
+from .monitoring.metrics import MetricsMiddleware, metrics_endpoint
 
 app = FastAPI(
     title="AI Image Classification API",
-    version="2.0.0"
+    version="3.1.0"
 )
 
-# =========================
-# Models Status
-# =========================
+app.add_middleware(MetricsMiddleware)
+
+ALLOWED_CONTENT_TYPES = {"image/jpeg", "image/png", "image/jpg"}
+MAX_FILE_SIZE = 10 * 1024 * 1024
+
+
+@app.get("/metrics")
+async def prometheus_metrics(request: Request):
+    return metrics_endpoint(request)
+
+
 @app.get("/models/status")
 async def models_status():
     classification_model = config.get_classification_model()
@@ -25,115 +44,190 @@ async def models_status():
     return {
         "classification_loaded": classification_model is not None,
         "segmentation_loaded": segmentation_model is not None,
+        "ab_testing": model_router.get_ab_status(),
+        "loaded_versions": config.get_loaded_versions(),
         "storage_paths": {
             "images": str(config.IMAGES_DIR),
-            "segments": str(config.SEGMENTS_DIR)
-        }
+            "segments": str(config.SEGMENTS_DIR),
+        },
     }
 
-# =========================
-# Prediction Endpoint
-# =========================
-@app.post("/predict", response_model=PredictionResponse)
-async def predict(
-    background_tasks: BackgroundTasks,
-    file: UploadFile = File(...)
-):
+
+@app.get("/api/ab-status")
+async def ab_testing_status():
+    return model_router.get_ab_status()
+
+
+@app.post("/api/upload", status_code=201, response_model=UploadResponse)
+async def upload_image(file: UploadFile = File(...)):
+    if file.content_type not in ALLOWED_CONTENT_TYPES:
+        raise HTTPException(
+            status_code=400,
+            detail=f"Unsupported content type: {file.content_type}. Allowed: {ALLOWED_CONTENT_TYPES}",
+        )
+
     file_bytes = await file.read()
 
-    # Validation
-    if len(file_bytes) > 10 * 1024 * 1024:
+    if len(file_bytes) > MAX_FILE_SIZE:
         raise HTTPException(status_code=400, detail="File too large (10MB max)")
 
+    ext = Path(file.filename or "image.jpg").suffix or ".jpg"
+    image_id = str(uuid.uuid4())
+
+    save_image(image_id, file_bytes, ext)
+
+    return UploadResponse(
+        image_id=image_id,
+        filename=file.filename or "image.jpg",
+        size_bytes=len(file_bytes),
+        content_type=file.content_type,
+        uploaded_at=datetime.now().isoformat(),
+        url=f"/api/images/{image_id}",
+    )
+
+
+@app.get("/api/images/{image_id}")
+async def get_image(image_id: str):
+    image_path = get_image_path(image_id)
+    if image_path is None:
+        raise HTTPException(status_code=404, detail="Image not found")
+
+    media_type = "image/jpeg"
+    if image_path.suffix.lower() in {".png"}:
+        media_type = "image/png"
+
+    return FileResponse(path=image_path, media_type=media_type)
+
+
+@app.post("/api/classify", response_model=ClassifyResponse)
+async def classify(body: ClassifyRequest):
+    image_bytes = get_image_bytes(body.image_id)
+    if image_bytes is None:
+        raise HTTPException(status_code=404, detail="Image not found")
+
+    existing = get_classification_result(body.image_id)
+    if existing and not body.model_version:
+        return ClassifyResponse(
+            image_id=body.image_id,
+            prediction=existing["prediction"],
+            confidence=existing["confidence"],
+            model_version=existing["model_version"],
+            status=existing.get("status", "completed"),
+        )
+
+    resolved_uri = model_router.resolve_classifier_version(body.model_version)
+
     try:
-        prediction, confidence, image_path = predict_image(
-            file_bytes,
-            file.filename or "image.jpg"
+        prediction, confidence, actual_version = classify_image(
+            image_bytes,
+            model_version=resolved_uri,
         )
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
-    # 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,
+    result = {
         "prediction": prediction,
         "confidence": confidence,
-        "model_version": model_version,
-        "timestamp": datetime.now().isoformat(),
-        "status": "classified",
-        "segmentation_path": None
+        "model_version": actual_version,
+        "status": "completed",
     }
+    save_classification_result(body.image_id, result)
+
+    return ClassifyResponse(
+        image_id=body.image_id,
+        **result,
+    )
+
+
+@app.get("/api/classify/{image_id}", response_model=ClassifyResponse)
+async def get_classify_result(image_id: str):
+    result = get_classification_result(image_id)
+    if result is None:
+        raise HTTPException(status_code=404, detail="Classification result not found")
+
+    return ClassifyResponse(
+        image_id=image_id,
+        prediction=result["prediction"],
+        confidence=result["confidence"],
+        model_version=result["model_version"],
+        status=result.get("status", "completed"),
+    )
+
+
+@app.post("/api/segment", response_model=SegmentResponse)
+async def segment(body: SegmentRequest):
+    image_bytes = get_image_bytes(body.image_id)
+    if image_bytes is None:
+        raise HTTPException(status_code=404, detail="Image not found")
+
+    existing = get_segmentation_result(body.image_id)
+    if existing:
+        return SegmentResponse(
+            image_id=body.image_id,
+            status=existing.get("status", "completed"),
+            masks_shape=existing.get("masks_shape"),
+            max_confidence=existing.get("max_confidence"),
+            result_url=existing.get("result_path"),
+            error=existing.get("error"),
+        )
 
-    config.request_history.append(history_item)
+    seg_result = segment_image(image_bytes)
+    seg_path = save_segmentation_result(body.image_id, seg_result)
 
-    # 🔥 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 SegmentResponse(
+        image_id=body.image_id,
+        status=seg_result.get("status", "completed"),
+        masks_shape=seg_result.get("masks_shape"),
+        max_confidence=seg_result.get("max_confidence"),
+        result_url=seg_path,
+        error=seg_result.get("error"),
+    )
 
-    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("/api/segment/{image_id}", response_model=SegmentResponse)
+async def get_segment_result(image_id: str):
+    result = get_segmentation_result(image_id)
+    if result is None:
+        raise HTTPException(status_code=404, detail="Segmentation result not found")
 
-@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)
+    return SegmentResponse(
+        image_id=image_id,
+        status=result.get("status", "completed"),
+        masks_shape=result.get("masks_shape"),
+        max_confidence=result.get("max_confidence"),
+        result_url=result.get("result_path"),
+        error=result.get("error"),
+    )
 
-    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 !"
+    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",
+        "version": "3.1.0",
         "classification": class_ready,
         "segmentation": seg_ready,
-        "endpoint": "POST /predict",
-        "history": f"GET /history ({len(config.request_history)} records)",
-        "docs": "/docs"
+        "endpoints": {
+            "upload": "POST /api/upload",
+            "classify": "POST /api/classify",
+            "segment": "POST /api/segment",
+            "health": "GET /health",
+            "metrics": "GET /metrics",
+            "ab_status": "GET /api/ab-status",
+            "docs": "/docs",
+        },
+        "models_status": "/models/status",
     }
 
-# =========================
-# Health Check
-# =========================
+
 @app.get("/health")
 async def health():
     return {
         "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.main:app",
-        host="0.0.0.0",
-        port=7860
-    )

app/models.py

@@ -1,24 +1,37 @@
 from pydantic import BaseModel
 from typing import Optional
 
-class PredictionResponse(BaseModel):
-    request_id: str
-    filename: str
-    image_path: str
-    prediction: str
-    confidence: float
-    model_version: str
-    timestamp: str
-    status: str
-    segmentation_path: Optional[str] = None
 
-class HistoryItem(BaseModel):
-    request_id: str
+class UploadResponse(BaseModel):
+    image_id: str
     filename: str
-    image_path: str
+    size_bytes: int
+    content_type: str
+    uploaded_at: str
+    url: str
+
+
+class ClassifyRequest(BaseModel):
+    image_id: str
+    model_version: Optional[str] = None
+
+
+class ClassifyResponse(BaseModel):
+    image_id: str
     prediction: str
     confidence: float
     model_version: str
-    timestamp: str
+    status: str = "completed"
+
+
+class SegmentRequest(BaseModel):
+    image_id: str
+
+
+class SegmentResponse(BaseModel):
+    image_id: str
     status: str
-    segmentation_path: Optional[str] = None
+    masks_shape: Optional[list] = None
+    max_confidence: Optional[float] = None
+    result_url: Optional[str] = None
+    error: Optional[str] = None

app/storage.py

@@ -1,25 +1,93 @@
-from pathlib import Path
 import json
-from typing import List, Dict, Any
-from app.configs import IMAGES_DIR, SEGMENTS_DIR, request_history
+import threading
+from pathlib import Path
+from typing import Optional
+from app.configs import IMAGES_DIR, SEGMENTS_DIR, RESULTS_DIR
+
+_results: dict = {}
+_results_lock = threading.Lock()
+_RESULTS_FILE = RESULTS_DIR / "results.json"
+
+
+def _load_results():
+    global _results
+    if _RESULTS_FILE.exists():
+        with open(_RESULTS_FILE, "r") as f:
+            _results = json.load(f)
+
 
-def save_image(image_bytes: bytes, filename: str) -> str:
-    """Save image to storage"""
-    image_path = IMAGES_DIR / filename
+def _persist_results():
+    with open(_RESULTS_FILE, "w") as f:
+        json.dump(_results, f, indent=2)
+
+
+def save_image(image_id: str, file_bytes: bytes, extension: str = ".jpg") -> Path:
+    image_path = IMAGES_DIR / f"{image_id}{extension}"
     with open(image_path, "wb") as f:
-        f.write(image_bytes)
-    return str(image_path)
+        f.write(file_bytes)
+    return image_path
+
+
+def get_image_path(image_id: str) -> Optional[Path]:
+    for ext in [".jpg", ".jpeg", ".png", ".JPG", ".JPEG", ".PNG"]:
+        path = IMAGES_DIR / f"{image_id}{ext}"
+        if path.exists():
+            return path
+    return None
+
 
-def save_segmentation_mask(mask_image, request_id: str) -> str:
-    """Save segmentation mask as PNG"""
-    seg_filename = f"{request_id}_mask.png"
+def get_image_bytes(image_id: str) -> Optional[bytes]:
+    path = get_image_path(image_id)
+    if path is None:
+        return None
+    with open(path, "rb") as f:
+        return f.read()
+
+
+def save_classification_result(image_id: str, result: dict) -> None:
+    with _results_lock:
+        if image_id not in _results:
+            _results[image_id] = {}
+        _results[image_id]["classification"] = result
+        _persist_results()
+
+
+def get_classification_result(image_id: str) -> Optional[dict]:
+    with _results_lock:
+        if not _results:
+            _load_results()
+        return _results.get(image_id, {}).get("classification")
+
+
+def save_segmentation_result(image_id: str, seg_result: dict) -> str:
+    seg_filename = f"{image_id}_segment.json"
     seg_path = SEGMENTS_DIR / seg_filename
-    mask_image.save(seg_path, format='PNG')
+    with open(seg_path, "w") as f:
+        json.dump(seg_result, f, indent=2)
+
+    with _results_lock:
+        if image_id not in _results:
+            _results[image_id] = {}
+        _results[image_id]["segmentation"] = {
+            "result_path": str(seg_path),
+            "status": seg_result.get("status", "completed"),
+            "masks_shape": seg_result.get("masks_shape"),
+            "max_confidence": seg_result.get("max_confidence"),
+            "error": seg_result.get("error"),
+        }
+        _persist_results()
+
     return str(seg_path)
 
-def update_history(request_id: str, **updates: Any) -> None:
-    """Update history item by request_id"""
-    for item in request_history:
-        if item["request_id"] == request_id:
-            item.update(updates)
-            break
+
+def get_segmentation_result(image_id: str) -> Optional[dict]:
+    with _results_lock:
+        if not _results:
+            _load_results()
+        return _results.get(image_id, {}).get("segmentation")
+
+
+def get_segmentation_result_path(image_id: str) -> Optional[str]:
+    for f in SEGMENTS_DIR.glob(f"{image_id}_segment.json"):
+        return str(f)
+    return None