aa1223 commited on
Commit
4cce507
1 Parent(s): 0fd4fc6

Upload 3 files

Browse files
Files changed (3) hide show
  1. .gitignore +1 -0
  2. app.py +83 -0
  3. requirements.txt +0 -0
.gitignore ADDED
@@ -0,0 +1 @@
 
 
1
+ *.pth
app.py ADDED
@@ -0,0 +1,83 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # from google.colab import drive
2
+ # drive.mount('/content/drive')
3
+
4
+ # import gradio as gr
5
+ import gradio as gr
6
+ import webbrowser
7
+ from threading import Timer
8
+ import torch
9
+ import torch.nn.functional as F
10
+ from facenet_pytorch import InceptionResnetV1
11
+ import cv2
12
+ from PIL import Image
13
+ import numpy as np
14
+ import warnings
15
+ warnings.filterwarnings("ignore")
16
+
17
+ DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
18
+
19
+ model = InceptionResnetV1(pretrained="vggface2", classify=True, num_classes=1).to(DEVICE).eval()
20
+
21
+ # checkpoint_path = "/content/drive/MyDrive/resnetinceptionv1_epoch_32.pth"
22
+ checkpoint_path = "resnetinceptionv1_epoch_32.pth"
23
+ checkpoint = torch.load(checkpoint_path, map_location=DEVICE)
24
+ if 'model_state_dict' in checkpoint:
25
+ state_dict = checkpoint['model_state_dict']
26
+ else:
27
+ state_dict = checkpoint
28
+
29
+ try:
30
+ model.load_state_dict(state_dict)
31
+ print("Model weights loaded successfully.")
32
+ except RuntimeError as e:
33
+ print(f"Error loading model weights: {e}")
34
+
35
+ def create_montage(frames, size=(512, 512)):
36
+ """Create a montage from selected frames."""
37
+ montage = Image.new('RGB', size)
38
+ num_images = len(frames)
39
+ montage_grid = int(np.ceil(np.sqrt(num_images)))
40
+ thumb_size = (size[0] // montage_grid, size[1] // montage_grid)
41
+
42
+ for i, frame in enumerate(frames):
43
+ thumbnail = frame.resize(thumb_size, Image.ANTIALIAS)
44
+ x_offset = (i % montage_grid) * thumb_size[0]
45
+ y_offset = (i // montage_grid) * thumb_size[1]
46
+ montage.paste(thumbnail, (x_offset, y_offset))
47
+
48
+ return montage
49
+
50
+ def predict(input_video):
51
+ cap = cv2.VideoCapture(input_video)
52
+ total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
53
+
54
+ selected_frames = []
55
+ frame_indices = np.linspace(0, total_frames - 1, 9, dtype=int)
56
+
57
+ for i in range(total_frames):
58
+ ret, frame = cap.read()
59
+ if not ret:
60
+ break
61
+ if i in frame_indices:
62
+ frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
63
+ selected_frames.append(Image.fromarray(frame))
64
+
65
+ cap.release()
66
+
67
+ video_label = "Fake" if np.random.rand() > 0.5 else "Real"
68
+ fake_ratio = np.random.rand()
69
+ detail = f"Placeholder ratio: {fake_ratio*100:.2f}%"
70
+
71
+ montage = create_montage(selected_frames)
72
+
73
+ return video_label, detail, montage
74
+
75
+ interface = gr.Interface(
76
+ fn=predict,
77
+ inputs=gr.Video(label="Input Video"),
78
+ outputs=[
79
+ gr.Text(label="Classification"),
80
+ gr.Text(label="Details"),
81
+ gr.Image(label="Montage of Selected Frames")
82
+ ],
83
+ ).launch(debug=True, share=True)
requirements.txt ADDED
Binary file (2.67 kB). View file