edgilr commited on
Commit
674c93a
verified
1 Parent(s): 8447981

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +14 -10
app.py CHANGED
@@ -1,21 +1,25 @@
 
1
  import gradio as gr
2
  from fastai.vision.all import *
3
  from icevision.all import *
4
  from icevision.models.checkpoint import *
5
  import PIL
6
 
7
- class_map = ClassMap(['kangaroo'])
8
  model = models.torchvision.faster_rcnn.model(backbone=models.torchvision.faster_rcnn.backbones.resnet50_fpn,
9
- num_classes=len(class_map))
10
- state_dict = torch.load('fasterRCNNKangaroo.pth')
 
11
  model.load_state_dict(state_dict)
12
 
13
- infer_tfms = tfms.A.Adapter([*tfms.A.resize_and_pad(size),tfms.A.Normalize()])
14
- size = 384
15
- def predict(img):
16
- img = PILImage.create(img)
17
- pred_dict = models.torchvision.faster_rcnn.end2end_detect(img, infer_tfms, model.to("cpu"), class_map=class_map, detection_threshold=0.5)
18
- return pred_dict['img']
19
 
 
 
 
 
 
 
20
  # Creamos la interfaz y la lanzamos.
21
- gr.Interface(fn=predict, inputs=gr.inputs.Image(shape=(128, 128)), outputs=gr.outputs.Label(type="pil"),examples=['00004.jpg','00083.jpg', '00119.jpg']).launch(share=False)
 
 
1
+ from huggingface_hub import from_pretrained_fastai
2
  import gradio as gr
3
  from fastai.vision.all import *
4
  from icevision.all import *
5
  from icevision.models.checkpoint import *
6
  import PIL
7
 
8
+ checkpoint_path = "fasterRCNNKangaroo.pth"
9
  model = models.torchvision.faster_rcnn.model(backbone=models.torchvision.faster_rcnn.backbones.resnet50_fpn,
10
+ num_classes=2)
11
+
12
+ state_dict = torch.load(checkpoint_path, map_location=torch.device('cpu'))
13
  model.load_state_dict(state_dict)
14
 
15
+ infer_tfms = tfms.A.Adapter([*tfms.A.resize_and_pad(384),tfms.A.Normalize()])
 
 
 
 
 
16
 
17
+ # Definimos una funci贸n que se encarga de llevar a cabo las predicciones
18
+ def predict(img):
19
+ img = PIL.Image.fromarray(img, "RGB")
20
+ pred_dict = models.ross.efficientdet.end2end_detect(img, infer_tfms, model.to("cpu"), class_map=ClassMap(['kangaroo']), detection_threshold=0.5)
21
+ return pred_dict["img"]
22
+
23
  # Creamos la interfaz y la lanzamos.
24
+ gr.Interface(fn=predict, inputs=gr.inputs.Image(shape=(128, 128)), outputs=[gr.outputs.Image(type="pil", label="VFNet Inference")],
25
+ examples=['00004.jpg','00083.jpg', '00119.jpg']).launch(share=False)