Minor fixes to make the system run

examples/ssd/demo.py

@@ -5,7 +5,7 @@
 from PIL import Image, ImageDraw
 from torch.autograd import Variable
 from torchcv.models.fpnssd import FPNSSD512
-from torchcv.models.ssd import SSD512, SSDBoxCoder
+from torchcv.models.ssd import SSDBoxCoder
 
 
 print('Loading model..')

examples/ssd/eval.py

@@ -7,19 +7,21 @@
 from torchcv.transforms import resize
 from torchcv.datasets import ListDataset
 from torchcv.evaluations.voc_eval import voc_eval
-from torchcv.models.ssd import SSD300, SSDBoxCoder
+from torchcv.models.fpnssd import FPNSSD512
+from torchcv.models.ssd import SSDBoxCoder
 
 from PIL import Image
+import numpy as np
 
 
 print('Loading model..')
-net = SSD300(num_classes=21)
+net = FPNSSD512(num_classes=21)
 net.load_state_dict(torch.load('./examples/ssd/checkpoint/params.pth'))
 net.cuda()
 net.eval()
 
 print('Preparing dataset..')
-img_size = 300
+img_size = 512
 def transform(img, boxes, labels):
     img, boxes = resize(img, boxes, size=(img_size,img_size))
     img = transforms.Compose([
@@ -32,7 +34,7 @@ def transform(img, boxes, labels):
                       list_file='torchcv/datasets/voc/voc07_test.txt',
                       transform=transform)
 dataloader = torch.utils.data.DataLoader(dataset, batch_size=1, shuffle=False, num_workers=2)
-box_coder = SSDBoxCoder()
+box_coder = SSDBoxCoder(net)
 
 pred_boxes = []
 pred_labels = []
@@ -44,7 +46,7 @@ def transform(img, boxes, labels):
     gt_difficults = []
     for line in f.readlines():
         line = line.strip().split()
-        d = [int(x) for x in line[1:]]
+        d = np.array([int(x) for x in line[1:]])
         gt_difficults.append(d)
 
 def eval(net, dataset):
@@ -62,10 +64,9 @@ def eval(net, dataset):
         pred_boxes.append(box_preds)
         pred_labels.append(label_preds)
         pred_scores.append(score_preds)
-
-    print voc_eval(
-        pred_boxes, pred_labels, pred_scores,
-        gt_boxes, gt_labels, gt_difficults,
-        iou_thresh=0.5, use_07_metric=True)
+
+    print(voc_eval(pred_boxes, pred_labels, pred_scores,
+                   gt_boxes, gt_labels, gt_difficults,
+                   iou_thresh=0.5, use_07_metric=True))
 
 eval(net, dataset)

torchcv/models/fpnssd/net.py

@@ -2,7 +2,7 @@
 import torch.nn as nn
 import torch.nn.functional as F
 
-from fpn import FPN50
+from .fpn  import FPN50
 from torch.autograd import Variable
 
 
@@ -44,7 +44,7 @@ def forward(self, x):
 
 
 def test():
-    net = SSD512(21)
+    net = FPNSSD512(21)
     loc_preds, cls_preds = net(Variable(torch.randn(1,3,512,512)))
     print(loc_preds.size(), cls_preds.size())
 

torchcv/models/retinanet/retinanet.py

@@ -1,7 +1,7 @@
 import torch
 import torch.nn as nn
 
-from fpn import FPN50
+from .fpn import FPN50
 from torch.autograd import Variable