inigo99 commited on
Commit
6deab91
·
1 Parent(s): 63e9335

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -3
app.py CHANGED
@@ -1,12 +1,13 @@
1
  from huggingface_hub import from_pretrained_fastai
2
  import gradio as gr
3
- from fastai.vision.all import *
4
 
5
  # repo_id = "YOUR_USERNAME/YOUR_LEARNER_NAME"
6
  repo_id = "inigo99/kangaroo-detector"
7
 
8
- learner = from_pretrained_fastai(repo_id)
9
- labels = learner.dls.vocab
 
 
10
 
11
  # Definimos una función que se encarga de llevar a cabo las predicciones
12
  def predict(img):
 
1
  from huggingface_hub import from_pretrained_fastai
2
  import gradio as gr
 
3
 
4
  # repo_id = "YOUR_USERNAME/YOUR_LEARNER_NAME"
5
  repo_id = "inigo99/kangaroo-detector"
6
 
7
+ model = models.torchvision.faster_rcnn.model(backbone=models.torchvision.faster_rcnn.backbones.resnet18_fpn,
8
+ num_classes=len(class_map))
9
+ state_dict = torch.load(repo_id)
10
+ model.load_state_dict(state_dict)
11
 
12
  # Definimos una función que se encarga de llevar a cabo las predicciones
13
  def predict(img):