Update app.py

app.py

@@ -1,27 +1,38 @@
-from huggingface_hub import from_pretrained_fastai
 import gradio as gr
-from fastai.vision.all import *
-from PIL import ImageFile
-from icevision.all import *
+import PIL
+import torch
+from fastai.vision.all import models
+from icevision.all import ClassMap, tfms
 
 # repo_id = "YOUR_USERNAME/YOUR_LEARNER_NAME"
 repo_id = "el-filatova/Practica2"
 
-learner = from_pretrained_fastai(repo_id)
-labels = learner.dls.vocab
-class_map = ClassMap(['kangaroo'])
-state_dict = torch.load('fasterRCNNFkangaroo.pth')
-model = models.torchvision.faster_rcnn.model(num_classes=len(class_map))
+
+class_map = ClassMap(["kangaroo"])
+state_dict = torch.load("fasterRCNNFkangaroo.pth")
+model = models.torchvision.faster_rcnn.model(
+    backbone=models.torchvision.faster_rcnn.backbones.resnet50_fpn,
+    num_classes=len(class_map),
+)
 model.load_state_dict(state_dict)
 
+size = 384
+
+
 # Definimos una función que se encarga de llevar a cabo las predicciones
-def predict(img):
-    img = PILImage.create(img)
-    infer_tfms = tfms.A.Adapter([*tfms.A.resize_and_pad(size),tfms.A.Normalize()])
-    pred_dict  = models.torchvision.faster_rcnn.end2end_detect(img, infer_tfms, model.to("cpu"), class_map=class_map, detection_threshold=0.5)
-    return pred_dict
-
-    
-# Creamos la interfaz y la lanzamos. 
-gr.Interface(fn=predict, inputs=gr.inputs.Image(shape=(128, 128)), outputs=gr.outputs.Label(num_top_classes=3),examples=['image.jpg']).launch(share=False)
- 
+def predict(img_file):
+    img = PIL.Image.open(img_file)
+    infer_tfms = tfms.A.Adapter([*tfms.A.resize_and_pad(size), tfms.A.Normalize()])
+    pred_dict = models.torchvision.faster_rcnn.end2end_detect(
+        img, infer_tfms, model.to("cpu"), class_map=class_map, detection_threshold=0.5
+    )
+    return pred_dict["img"]
+
+
+# Creamos la interfaz y la lanzamos.
+gr.Interface(
+    fn=predict,
+    inputs=gr.inputs.Image(shape=(128, 128)),
+    outputs=gr.outputs.Label(num_top_classes=3),
+    examples=["image.jpg"],
+).launch(share=False)