From a8c871855be3ae535f05905dbb95c40396e64c79 Mon Sep 17 00:00:00 2001
From: wuzewu <wuzewu@baidu.com>
Date: Wed, 29 Apr 2020 14:40:02 +0800
Subject: [PATCH] update detection demo

---
 demo/object_detection/predict_yolo.py      | 11 +++++------
 demo/object_detection/train_faster_rcnn.py |  4 ++--
 demo/object_detection/train_yolo.py        |  2 +-
 3 files changed, 8 insertions(+), 9 deletions(-)

diff --git a/demo/object_detection/predict_yolo.py b/demo/object_detection/predict_yolo.py
index 4aa8732c..e001c9ef 100644
--- a/demo/object_detection/predict_yolo.py
+++ b/demo/object_detection/predict_yolo.py
@@ -12,21 +12,21 @@ from paddlehub.dataset.base_cv_dataset import ObjectDetectionDataset
 # yapf: disable
 parser = argparse.ArgumentParser(__doc__)
 parser.add_argument("--use_gpu",            type=ast.literal_eval,  default=True,                         help="Whether use GPU for fine-tuning.")
-parser.add_argument("--checkpoint_dir",     type=str,               default="ssd_finetune_ckpt",          help="Path to save log data.")
+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="ssd_vgg16_512_coco2017",     help="Module used as feature extractor.")
+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('ssd')
+    dataset = hub.dataset.Coco10('yolo')
 
     print("dataset.num_labels:", dataset.num_labels)
 
     # define batch reader
-    data_reader = ObjectDetectionReader(dataset=dataset, model_type='ssd')
+    data_reader = ObjectDetectionReader(dataset=dataset, model_type='yolo')
 
     input_dict, output_dict, program = module.context(trainable=True)
     feed_list = [input_dict["image"].name, input_dict["im_size"].name]
@@ -41,12 +41,11 @@ def predict(args):
         checkpoint_dir=args.checkpoint_dir,
         strategy=hub.finetune.strategy.DefaultFinetuneStrategy())
 
-    task = hub.SSDTask(
+    task = hub.YOLOTask(
         data_reader=data_reader,
         num_classes=dataset.num_labels,
         feed_list=feed_list,
         feature=feature,
-        multi_box_head_config=module.multi_box_head_config,
         config=config)
 
     data = [
diff --git a/demo/object_detection/train_faster_rcnn.py b/demo/object_detection/train_faster_rcnn.py
index e64600ec..b379b16a 100644
--- a/demo/object_detection/train_faster_rcnn.py
+++ b/demo/object_detection/train_faster_rcnn.py
@@ -46,11 +46,11 @@ def finetune(args):
     ]
 
     feature = [
-        output_dict['head_feat'], output_dict['rpn_cls_loss'],
+        output_dict['head_features'], output_dict['rpn_cls_loss'],
         output_dict['rpn_reg_loss'], output_dict['generate_proposal_labels']
     ]
 
-    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(
         log_interval=10,
diff --git a/demo/object_detection/train_yolo.py b/demo/object_detection/train_yolo.py
index f9e0711b..df835b4e 100644
--- a/demo/object_detection/train_yolo.py
+++ b/demo/object_detection/train_yolo.py
@@ -32,7 +32,7 @@ def finetune(args):
 
     input_dict, output_dict, program = module.context(trainable=True)
     feed_list = [input_dict["image"].name, input_dict["im_size"].name]
-    feature = output_dict['head_features']
+    feature = output_dict['body_features']
 
     config = hub.RunConfig(
         log_interval=10,
-- 
GitLab