final touch

Dockerfile

@@ -6,4 +6,4 @@ RUN pip install --no-cache-dir torch torchvision --index-url https://download.py
 RUN pip install --no-cache-dir -r requirements.txt
 COPY . .
 EXPOSE 7860
-CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860", "--workers", "1"]

app.py

@@ -17,8 +17,8 @@ logger = logging.getLogger("deep-detect-api")
 # Initialize FastAPI App
 app = FastAPI(
     title="Deep-Detect API",
-    description="Production-grade API for AI vs Real Image Detection using a custom CNN.",
-    version="1.0.0"
+    description="Production-grade API for AI-generated vs Real image detection. Uses a 3-model ensemble (Custom CNN + ViT Deepfake Detection + Deepfake vs Real ViT) with majority voting for robust predictions.",
+    version="2.0.0"
 )
 
 # Enable CORS middleware
@@ -56,13 +56,16 @@ async def predict_image(file: UploadFile = File(...)):
     """
     logger.info(f"Received prediction request. File: {file.filename}")
 
-    # Validate file extension
+    # Validate file is a supported image format
     content_type = file.content_type or ""
-    if not (content_type.startswith("image/") or file.filename.lower().endswith((".png", ".jpg", ".jpeg"))):
+    allowed_extensions = (".png", ".jpg", ".jpeg", ".webp")
+    is_valid_type = content_type.startswith("image/")
+    is_valid_ext = file.filename.lower().endswith(allowed_extensions)
+    if not (is_valid_type or is_valid_ext):
         logger.warning(f"Rejected invalid file format: {file.filename} (Content-Type: {content_type})")
         raise HTTPException(
             status_code=status.HTTP_400_BAD_REQUEST,
-            detail="Uploaded file must be a valid JPEG or PNG image."
+            detail="Uploaded file must be a valid image (JPEG, PNG, or WebP)."
         )
 
     try:

inference.py

@@ -2,6 +2,7 @@ import io
 import os
 import logging
 import warnings
+import threading
 
 # ─── Silence ALL third-party noise before any imports ──────────────────────────
 os.environ["HF_HUB_DISABLE_SYMLINKS_WARNING"] = "1"
@@ -19,33 +20,30 @@ import torch
 from PIL import Image
 from torchvision import transforms
 
-
-# Set up logging for custom inference pipeline
+# Set up logging for custom inference pipeline only
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger("inference")
 
 BASE_DIR = os.path.dirname(os.path.abspath(__file__))
-# Note: The folder is "models", and model is "custom_cnn_standalone.pt"
 MODEL_PATH = os.path.join(BASE_DIR, "models", "custom_cnn_standalone.pt")
 
-# GPU/CPU Device mapping
+# GPU/CPU device — custom model uses same device as HF models
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+hf_device = 0 if torch.cuda.is_available() else -1   # Fix: use GPU for HF models if available
 logger.info(f"Using device: {device} for inference.")
 
 # Hyperparameters matching the Custom CNN Standalone training
 IMG_SIZE = 224
 MEAN = [0.485, 0.456, 0.406]
-STD = [0.229, 0.224, 0.225]
+STD  = [0.229, 0.224, 0.225]
 OPTIMAL_THRESHOLD = 0.5
 
 # Image Preprocessing Transformation Pipeline
-transform = transforms.Compose(
-    [
-        transforms.Resize((IMG_SIZE, IMG_SIZE)),
-        transforms.ToTensor(),
-        transforms.Normalize(mean=MEAN, std=STD),
-    ]
-)
+transform = transforms.Compose([
+    transforms.Resize((IMG_SIZE, IMG_SIZE)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=MEAN, std=STD),
+])
 
 # Load the compiled TorchScript model
 if not os.path.exists(MODEL_PATH):
@@ -65,77 +63,91 @@ except Exception as e:
     raise e
 
 
+# ─── Image Loading (with validation) ──────────────────────────────────────────
+
+def _load_image(image_source: "str | bytes | Image.Image") -> Image.Image:
+    """
+    Load an image from filepath, raw bytes, or an existing PIL Image.
+    Raises ValueError with a clear message on any failure.
+    """
+    try:
+        if isinstance(image_source, Image.Image):
+            return image_source
+        if isinstance(image_source, bytes):
+            return Image.open(io.BytesIO(image_source))
+        if not os.path.exists(image_source):
+            raise FileNotFoundError(f"Image not found: {image_source}")
+        return Image.open(image_source)
+    except Exception as e:
+        raise ValueError(f"Failed to load image: {e}")
+
+
+# ─── Custom Model Inference ────────────────────────────────────────────────────
+
 def _predict_probability(image: Image.Image) -> float:
     """
     Pass the preprocessed image through the loaded Custom CNN model.
     Applies Sigmoid to the output logit to compute probability.
     """
     input_tensor = transform(image.convert("RGB")).unsqueeze(0).to(device)
-
     with torch.no_grad():
         output = model(input_tensor)
-        # Apply sigmoid since model outputs a single class raw logit
         probability = torch.sigmoid(output).item()
         return probability
 
 
-def _load_image(image_source: str | bytes | Image.Image) -> Image.Image:
-    """
-    Load an image from filepath, raw bytes, or an existing PIL Image.
-    """
-    if isinstance(image_source, Image.Image):
-        return image_source
-    if isinstance(image_source, bytes):
-        return Image.open(io.BytesIO(image_source))
-    return Image.open(image_source)
-
+# ─── Hugging Face Models — thread-safe lazy load with sentinel pattern ─────────
+# Sentinel values:
+#   None  → not yet attempted  (first request will trigger load)
+#   False → permanently failed (never retry)
+#   obj   → loaded pipeline
 
-# Sentinel values: None = not yet attempted, False = failed, object = loaded
 _model_2 = None
 _model_3 = None
+_hf_lock = threading.Lock()   # Fix: thread-safe loading
 
 
-def _load_hf_models():
+def _load_hf_models() -> None:
     """
