Merging train.py

eval.py

@@ -14,7 +14,7 @@
 from data import VOC_CLASSES as labelmap
 import torch.utils.data as data
 
-from data import AnnotationTransform, VOCDetection, BaseTransform
+from data import AnnotationTransform, VOCDetection, BaseTransform, VOC_CLASSES
 from ssd import build_ssd
 
 import sys
@@ -408,7 +408,8 @@ def evaluate_detections(box_list, output_dir, dataset):
 
 if __name__ == '__main__':
     # load net
-    net = build_ssd('test', 300, 21)    # initialize SSD
+    num_classes = len(VOC_CLASSES) + 1 # +1 background
+    net = build_ssd('test', 300, num_classes) # initialize SSD
     net.load_state_dict(torch.load(args.trained_model))
     net.eval()
     print('Finished loading model!')

test.py

@@ -9,7 +9,7 @@
 from torch.autograd import Variable
 from data import VOCroot, VOC_CLASSES as labelmap
 from PIL import Image
-from data import AnnotationTransform, VOCDetection, BaseTransform
+from data import AnnotationTransform, VOCDetection, BaseTransform, VOC_CLASSES
 import torch.utils.data as data
 from ssd import build_ssd
 
@@ -73,7 +73,8 @@ def test_net(save_folder, net, cuda, testset, transform, thresh):
 
 if __name__ == '__main__':
     # load net
-    net = build_ssd('test', 300, 21)    # initialize SSD
+    num_classes = len(VOC_CLASSES) + 1 # +1 background
+    net = build_ssd('test', 300, num_classes) # initialize SSD
     net.load_state_dict(torch.load(args.trained_model))
     net.eval()
     print('Finished loading model!')

train.py

@@ -8,7 +8,7 @@
 import argparse
 from torch.autograd import Variable
 import torch.utils.data as data
-from data import v2, v1, AnnotationTransform, VOCDetection, detection_collate, VOCroot
+from data import v2, v1, AnnotationTransform, VOCDetection, detection_collate, VOCroot, VOC_CLASSES
 from utils.augmentations import SSDAugmentation
 from layers.modules import MultiBoxLoss
 from ssd import build_ssd
@@ -54,7 +54,7 @@ def str2bool(v):
 # train_sets = 'train'
 ssd_dim = 300  # only support 300 now
 means = (104, 117, 123)  # only support voc now
-num_classes = 21
+num_classes = len(VOC_CLASSES) + 1
 batch_size = args.batch_size
 accum_batch_size = 32
 iter_size = accum_batch_size / batch_size
@@ -68,7 +68,7 @@ def str2bool(v):
     import visdom
     viz = visdom.Visdom()
 
-ssd_net = build_ssd('train', 300, 21)
+ssd_net = build_ssd('train', 300, num_classes)
 net = ssd_net
 
 if args.cuda: