Update app.py

app.py

@@ -1,37 +1,32 @@
-import gradio as gr
-from transformers import AutoModelForImageClassification, AutoImageProcessor
-from PIL import Image
-import torch
+import gradio as gr 
+from transformers import AutoModelForImageClassification, AutoImageProcessor 
+from PIL import Image 
+import torch 
 import os 
-
-MODEL_ID = "nicolasrl/deepfake_vs_real_ViTlarge"
+ 
+MODEL_ID = "nicolasrl/deepfake_vs_real_ViTlarge" 
 HF_TOKEN = os.getenv("HF_TOKEN")  
-
-model = AutoModelForImageClassification.from_pretrained(MODEL_ID, use_auth_token=HF_TOKEN)
+ 
+model = AutoModelForImageClassification.from_pretrained(MODEL_ID, use_auth_token=HF_TOKEN) 
 processor = AutoImageProcessor.from_pretrained(MODEL_ID, use_auth_token=HF_TOKEN) 
-
-def predict(image_path: str):
-    image = Image.open(image_path)
-    
-    inputs = processor(images=image, return_tensors="pt")
-    
-    with torch.no_grad():
-        outputs = model(**inputs)
-        
-    logits = outputs.logits
-    probs = torch.nn.functional.softmax(logits, dim=-1)[0]
-    results = {model.config.id2label[i]: float(probs[i]) for i in range(len(probs))}
-    
+ 
+def predict(image: Image.Image):     
+    inputs = processor(images=image, return_tensors="pt")     
+    with torch.no_grad():         
+        outputs = model(**inputs)     
+    logits = outputs.logits     
+    probs = torch.nn.functional.softmax(logits, dim=-1)[0]     
+    results = {model.config.id2label[i]: float(probs[i]) for i in range(len(probs))}     
     return dict(sorted(results.items(), key=lambda x: x[1], reverse=True)) 
-
-iface = gr.Interface(
-    fn=predict,
-    inputs=gr.Image(type="filepath"), #PIL
-    outputs=gr.Label(num_top_classes=2),
-    title="Deepfake o Real",
-    description="Sube una imagen y el modelo hace magia para predecir si es Deepfake o Real",
-    api_name="predict"
+ 
+iface = gr.Interface(     
+    fn=predict,     
+    inputs=gr.Image(type="pil"),     
+    outputs=gr.Label(num_top_classes=2),     
+    title="Deepfake o Real",     
+    description="Sube una imagen y el modelo hace magia para predecir si es Deepfake o Real",     
+    api_name="predict" 
 ) 
-
-if __name__ == "__main__":
+ 
+if __name__ == "__main__":     
     iface.launch(show_error=True)