Alesteba commited on
Commit
0d067fc
·
1 Parent(s): bf0255c

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +4 -3
app.py CHANGED
@@ -9,9 +9,12 @@ import PIL
9
 
10
  class_map = ClassMap(['raccoon','banana'])
11
 
 
 
12
  model_2 = models.torchvision.retinanet.model(
13
  backbone=models.torchvision.retinanet.backbones.resnext50_32x4d_fpn (pretrained=True),
14
- num_classes=len(class_map)
 
15
  )
16
 
17
  # load from model_repo:
@@ -25,8 +28,6 @@ model_2.load_state_dict(state_dict)
25
 
26
  # use test img:
27
 
28
- size = 384
29
-
30
  infer_tfms = tfms.A.Adapter([*tfms.A.resize_and_pad(size),tfms.A.Normalize()])
31
 
32
  def predict(img):
 
9
 
10
  class_map = ClassMap(['raccoon','banana'])
11
 
12
+ size = 384
13
+
14
  model_2 = models.torchvision.retinanet.model(
15
  backbone=models.torchvision.retinanet.backbones.resnext50_32x4d_fpn (pretrained=True),
16
+ num_classes=len(class_map),
17
+ img_size = size
18
  )
19
 
20
  # load from model_repo:
 
28
 
29
  # use test img:
30
 
 
 
31
  infer_tfms = tfms.A.Adapter([*tfms.A.resize_and_pad(size),tfms.A.Normalize()])
32
 
33
  def predict(img):