Update app.py

app.py

@@ -1,7 +1,99 @@
 import gradio as gr
+import torch
+import torch.nn.functional as F
+from facenet_pytorch import MTCNN, InceptionResnetV1
+import os
+import numpy as np
+from PIL import Image
+import matplotlib.pyplot as plt
 
-def saludar(nombre):
-  return f'Hola {nombre}'
-inf=gr.Interface(saludar,inputs='text',outputs='text')
+DEVICE = 'cuda:0' if torch.cuda.is_available() else 'cpu'
+print(f'Running on device: {DEVICE.upper()}')
 
-inf.launch()
+torch.load('./models/resnetinceptionv1_final.pth',map_location='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
+)
+model.load_state_dict(torch.load('./models/resnetinceptionv1_final.pth',map_location='cpu'))
+model.to(DEVICE)
+model.eval()
+print("MTCNN & Classfier models loaded")
+
+
+EXAMPLES_FOLDER = 'examples'
+examples_names = os.listdir(EXAMPLES_FOLDER)
+examples = []
+for example_name in examples_names:
+    example_path = os.path.join(EXAMPLES_FOLDER, example_name)
+    label = example_name.split('_')[0]
+    example = {
+        'path': example_path,
+        'label': label
+    }
+    examples.append(example)
+ 
+   
+       
+def predict(input_image:Image.Image):
+    """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
+    face_image_to_plot = face.squeeze(0).permute(1, 2, 0).cpu().detach().int().numpy()
+
+    face = face.to(DEVICE)
+    face = face.to(torch.float32)
+    face = face / 255.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
+        }
+    return confidences, face_image_to_plot 
+    
+  for i in range(10):
+    example = examples[8]
+    example_img = example['path']
+    example_label = example['label']
+
+    print(f"True label: {example_label}")
+
+    example_img = Image.open(example_img)
+    confidences, _ = predict(example_img)
+    if confidences['real'] > 0.5:
+        print("Predicted label: real")
+    else:
+        print("Predicted label: fake")
+
+    print()     
+    
+    
+    interface = gr.Interface(
+    fn=predict,
+    inputs=gr.inputs.Image(label="Input Image", type="pil"),
+    outputs=[
+        gr.outputs.Label(label="Class"),
+        gr.outputs.Image(label="Face")
+    ],
+    examples=[examples[i]["path"] for i in range(8)] # fake examples
+).launch()