hbhw

app.py

@@ -1,20 +1,17 @@
 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
-from huggingface_hub import hf_hub_download
+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")
-print("Initializing model...")
 
 # Define image transformations
 transform = transforms.Compose([
@@ -23,27 +20,17 @@ transform = transforms.Compose([
     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
 ])
 
-# Hugging Face model repository details
-HF_MODEL_REPO = "Maddy21/vit-deepfake-model"  # Replace with your repo
-HF_MODEL_FILE = "best_vit_model.pth"
-
-# Download the model from Hugging Face
-model_path = hf_hub_download(repo_id=HF_MODEL_REPO, filename=HF_MODEL_FILE)
+# Load model from Hugging Face Hub
+MODEL_URL = "https://huggingface.co/Maddy21/deepfake-detection-api/blob/main/best_vit_model.pth"
 
-# 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.load_state_dict(torch.hub.load_state_dict_from_url(MODEL_URL, map_location=device))
 model.to(device)
 model.eval()
-print("Model loaded successfully.")
 
-# Function to process video and classify frames
 def predict_video(video_path):
-    print("Processing video for deepfake detection...")
     cap = cv2.VideoCapture(video_path)
-    frame_count = 0
-    real_count = 0
-    manipulated_count = 0
+    frame_count, real_count, manipulated_count = 0, 0, 0
 
     while cap.isOpened():
         ret, frame = cap.read()
@@ -64,33 +51,22 @@ 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 if API is running
 @app.get("/")
-def read_root():
+def home():
     return {"message": "Deepfake Detection API is running!"}
 
-# API Endpoint to upload a video and get predictions
 @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)