Update app.py
Browse files
app.py
CHANGED
@@ -2,10 +2,10 @@ from fastai.vision.all import *
|
|
2 |
from icevision.all import *
|
3 |
import gradio as gr
|
4 |
|
|
|
5 |
|
6 |
# Cargamos el learner
|
7 |
-
model = models.torchvision.faster_rcnn.model(backbone=models.torchvision.faster_rcnn.backbones.resnet18_fpn,
|
8 |
-
num_classes=2)
|
9 |
state_dict = torch.load('fasterRCNNkangaroo.pth')
|
10 |
model.load_state_dict(state_dict)
|
11 |
|
@@ -15,8 +15,8 @@ model.load_state_dict(state_dict)
|
|
15 |
infer_tfms = tfms.A.Adapter([*tfms.A.resize_and_pad(384),tfms.A.Normalize()])
|
16 |
def predict(img):
|
17 |
img = PILImage.create(img)
|
18 |
-
pred_dict = models.torchvision.faster_rcnn.end2end_detect(img, infer_tfms, model.to("cpu"), class_map=
|
19 |
return pred_dict['img']
|
20 |
|
21 |
# Creamos la interfaz y la lanzamos.
|
22 |
-
gr.Interface(fn=predict, inputs=gr.inputs.Image(shape=(128, 128)), outputs=gr.outputs.Image(
|
|
|
2 |
from icevision.all import *
|
3 |
import gradio as gr
|
4 |
|
5 |
+
class_map=ClassMap(['kangaroo'])
|
6 |
|
7 |
# Cargamos el learner
|
8 |
+
model = models.torchvision.faster_rcnn.model(backbone=models.torchvision.faster_rcnn.backbones.resnet18_fpn,num_classes=len(class_map))
|
|
|
9 |
state_dict = torch.load('fasterRCNNkangaroo.pth')
|
10 |
model.load_state_dict(state_dict)
|
11 |
|
|
|
15 |
infer_tfms = tfms.A.Adapter([*tfms.A.resize_and_pad(384),tfms.A.Normalize()])
|
16 |
def predict(img):
|
17 |
img = PILImage.create(img)
|
18 |
+
pred_dict = models.torchvision.faster_rcnn.end2end_detect(img, infer_tfms, model.to("cpu"), class_map=class_map, detection_threshold=0.5)
|
19 |
return pred_dict['img']
|
20 |
|
21 |
# Creamos la interfaz y la lanzamos.
|
22 |
+
gr.Interface(fn=predict, inputs=gr.inputs.Image(shape=(128, 128)), outputs=gr.outputs.Image(),examples=['00004.jpg','00014.jpg']).launch(share=False)
|