Update detection demo

demo/object_detection/predict_faster_rcnn.py

@@ -15,27 +15,26 @@ parser.add_argument("--use_gpu",            type=ast.literal_eval,  default=True
 parser.add_argument("--checkpoint_dir",     type=str,               default="faster_rcnn_finetune_ckpt",         help="Path to save log data.")
 parser.add_argument("--batch_size",         type=int,               default=8,                                   help="Total examples' number in batch for training.")
 parser.add_argument("--module",             type=str,               default="faster_rcnn_resnet50_coco2017",     help="Module used as feature extractor.")
-parser.add_argument("--dataset",            type=str,               default="coco_10",                           help="Dataset to finetune.")
 # yapf: enable.
 
 
 def predict(args):
     module = hub.Module(name=args.module)
-    dataset = hub.dataset.Coco10('rcnn')
+    dataset = hub.dataset.Balloon('rcnn')
 
     print("dataset.num_labels:", dataset.num_labels)
 
     # define batch reader
     data_reader = ObjectDetectionReader(dataset=dataset, model_type='rcnn')
     pred_input_dict, pred_output_dict, pred_program = module.context(
-        trainable=False, phase='predict')
+        trainable=False, phase='predict', num_classes=dataset.num_labels)
 
     pred_feed_list = [
         pred_input_dict['image'].name, pred_input_dict['im_info'].name,
         pred_input_dict['im_shape'].name
     ]
 
-    pred_feature = [pred_output_dict['head_feat'], pred_output_dict['rois']]
+    pred_feature = [pred_output_dict['head_features'], pred_output_dict['rois']]
 
     config = hub.RunConfig(
         use_data_parallel=False,
@@ -54,8 +53,8 @@ def predict(args):
         config=config)
 
     data = [
-        "./test/test_img_bird.jpg",
-        "./test/test_img_cat.jpg",
+        "./test/balloon1.jpg",
+        "./test/balloon2.jpg",
     ]
     label_map = dataset.label_dict()
     results = task.predict(data=data, return_result=True, accelerate_mode=False)

demo/object_detection/predict_ssd.py

@@ -15,13 +15,12 @@ parser.add_argument("--use_gpu",            type=ast.literal_eval,  default=True
 parser.add_argument("--checkpoint_dir",     type=str,               default="ssd_finetune_ckpt",          help="Path to save log data.")
 parser.add_argument("--batch_size",         type=int,               default=8,                            help="Total examples' number in batch for training.")
 parser.add_argument("--module",             type=str,               default="ssd_vgg16_512_coco2017",     help="Module used as feature extractor.")
-parser.add_argument("--dataset",            type=str,               default="coco_10",                    help="Dataset to finetune.")
 # yapf: enable.
 
 
 def predict(args):
     module = hub.Module(name=args.module)
-    dataset = hub.dataset.Coco10('ssd')
+    dataset = hub.dataset.Balloon('ssd')
 
     print("dataset.num_labels:", dataset.num_labels)
 
@@ -50,8 +49,8 @@ def predict(args):
         config=config)
 
     data = [
-        "./test/test_img_bird.jpg",
-        "./test/test_img_cat.jpg",
+        "./test/balloon1.jpg",
+        "./test/balloon2.jpg",
     ]
     label_map = dataset.label_dict()
     results = task.predict(data=data, return_result=True, accelerate_mode=False)

demo/object_detection/predict_yolo.py

@@ -15,13 +15,12 @@ parser.add_argument("--use_gpu",            type=ast.literal_eval,  default=True
 parser.add_argument("--checkpoint_dir",     type=str,               default="yolo_finetune_ckpt",         help="Path to save log data.")
 parser.add_argument("--batch_size",         type=int,               default=8,                            help="Total examples' number in batch for training.")
 parser.add_argument("--module",             type=str,               default="yolov3_darknet53_coco2017",  help="Module used as feature extractor.")
-parser.add_argument("--dataset",            type=str,               default="coco_10",                    help="Dataset to finetune.")
 # yapf: enable.
 
 
 def predict(args):
     module = hub.Module(name=args.module)
-    dataset = hub.dataset.Coco10('yolo')
+    dataset = hub.dataset.Balloon('yolo')
 
     print("dataset.num_labels:", dataset.num_labels)
 
@@ -49,8 +48,8 @@ def predict(args):
         config=config)
 
     data = [
-        "./test/test_img_bird.jpg",
-        "./test/test_img_cat.jpg",
+        "./test/balloon1.jpg",
+        "./test/balloon2.jpg",
     ]
     label_map = dataset.label_dict()
     results = task.predict(data=data, return_result=True, accelerate_mode=False)