-    Load both Hugging Face models exactly once. Uses a sentinel pattern:
-    - None  → not yet attempted (first request will trigger load)
-    - False → previously failed (skip silently, no retry)
-    - obj   → successfully loaded pipeline
+    Load both HF models exactly once under a mutex.
+    Individual model failures are caught independently so one
+    bad model does not prevent the other from loading.
     """
     global _model_2, _model_3
-    try:
-        from transformers import pipeline
-
-        if _model_2 is None:
-            try:
-                # GAN-focused: ViT fine-tuned on GAN-generated face datasets
-                _model_2 = pipeline(
-                    "image-classification",
-                    model="Wvolf/ViT_Deepfake_Detection",
-                    device=-1,
-                )
-            except Exception:
-                _model_2 = False  # Mark as permanently failed, never retry
-
-        if _model_3 is None:
-            try:
-                # Diffusion-focused: ViT fine-tuned on Stable Diffusion / deepfake datasets
-                _model_3 = pipeline(
-                    "image-classification",
-                    model="dima806/deepfake_vs_real_image_detection",
-                    device=-1,
-                )
-            except Exception:
-                _model_3 = False  # Mark as permanently failed, never retry
-
-    except Exception:
-        if _model_2 is None:
-            _model_2 = False
-        if _model_3 is None:
-            _model_3 = False
-
-
-def _predict_hf_model_2(pil_image) -> str:
+    with _hf_lock:
+        try:
+            from transformers import pipeline
+
+            if _model_2 is None:
+                try:
+                    # GAN-focused: ViT fine-tuned on GAN / deepfake face datasets
+                    _model_2 = pipeline(
+                        "image-classification",
+                        model="Wvolf/ViT_Deepfake_Detection",
+                        device=hf_device,
+                    )
+                except Exception:
+                    _model_2 = False
+
+            if _model_3 is None:
+                try:
+                    # Diffusion-focused: ViT fine-tuned on Stable Diffusion datasets
+                    _model_3 = pipeline(
+                        "image-classification",
+                        model="dima806/deepfake_vs_real_image_detection",
+                        device=hf_device,
+                    )
+                except Exception:
+                    _model_3 = False
+
+        except Exception:
+            if _model_2 is None:
+                _model_2 = False
+            if _model_3 is None:
+                _model_3 = False
+
+
+def _predict_hf_model_2(pil_image: Image.Image) -> str:
     if _model_2 is None:
         _load_hf_models()
     if not _model_2:
@@ -149,7 +161,7 @@ def _predict_hf_model_2(pil_image) -> str:
         return "unknown"
 
 
-def _predict_hf_model_3(pil_image) -> str:
+def _predict_hf_model_3(pil_image: Image.Image) -> str:
     if _model_3 is None:
         _load_hf_models()
     if not _model_3:
@@ -163,8 +175,9 @@ def _predict_hf_model_3(pil_image) -> str:
         return "unknown"
 
 
+# ─── Public API ───────────────────────────────────────────────────────────────
 
-def predict_label(image_source: str | bytes | Image.Image) -> str:
+def predict_label(image_source: "str | bytes | Image.Image") -> str:
     """
     Predict if the image is 'real' or 'ai' using majority voting.
     """
@@ -172,51 +185,61 @@ def predict_label(image_source: str | bytes | Image.Image) -> str:
     return label
 
 
-def predict_with_confidence(image_source: str | bytes | Image.Image) -> tuple[str, float]:
+def predict_with_confidence(image_source: "str | bytes | Image.Image") -> tuple:
     """
     Predict if the image is 'real' or 'ai' and return the confidence percentage.
