igmarco commited on
Commit
1a2f706
·
1 Parent(s): 5a6731c

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -7,7 +7,7 @@ import PIL
7
 
8
  checkpoint_path = "fasterRCNN_resnet18_Raccoons.pth"
9
  model = models.torchvision.faster_rcnn.model(backbone=models.torchvision.faster_rcnn.backbones.resnet18_fpn(pretrained=True),
10
- num_classes=1)
11
 
12
  state_dict = torch.load(checkpoint_path, map_location=torch.device('cpu'))
13
  model.load_state_dict(state_dict)
 
7
 
8
  checkpoint_path = "fasterRCNN_resnet18_Raccoons.pth"
9
  model = models.torchvision.faster_rcnn.model(backbone=models.torchvision.faster_rcnn.backbones.resnet18_fpn(pretrained=True),
10
+ num_classes=2)
11
 
12
  state_dict = torch.load(checkpoint_path, map_location=torch.device('cpu'))
13
  model.load_state_dict(state_dict)