demo/object_detection/train_faster_rcnn.py

@@ -16,23 +16,23 @@ parser.add_argument("--use_gpu",            type=ast.literal_eval,  default=True
 parser.add_argument("--checkpoint_dir",     type=str,               default="faster_rcnn_finetune_ckpt",      help="Path to save log data.")
 parser.add_argument("--batch_size",         type=int,               default=1,                                help="Total examples' number in batch for training.")
 parser.add_argument("--module",             type=str,               default="faster_rcnn_resnet50_coco2017",  help="Module used as feature extractor.")
-parser.add_argument("--dataset",            type=str,               default="coco_10",                        help="Dataset to finetune.")
 parser.add_argument("--use_data_parallel",  type=ast.literal_eval,  default=False,                            help="Whether use data parallel.")
 # yapf: enable.
 
 
 def finetune(args):
     module = hub.Module(name=args.module)
-    dataset = hub.dataset.Coco10('rcnn')
+    dataset = hub.dataset.Balloon('rcnn')
 
     print("dataset.num_labels:", dataset.num_labels)
 
     # define batch reader
     data_reader = ObjectDetectionReader(dataset=dataset, model_type='rcnn')
 
-    input_dict, output_dict, program = module.context(trainable=True)
+    input_dict, output_dict, program = module.context(
+        trainable=True, num_classes=dataset.num_labels)
     pred_input_dict, pred_output_dict, pred_program = module.context(
-        trainable=False, phase='predict')
+        trainable=False, phase='predict', num_classes=dataset.num_labels)
 
     feed_list = [
         input_dict["image"].name, input_dict["im_info"].name,
@@ -54,7 +54,7 @@ def finetune(args):
 
     config = hub.RunConfig(
         log_interval=10,
-        eval_interval=100,
+        eval_interval=10,
         use_data_parallel=args.use_data_parallel,
         use_pyreader=True,
         use_cuda=args.use_gpu,

demo/object_detection/train_ssd.py

@@ -16,14 +16,13 @@ parser.add_argument("--use_gpu",            type=ast.literal_eval,  default=True
 parser.add_argument("--checkpoint_dir",     type=str,               default="ssd_finetune_ckpt",          help="Path to save log data.")
 parser.add_argument("--batch_size",         type=int,               default=8,                            help="Total examples' number in batch for training.")
 parser.add_argument("--module",             type=str,               default="ssd_vgg16_512_coco2017",     help="Module used as feature extractor.")
-parser.add_argument("--dataset",            type=str,               default="coco_10",                    help="Dataset to finetune.")
 parser.add_argument("--use_data_parallel",  type=ast.literal_eval,  default=False,                        help="Whether use data parallel.")
 # yapf: enable.
 
 
 def finetune(args):
     module = hub.Module(name=args.module)
-    dataset = hub.dataset.Coco10('ssd')
+    dataset = hub.dataset.Balloon('ssd')
 
     print("dataset.num_labels:", dataset.num_labels)
 

demo/object_detection/train_yolo.py

@@ -16,14 +16,13 @@ parser.add_argument("--use_gpu",            type=ast.literal_eval,  default=True
 parser.add_argument("--checkpoint_dir",     type=str,               default="yolo_finetune_ckpt",         help="Path to save log data.")
 parser.add_argument("--batch_size",         type=int,               default=8,                            help="Total examples' number in batch for training.")
 parser.add_argument("--module",             type=str,               default="yolov3_darknet53_coco2017",  help="Module used as feature extractor.")
-parser.add_argument("--dataset",            type=str,               default="coco_10",                    help="Dataset to finetune.")
 parser.add_argument("--use_data_parallel",  type=ast.literal_eval,  default=False,                        help="Whether use data parallel.")
 # yapf: enable.
 
 
 def finetune(args):
     module = hub.Module(name=args.module)
-    dataset = hub.dataset.Coco10('yolo')
+    dataset = hub.dataset.Balloon('yolo')
 
     print("dataset.num_labels:", dataset.num_labels)