switched to pipeline

app.py

@@ -10,29 +10,11 @@ import cv2
 from pytorch_grad_cam import GradCAM
 from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
 from pytorch_grad_cam.utils.image import show_cam_on_image
-
+from transformers import pipeline
 with zipfile.ZipFile("examples.zip","r") as zip_ref:
     zip_ref.extractall(".")
 
-DEVICE = 'cuda:0' if torch.cuda.is_available() else 'cpu'
-
-mtcnn = MTCNN(
-    select_largest=False,
-    post_process=False,
-    device=DEVICE
-).to(DEVICE).eval()
-
-model = InceptionResnetV1(
-    pretrained="vggface2",
-    classify=True,
-    num_classes=1,
-    device=DEVICE
-)
-
-checkpoint = torch.load("resnetinceptionv1_epoch_32.pth", map_location=torch.device('cpu'))
-model.load_state_dict(checkpoint['model_state_dict'])
-model.to(DEVICE)
-model.eval()
+pipe = pipeline(model="not-lain/deepfake",trust_remote_code=True)
 
 EXAMPLES_FOLDER = 'examples'
 examples_names = os.listdir(EXAMPLES_FOLDER)
@@ -48,55 +30,20 @@ for example_name in examples_names:
 np.random.shuffle(examples) # shuffle
 
 def predict(input_image:Image.Image, true_label:str):
-    """Predict the label of the input_image"""
-    face = mtcnn(input_image)
-    if face is None:
-        raise Exception('No face detected')
-    face = face.unsqueeze(0) # add the batch dimension
-    face = F.interpolate(face, size=(256, 256), mode='bilinear', align_corners=False)
-    
-    # convert the face into a numpy array to be able to plot it
-    prev_face = face.squeeze(0).permute(1, 2, 0).cpu().detach().int().numpy()
-    prev_face = prev_face.astype('uint8')
-
-    face = face.to(DEVICE)
-    face = face.to(torch.float32)
-    face = face / 255.0
-    face_image_to_plot = face.squeeze(0).permute(1, 2, 0).cpu().detach().int().numpy()
-
-    target_layers=[model.block8.branch1[-1]]
-    use_cuda = True if torch.cuda.is_available() else False
-    cam = GradCAM(model=model, target_layers=target_layers, use_cuda=use_cuda)
-    targets = [ClassifierOutputTarget(0)]
-
-    grayscale_cam = cam(input_tensor=face, targets=targets, eigen_smooth=True)
-    grayscale_cam = grayscale_cam[0, :]
-    visualization = show_cam_on_image(face_image_to_plot, grayscale_cam, use_rgb=True)
-    face_with_mask = cv2.addWeighted(prev_face, 1, visualization, 0.5, 0)
-
-    with torch.no_grad():
-        output = torch.sigmoid(model(face).squeeze(0))
-        prediction = "real" if output.item() < 0.5 else "fake"
-        
-        real_prediction = 1 - output.item()
-        fake_prediction = output.item()
-        
-        confidences = {
-            'real': real_prediction,
-            'fake': fake_prediction
-        }
+    out = pipe.predict(input_image)
+    confidences,face_with_mask = out["confidences"], out["face_with_mask"]
     return confidences, true_label, face_with_mask
 
 interface = gr.Interface(
     fn=predict,
     inputs=[
-        gr.inputs.Image(label="Input Image", type="pil"),
+        gr.Image(label="Input Image", type="filepath"),
         "text"
     ],
     outputs=[
-        gr.outputs.Label(label="Class"),
+        gr.Label(label="Class"),
         "text",
-        gr.outputs.Image(label="Face with Explainability")
+        gr.Image(label="Face with Explainability")
     ],
     examples=[[examples[i]["path"], examples[i]["label"]] for i in range(10)]
 ).launch()

resnetinceptionv1_epoch_32.pth

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:794ebe83c6a7d7959c30c175030b4885e2b9fa175f1cc3e582236595d119f52b
-size 282395989