Update app.py

app.py

@@ -9,9 +9,11 @@ from PIL import Image
 import base64
 import torch
 import torch.nn.functional as F
-from transformers import ViTImageProcessor, SwinForImageClassification,AutoImageProcessor
+from transformers import ViTImageProcessor, SwinForImageClassification
 import lightning as L
 import uuid
+import cv2
+
 # Set device
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -27,11 +29,7 @@ hyper_params = {
     "label2id": label2id,
 }
 
-
-
-# Load the processor manually
-vit_img_processor = AutoImageProcessor.from_pretrained('microsoft/swin-small-patch4-window7-224')
-
+vit_img_processor = ViTImageProcessor.from_pretrained(hyper_params['MODEL_CKPT'])
 
 class DeepFakeModel(L.LightningModule):
     def __init__(self, hyperparams: dict):
@@ -51,7 +49,7 @@ class DeepFakeModel(L.LightningModule):
 
 # Load trained model
 model = DeepFakeModel(hyper_params)
-state_dict = torch.load("deepfake_new_trained.pth", map_location=torch.device(device))
+state_dict = torch.load("trained_model.pth", map_location=torch.device(device))
 model.load_state_dict(state_dict)
 model.to(device)
 model.eval()
@@ -84,8 +82,51 @@ def preprocess_image(img):
     return img
 
 
+
+# Load the face detector once
+face_net = cv2.dnn.readNetFromCaffe(
+    "deploy.prototxt",
+    "res10_300x300_ssd_iter_140000.caffemodel"
+)
+
+def detect_face_opencv(image: Image.Image) -> bool:
+    """Detect face using OpenCV DNN"""
+    try:
+        # Convert PIL Image to OpenCV format
+        open_cv_image = np.array(image)
+        open_cv_image = open_cv_image[:, :, ::-1].copy()  # RGB to BGR
+        
+        (h, w) = open_cv_image.shape[:2]
+        blob = cv2.dnn.blobFromImage(open_cv_image, 1.0, (300, 300),
+                                     (104.0, 177.0, 123.0))
+
+        face_net.setInput(blob)
+        detections = face_net.forward()
+
+        # Check if any detection has confidence > 0.5
+        for i in range(detections.shape[2]):
+            confidence = detections[0, 0, i, 2]
+            if confidence > 0.5:
+                return True  # Face detected
+        
+        return False  # No face detected
+    except Exception as e:
+        print(f"Face detection error: {e}")
+        return False  # Fail safe: assume no face
+
 def predict_deepfake(image):
     try:
+        # Step 1: Face Detection
+        has_face = detect_face_opencv(image)
+        if not has_face:
+            return {
+                "id": str(uuid.uuid4()),
+                "isDeepfake": None,
+                "confidence": 0.0,
+                "details": "No face detected in the image. Cannot proceed with deepfake analysis."
+            }
+        
+        # Step 2: Deepfake Prediction (your original logic)
         img_tensor = preprocess_image(image)
         with torch.inference_mode():
             logits = model(img_tensor)
@@ -104,17 +145,19 @@ def predict_deepfake(image):
         raise HTTPException(status_code=500, detail=f"Error during prediction: {str(e)}")
 
 
-@app.post("/api/analyze", response_model=AnalysisResult)
-async def analyze_image(file: UploadFile = File(...)):
-    if not file.content_type.startswith("image/"):
-        raise HTTPException(status_code=400, detail="File must be an image")
-    try:
-        contents = await file.read()
-        image = Image.open(io.BytesIO(contents)).convert("RGB")
-        result = predict_deepfake(image)
-        return JSONResponse(content=result)
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")
+
+
+# @app.post("/api/analyze", response_model=AnalysisResult)
+# async def analyze_image(file: UploadFile = File(...)):
+#     if not file.content_type.startswith("image/"):
+#         raise HTTPException(status_code=400, detail="File must be an image")
+#     try:
+#         contents = await file.read()
+#         image = Image.open(io.BytesIO(contents)).convert("RGB")
+#         result = predict_deepfake(image)
+#         return JSONResponse(content=result)
+#     except Exception as e:
+#         raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")
 
 
 @app.post("/api/analyze-base64", response_model=AnalysisResult)
@@ -134,4 +177,7 @@ async def root():
     return {"message": "DeepFake Detector API is running"}
 
 if __name__ == "__main__":
-      uvicorn.run("app:app", host="0.0.0.0", port=7860)
+    # Remove uvicorn.run for Hugging Face Spaces
+    uvicorn.run("app:app", host="0.0.0.0", port=8000, reload=True)
+
+