truens66 commited on
Commit
e5a0b5e
·
verified ·
1 Parent(s): a6b0e12

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +23 -6
app.py CHANGED
@@ -6,16 +6,27 @@ import numpy as np
6
  from imutils import face_utils
7
  from torchvision import models, transforms
8
  from tempfile import NamedTemporaryFile
9
-
10
  # Load face detector and landmark predictor
11
  face_detector = dlib.get_frontal_face_detector()
12
- PREDICTOR_PATH = "./lib/shape_predictor_81_face_landmarks.dat"
13
  face_predictor = dlib.shape_predictor(PREDICTOR_PATH)
14
 
 
 
 
 
 
 
 
 
 
 
 
 
15
  # Load deepfake detection model
16
  model = models.resnet34()
17
  model.fc = torch.nn.Linear(model.fc.in_features, 2)
18
- ckpt_path = "./resnet34.pkl"
19
  model.load_state_dict(torch.load(ckpt_path, map_location="cpu"))
20
  model.eval()
21
 
@@ -73,11 +84,17 @@ def process_video(video_path: str):
73
  return output_path
74
 
75
  def gradio_interface(video_file):
 
 
 
 
76
  with NamedTemporaryFile(delete=False, suffix=".mp4") as temp_file:
77
- temp_file.write(video_file.read())
78
- temp_path = temp_file.name
 
 
79
 
80
- output_path = process_video(temp_path)
81
  return output_path
82
 
83
  # Gradio UI
 
6
  from imutils import face_utils
7
  from torchvision import models, transforms
8
  from tempfile import NamedTemporaryFile
9
+ import shutil
10
  # Load face detector and landmark predictor
11
  face_detector = dlib.get_frontal_face_detector()
12
+ PREDICTOR_PATH = "./shape_predictor_81_face_landmarks.dat"
13
  face_predictor = dlib.shape_predictor(PREDICTOR_PATH)
14
 
15
+ import torch
16
+ import torchvision.models as models
17
+
18
+ # Load pretrained ResNet-34 model
19
+ resnet34 = models.resnet34(weights=models.ResNet34_Weights.IMAGENET1K_V1)
20
+ resnet34.fc = torch.nn.Linear(resnet34.fc.in_features, 2)
21
+ ckpt_path = "./resnet34.pkl"
22
+
23
+ # Save model state dict
24
+ torch.save(resnet34.state_dict(), ckpt_path)
25
+ print(f"✅ Model saved at {ckpt_path}")
26
+
27
  # Load deepfake detection model
28
  model = models.resnet34()
29
  model.fc = torch.nn.Linear(model.fc.in_features, 2)
 
30
  model.load_state_dict(torch.load(ckpt_path, map_location="cpu"))
31
  model.eval()
32
 
 
84
  return output_path
85
 
86
  def gradio_interface(video_file):
87
+ if video_file is None:
88
+ return "Error: No video uploaded."
89
+
90
+ # Create a temporary file and copy the uploaded video content
91
  with NamedTemporaryFile(delete=False, suffix=".mp4") as temp_file:
92
+ temp_file_path = temp_file.name
93
+ # Read the uploaded video file using its path
94
+ with open(video_file, "rb") as uploaded_file:
95
+ temp_file.write(uploaded_file.read())
96
 
97
+ output_path = process_video(temp_file_path)
98
  return output_path
99
 
100
  # Gradio UI