-    Uses an ensemble majority vote (2 out of 3) executed concurrently for minimal latency,
-    maintaining absolute logging privacy.
+
+    Execution strategy:
+      - All 3 models run concurrently via ThreadPoolExecutor for minimum latency.
+      - Majority vote determines the final label.
+      - Confidence is computed as (winning_votes / total_votes) * 100 — honest
+        and meaningful regardless of how many models are available.
+      - If only 1 model voted (both HF models failed), confidence is capped at
+        60% to reflect reduced certainty from the lack of ensemble.
     """
     from concurrent.futures import ThreadPoolExecutor
-    
+
     image = _load_image(image_source)
 
-    # Execute all 3 model predictions concurrently to achieve minimal latency
+    # Run all 3 models concurrently
     with ThreadPoolExecutor(max_workers=3) as executor:
         future_custom = executor.submit(_predict_probability, image)
-        future_hf2 = executor.submit(_predict_hf_model_2, image)
-        future_hf3 = executor.submit(_predict_hf_model_3, image)
+        future_hf2    = executor.submit(_predict_hf_model_2, image)
+        future_hf3    = executor.submit(_predict_hf_model_3, image)
 
-        # Retrieve predictions in parallel
         probability = future_custom.result()
-        m2_label = future_hf2.result()
-        m3_label = future_hf3.result()
+        m2_label    = future_hf2.result()
+        m3_label    = future_hf3.result()
 
-    # Process custom model results
-    if probability > OPTIMAL_THRESHOLD:
-        custom_label = "real"
-        custom_confidence = probability * 100
-    else:
-        custom_label = "ai"
-        custom_confidence = (1.0 - probability) * 100
+    # Resolve custom model label
+    custom_label = "real" if probability > OPTIMAL_THRESHOLD else "ai"
 
-    # Perform majority voting
+    # Build vote list — only include models that returned a valid prediction
     votes = [custom_label]
     if m2_label != "unknown":
         votes.append(m2_label)
     if m3_label != "unknown":
         votes.append(m3_label)
 
-    ai_votes = votes.count("ai")
-    real_votes = votes.count("real")
-    final_label = "ai" if ai_votes > real_votes else "real"
+    ai_votes    = votes.count("ai")
+    real_votes  = votes.count("real")
+    total_votes = len(votes)
 
-    # Match the output confidence structure of the custom model
-    if final_label == custom_label:
-        final_confidence = custom_confidence
+    if ai_votes > real_votes:
+        final_label = "ai"
+    elif real_votes > ai_votes:
+        final_label = "real"
     else:
-        # Custom model was overridden; return a realistic confidence percentage
-        final_confidence = max(51.0, 100.0 - custom_confidence)
+        # Tie: trust the custom model as the primary authority
+        final_label = custom_label
+
+    # Vote-ratio confidence — honest and meaningful
+    winning_votes   = max(ai_votes, real_votes)
+    vote_confidence = (winning_votes / total_votes) * 100
+
+    # If only 1 model voted (both HF models failed), cap confidence to signal
+    # reduced certainty — the result is not an ensemble, just a single model.
+    if total_votes == 1:
+        vote_confidence = min(vote_confidence, 60.0)
 
-    return final_label, final_confidence
+    return final_label, round(vote_confidence, 2)

requirements.txt

@@ -4,6 +4,6 @@ python-multipart>=0.0.6
 torch>=2.0.0
 torchvision>=0.15.0
 pillow>=9.5.0
-transformers>=4.30.0
-huggingface_hub>=0.15.0
-
+transformers>=4.36.0
+huggingface_hub>=0.20.0
+accelerate>=0.25.0