ffr

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+

app.py

@@ -1,18 +1,41 @@
 import torch
+import torch.nn as nn
+from torchvision import transforms
+from PIL import Image
+import cv2
 import timm
-import uvicorn
 from fastapi import FastAPI, File, UploadFile
 import shutil
 import os
-import cv2
-from PIL import Image
-from torchvision import transforms
 
+# FastAPI app instance
 app = FastAPI()
 
 # Device configuration
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+# Create model storage directory
+MODEL_DIR = "./models"
+os.makedirs(MODEL_DIR, exist_ok=True)
+
+# Model URL from Hugging Face
+MODEL_URL = "https://huggingface.co/Maddy21/deepfake-detection-api/resolve/main/best_vit_model.pth"
+
+# Define model path
+model_path = os.path.join(MODEL_DIR, "best_vit_model.pth")
+
+# Download model if not already present
+if not os.path.exists(model_path):
+    print("Downloading model...")
+    torch.hub.download_url_to_file(MODEL_URL, model_path)
+    print("Model downloaded successfully.")
+
+# Load the trained ViT model
+model = timm.create_model('vit_large_patch16_224', pretrained=False, num_classes=2)
+model.load_state_dict(torch.load(model_path, map_location=device))
+model.to(device)
+model.eval()
+
 # Define image transformations
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
@@ -20,17 +43,12 @@ transform = transforms.Compose([
     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
 ])
 
-# Load model from Hugging Face Hub
-MODEL_URL = "https://huggingface.co/Maddy21/deepfake-detection-api/blob/main/best_vit_model.pth"
-
-model = timm.create_model('vit_large_patch16_224', pretrained=False, num_classes=2)
-model.load_state_dict(torch.hub.load_state_dict_from_url(MODEL_URL, map_location=device))
-model.to(device)
-model.eval()
-
+# Function to process video and classify frames
 def predict_video(video_path):
     cap = cv2.VideoCapture(video_path)
-    frame_count, real_count, manipulated_count = 0, 0, 0
+    frame_count = 0
+    real_count = 0
+    manipulated_count = 0
 
     while cap.isOpened():
         ret, frame = cap.read()
@@ -51,22 +69,26 @@ def predict_video(video_path):
             manipulated_count += 1
 
     cap.release()
-    return {"frames": frame_count, "real": real_count, "manipulated": manipulated_count, 
-            "result": "Real" if real_count > manipulated_count else "Fake"}
+    result = "Real" if real_count > manipulated_count else "Manipulated"
+    return {"total_frames": frame_count, "real_frames": real_count, "manipulated_frames": manipulated_count, "result": result}
 
+# API Endpoint to check API status
 @app.get("/")
-def home():
+def read_root():
     return {"message": "Deepfake Detection API is running!"}
 
+# API Endpoint to receive and process video
 @app.post("/predict/")
 async def predict(file: UploadFile = File(...)):
     file_path = f"temp_{file.filename}"
+    
+    # Save uploaded video
     with open(file_path, "wb") as buffer:
         shutil.copyfileobj(file.file, buffer)
     
+    # Run prediction
     result = predict_video(file_path)
+
+    # Delete temp file after processing
     os.remove(file_path)
     return result
-
-if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)