command line configuration

Author: yhenon

README.md

@@ -5,6 +5,7 @@
 
 Pytorch implementation of RetinaNet object detection as described in [Focal Loss for Dense Object Detection](https://arxiv.org/abs/1708.02002) by Tsung-Yi Lin, Priya Goyal, Ross Girshick, Kaiming He and Piotr Dollár.
 
+This implementation is primarily designed to be easy to read and simple to modify.
 
 ## Results
 Currently, this repo achieves 33.7% mAP at 600px resolution with a Resnet-50 backbone. The published result is 34.0% mAP. The difference is likely due to the use of Adam optimizer instead of SGD with weight decay.
@@ -51,18 +52,33 @@ cd ../
 The network can be trained using the `train.py` script. Currently, two dataloaders are available: COCO and CSV. For training on coco, use
 
 ```
-python train.py coco <path/to/coco>
+python train.py --dataset coco --coco_path ../coco --depth 50
 ```
 
 For training using a custom dataset, with annotations in CSV format (see below), use
 
 ```
-python train.py csv <path/to/annotations.csv> <path/to/classes.csv>
+python train.py --dataset csv --csv_train <path/to/train_annots.csv>  --csv_classes <path/to/train/class_list.csv> --csv_val <path/to/val_annots.csv>
 ```
 
+Note that the --csv_val argument is optional, in which case no validation will be performed.
+
 ## Visualization
 
-To visualize the network detection, use `test.py`.
+To visualize the network detection, use `visualize.py`:
+
+```
+python visualize.py --dataset coco --coco_path ../coco --model <path/to/model.pt>
+```
+This will visualize bounding boxes on the validation set. To visualise with a CSV dataset, use:
+
+```
+python visualize.py --dataset csv --csv_classes <path/to/train/class_list.csv> --csv_val <path/to/val_annots.csv> --model <path/to/model.pt>
+```
+
+## Model
+
+The retinanet model uses a resnet backbone. You can set the depth of the resnet model using the --depth argument. Depth must be one of 18, 34, 50, 101 or 152. Note that deeper models are more accurate but are slower and use more memory.
 
 ## CSV datasets
 The `CSVGenerator` provides an easy way to define your own datasets.

dataloader.py

@@ -93,6 +93,7 @@ def load_annotations(self, image_index):
  # parse annotations
  coco_annotations = self.coco.loadAnns(annotations_ids)
  for idx, a in enumerate(coco_annotations):
+
  # some annotations have basically no width / height, skip them
  if a['bbox'][2] < 1 or a['bbox'][3] < 1:
  continue

model.py

@@ -8,14 +8,11 @@
 
 from lib.nms.pth_nms import pth_nms
 
-import pdb
-
 def nms(dets, thresh):
  "Dispatch to either CPU or GPU NMS implementations.\
  Accept dets as tensor"""
  return pth_nms(dets, thresh)
 
-
 model_urls = {
  'resnet18': 'https://download.pytorch.org/models/resnet18-5c106cde.pth',
  'resnet34': 'https://download.pytorch.org/models/resnet34-333f7ec4.pth',
@@ -310,3 +307,24 @@ def resnet50(num_classes, pretrained=False, **kwargs):
  if pretrained:
  model.load_state_dict(model_zoo.load_url(model_urls['resnet50'], model_dir='.'), strict=False)
  return model
+
+def resnet101(num_classes, pretrained=False, **kwargs):
+ """Constructs a ResNet-101 model.
+ Args:
+ pretrained (bool): If True, returns a model pre-trained on ImageNet
+ """
+ model = ResNet(num_classes, Bottleneck, [3, 4, 23, 3], **kwargs)
+ if pretrained:
+ model.load_state_dict(model_zoo.load_url(model_urls['resnet101'], model_dir='.'), strict=False)
+ return model
+
+
+def resnet152(num_classes, pretrained=False, **kwargs):
+ """Constructs a ResNet-152 model.
+ Args:
+ pretrained (bool): If True, returns a model pre-trained on ImageNet
+ """
+ model = ResNet(num_classes, Bottleneck, [3, 8, 36, 3], **kwargs)
+ if pretrained:
+ model.load_state_dict(model_zoo.load_url(model_urls['resnet152'], model_dir='.'), strict=False)
+ return model