Update app.py

app.py

@@ -19,6 +19,7 @@ warnings.filterwarnings('ignore')
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 sbcl = create_model("Weights/weights.tar")
+sbcl.to(device)
 
 face_detector = get_model("resnet50_2020-07-20", max_size=1024, device=device)
 face_detector.eval()
@@ -44,13 +45,13 @@ def predict_image(inp):
         return {'No face detected!': 1}, None
 
     with torch.no_grad():
-        img = torch.tensor(face_list).to(device).float()/255
+        img = torch.tensor(face_list).to(device).float() / 255
         pred = sbcl(img).softmax(1)[:, 1].cpu().data.numpy().tolist()[0]
-        confidences = {'Real': 1-pred, 'Fake': pred}
+        confidences = {'Real': 1 - pred, 'Fake': pred}
 
     grayscale_cam = cam_sbcl(input_tensor=img, targets=targets, aug_smooth=True)
     grayscale_cam = grayscale_cam[0, :]
-    cam_image = show_cam_on_image(face_list[0].transpose(1, 2, 0)/255, grayscale_cam, use_rgb=True)
+    cam_image = show_cam_on_image(face_list[0].transpose(1, 2, 0) / 255, grayscale_cam, use_rgb=True)
 
     return confidences, cam_image
 
@@ -59,7 +60,7 @@ def predict_video(inp):
     face_list, idx_list = extract_frames(inp, 10, face_detector)
 
     with torch.no_grad():
-        img = torch.tensor(face_list).to(device).float()/255
+        img = torch.tensor(face_list).to(device).float() / 255
         pred = sbcl(img).softmax(1)[:, 1]
 
     pred_list = []
@@ -77,9 +78,9 @@ def predict_video(inp):
     most_fake = np.argmax(pred_res)
     grayscale_cam = cam_sbcl(input_tensor=img[most_fake].unsqueeze(0), targets=targets, aug_smooth=True)
     grayscale_cam = grayscale_cam[0, :]
-    cam_image = show_cam_on_image(face_list[most_fake].transpose(1, 2, 0)/255, grayscale_cam, use_rgb=True)
+    cam_image = show_cam_on_image(face_list[most_fake].transpose(1, 2, 0) / 255, grayscale_cam, use_rgb=True)
 
-    return {'Real': 1-pred, 'Fake': pred}, cam_image
+    return {'Real': 1 - pred, 'Fake': pred}, cam_image
 
 with gr.Blocks(title="Deepfake Detection CL", theme='upsatwal/mlsc_tiet', css="""
     @import url('https://fonts.googleapis.com/css?family=Source+Code+Pro:200');
@@ -176,7 +177,7 @@ with gr.Blocks(title="Deepfake Detection CL", theme='upsatwal/mlsc_tiet', css=""
             inputs=input_video,
             outputs=output_image_video,
             fn=predict_video,
-            cache_examples=False,
+            cache_examples=True,
         )
         btn_video.click(predict_video, inputs=input_video, outputs=[label_probs_video, output_image_video], api_name="/predict_video")