Specify num_classes = 2 for my custom dataset

facebookresearch#108 (comment)
woctezuma · Aug 13, 2020 · 6ecf58f · 6ecf58f
1 parent 40ed80c
commit 6ecf58f
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/models/detr.py b/models/detr.py
@@ -305,6 +305,10 @@ def build(args):
     num_classes = 20 if args.dataset_file != 'coco' else 91
     if args.dataset_file == "coco_panoptic":
         num_classes = 250
+    if args.dataset_file == 'custom':
+        # "You should always use num_classes = max_id + 1 where max_id is the highest class ID that you have in your dataset."
+        # Reference: https://github.com/facebookresearch/detr/issues/108#issuecomment-650269223
+        num_classes = 2
     device = torch.device(args.device)
 
     backbone = build_backbone(args)