update YOLO act dataloader (PaddlePaddle#1336)

example/auto_compression/pytorch_yolov5/configs/yolov5s_qat_dis.yaml

@@ -1,10 +1,11 @@
-
 Global:
-  reader_config: configs/yolov5_reader.yml
-  input_list: {'image': 'x2paddle_images'}
-  Evaluation: True
-  arch: 'YOLOv5'
   model_dir: ./yolov5s.onnx
+  dataset_dir: dataset/coco/
+  train_image_dir: train2017
+  val_image_dir: val2017
+  train_anno_path: annotations/instances_train2017.json
+  val_anno_path: annotations/instances_val2017.json
+  Evaluation: True
 
 Distillation:
   alpha: 1.0

example/auto_compression/pytorch_yolov5/dataset.py

@@ -0,0 +1,113 @@
+from pycocotools.coco import COCO
+import cv2
+import os
+import numpy as np
+import paddle
+
+
+class COCOValDataset(paddle.io.Dataset):
+    def __init__(self,
+                 dataset_dir=None,
+                 image_dir=None,
+                 anno_path=None,
+                 img_size=[640, 640]):
+        self.dataset_dir = dataset_dir
+        self.image_dir = image_dir
+        self.img_size = img_size
+        self.ann_file = os.path.join(dataset_dir, anno_path)
+        self.coco = COCO(self.ann_file)
+        ori_ids = list(sorted(self.coco.imgs.keys()))
+        # check gt bbox
+        clean_ids = []
+        for idx in ori_ids:
+            ins_anno_ids = self.coco.getAnnIds(imgIds=[idx], iscrowd=False)
+            instances = self.coco.loadAnns(ins_anno_ids)
+            num_bbox = 0
+            for inst in instances:
+                if inst.get('ignore', False):
+                    continue
+                if 'bbox' not in inst.keys():
+                    continue
+                elif not any(np.array(inst['bbox'])):
+                    continue
+                else:
+                    num_bbox += 1
+            if num_bbox > 0:
+                clean_ids.append(idx)
+        self.ids = clean_ids
+
+    def __getitem__(self, idx):
+        img_id = self.ids[idx]
+        img = self._get_img_data_from_img_id(img_id)
+        img, scale_factor = self.image_preprocess(img, self.img_size)
+        return {
+            'image': img,
+            'im_id': np.array([img_id]),
+            'scale_factor': scale_factor
+        }
+
+    def __len__(self):
+        return len(self.ids)
+
+    def _get_img_data_from_img_id(self, img_id):
+        img_info = self.coco.loadImgs(img_id)[0]
+        img_path = os.path.join(self.dataset_dir, self.image_dir,
+                                img_info['file_name'])
+        img = cv2.imread(img_path)
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        return img
+
+    def _generate_scale(self, im, target_shape, keep_ratio=True):
+        """
+            Args:
+                im (np.ndarray): image (np.ndarray)
+            Returns:
+                im_scale_x: the resize ratio of X
+                im_scale_y: the resize ratio of Y
+            """
+        origin_shape = im.shape[:2]
+        if keep_ratio:
+            im_size_min = np.min(origin_shape)
+            im_size_max = np.max(origin_shape)
+            target_size_min = np.min(target_shape)
+            target_size_max = np.max(target_shape)
+            im_scale = float(target_size_min) / float(im_size_min)
+            if np.round(im_scale * im_size_max) > target_size_max:
+                im_scale = float(target_size_max) / float(im_size_max)
+            im_scale_x = im_scale
+            im_scale_y = im_scale
+        else:
+            resize_h, resize_w = target_shape
+            im_scale_y = resize_h / float(origin_shape[0])
+            im_scale_x = resize_w / float(origin_shape[1])
+        return im_scale_y, im_scale_x
+
+    def image_preprocess(self, img, target_shape):
+        # Resize image
+        im_scale_y, im_scale_x = self._generate_scale(img, target_shape)
+        img = cv2.resize(
+            img,
+            None,
+            None,
+            fx=im_scale_x,
+            fy=im_scale_y,
+            interpolation=cv2.INTER_LINEAR)
+        # Pad
+        im_h, im_w = img.shape[:2]
+        h, w = target_shape[:]
+        if h != im_h or w != im_w:
+            canvas = np.ones((h, w, 3), dtype=np.float32)
+            canvas *= np.array([114.0, 114.0, 114.0], dtype=np.float32)
+            canvas[0:im_h, 0:im_w, :] = img.astype(np.float32)
+            img = canvas
+        img = np.transpose(img / 255, [2, 0, 1])
+        scale_factor = np.array([im_scale_y, im_scale_x])
+        return img.astype(np.float32), scale_factor
+
+
+class COCOTrainDataset(COCOValDataset):
+    def __getitem__(self, idx):
+        img_id = self.ids[idx]
+        img = self._get_img_data_from_img_id(img_id)
+        img, scale_factor = self.image_preprocess(img, self.img_size)
+        return {'x2paddle_images': img}

example/auto_compression/pytorch_yolov5/eval.py

@@ -16,13 +16,12 @@
 import sys
 import numpy as np
 import argparse
+from tqdm import tqdm
 import paddle
-from ppdet.core.workspace import load_config, merge_config
-from ppdet.core.workspace import create
-from ppdet.metrics import COCOMetric, VOCMetric
 from paddleslim.auto_compression.config_helpers import load_config as load_slim_config
 from paddleslim.common import load_onnx_model
-from post_process import YOLOv5PostProcess
+from post_process import YOLOv5PostProcess, coco_metric
+from dataset import COCOValDataset
 
 
 def argsparser():
@@ -42,104 +41,50 @@ def argsparser():
     return parser
 
 
-def reader_wrapper(reader, input_list):
-    def gen():
-        for data in reader:
-            in_dict = {}
-            if isinstance(input_list, list):
-                for input_name in input_list:
-                    in_dict[input_name] = data[input_name]
-            elif isinstance(input_list, dict):
-                for input_name in input_list.keys():
-                    in_dict[input_list[input_name]] = data[input_name]
-            yield in_dict
-
-    return gen
-
-
-def convert_numpy_data(data, metric):
-    data_all = {}
-    data_all = {k: np.array(v) for k, v in data.items()}
-    if isinstance(metric, VOCMetric):
-        for k, v in data_all.items():
-            if not isinstance(v[0], np.ndarray):
-                tmp_list = []
-                for t in v:
-                    tmp_list.append(np.array(t))
-                data_all[k] = np.array(tmp_list)
-    else:
-        data_all = {k: np.array(v) for k, v in data.items()}
-    return data_all
-
-
 def eval():
 
     place = paddle.CUDAPlace(0) if FLAGS.devices == 'gpu' else paddle.CPUPlace()
     exe = paddle.static.Executor(place)
+
     val_program, feed_target_names, fetch_targets = load_onnx_model(
         global_config["model_dir"])
 
-    metric = global_config['metric']
-    for batch_id, data in enumerate(val_loader):
-        data_all = convert_numpy_data(data, metric)
-        data_input = {}
-        for k, v in data.items():
-            if isinstance(global_config['input_list'], list):
-                if k in global_config['input_list']:
-                    data_input[k] = np.array(v)
-            elif isinstance(global_config['input_list'], dict):
-                if k in global_config['input_list'].keys():
-                    data_input[global_config['input_list'][k]] = np.array(v)
-
-        outs = exe.run(val_program,
-                       feed=data_input,
-                       fetch_list=fetch_targets,
-                       return_numpy=False)
-        res = {}
-        if 'arch' in global_config and global_config['arch'] == 'YOLOv5':
+    bboxes_list, bbox_nums_list, image_id_list = [], [], []
+    with tqdm(
+            total=len(val_loader),
+            bar_format='Evaluation stage, Run batch:|{bar}| {n_fmt}/{total_fmt}',
+            ncols=80) as t:
+        for data in val_loader:
+            data_all = {k: np.array(v) for k, v in data.items()}
+            outs = exe.run(val_program,
+                           feed={feed_target_names[0]: data_all['image']},
+                           fetch_list=fetch_targets,
+                           return_numpy=False)
+            res = {}
             postprocess = YOLOv5PostProcess(
-                score_threshold=0.001, nms_threshold=0.6, multi_label=True)
+                score_threshold=0.001, nms_threshold=0.65, multi_label=True)
             res = postprocess(np.array(outs[0]), data_all['scale_factor'])
-        else:
-            for out in outs:
-                v = np.array(out)
-                if len(v.shape) > 1:
-                    res['bbox'] = v
-                else:
-                    res['bbox_num'] = v
-        metric.update(data_all, res)
-        if batch_id % 100 == 0:
-            print('Eval iter:', batch_id)
-    metric.accumulate()
-    metric.log()
-    metric.reset()
+            bboxes_list.append(res['bbox'])
+            bbox_nums_list.append(res['bbox_num'])
+            image_id_list.append(np.array(data_all['im_id']))
+            t.update()
+
+    coco_metric(anno_file, bboxes_list, bbox_nums_list, image_id_list)
 
 
 def main():
     global global_config
     all_config = load_slim_config(FLAGS.config_path)
     global_config = all_config["Global"]
-    reader_cfg = load_config(global_config['reader_config'])
 
-    dataset = reader_cfg['EvalDataset']
     global val_loader
-    val_loader = create('EvalReader')(reader_cfg['EvalDataset'],
-                                      reader_cfg['worker_num'],
-                                      return_list=True)
-    metric = None
-    if reader_cfg['metric'] == 'COCO':
-        clsid2catid = {v: k for k, v in dataset.catid2clsid.items()}
-        anno_file = dataset.get_anno()
-        metric = COCOMetric(
-            anno_file=anno_file, clsid2catid=clsid2catid, IouType='bbox')
-    elif reader_cfg['metric'] == 'VOC':
-        metric = VOCMetric(
-            label_list=dataset.get_label_list(),
-            class_num=reader_cfg['num_classes'],
-            map_type=reader_cfg['map_type'])
-    else:
-        raise ValueError("metric currently only supports COCO and VOC.")
-    global_config['metric'] = metric
+    dataset = COCOValDataset(
+        dataset_dir=global_config['dataset_dir'],
+        image_dir=global_config['val_image_dir'],
+        anno_path=global_config['val_anno_path'])
+    global anno_file
+    anno_file = dataset.ann_file
+    val_loader = paddle.io.DataLoader(dataset, batch_size=1)
 
     eval()
 

example/auto_compression/pytorch_yolov5/paddle_trt_infer.py

@@ -241,12 +241,9 @@ def predict_image(predictor,
                   image_shape=[640, 640],
                   warmup=1,
                   repeats=1,
-                  threshold=0.5,
-                  arch='YOLOv5'):
+                  threshold=0.5):
     img, scale_factor = image_preprocess(image_file, image_shape)
-    inputs = {}
-    if arch == 'YOLOv5':
-        inputs['x2paddle_images'] = img
+    inputs['x2paddle_images'] = img
     input_names = predictor.get_input_names()
     for i in range(len(input_names)):
         input_tensor = predictor.get_input_handle(input_names[i])

example/auto_compression/pytorch_yolov5/post_process.py

@@ -14,6 +14,8 @@
 
 import numpy as np
 import cv2
+import json
+import sys
 
 
 def box_area(boxes):
@@ -171,3 +173,59 @@ def __call__(self, outs, scale_factor):
         bboxs = np.concatenate(bboxs, axis=0)
         box_nums = np.array(box_nums)
         return {'bbox': bboxs, 'bbox_num': box_nums}
+
+
+def coco_metric(anno_file, bboxes_list, bbox_nums_list, image_id_list):
+    try:
+        from pycocotools.coco import COCO
+        from pycocotools.cocoeval import COCOeval
+    except:
+        print(
+            "[ERROR] Not found pycocotools, please install by `pip install pycocotools`"
+        )
+        sys.exit(1)
+
+    coco_gt = COCO(anno_file)
+    cats = coco_gt.loadCats(coco_gt.getCatIds())
+    clsid2catid = {i: cat['id'] for i, cat in enumerate(cats)}
+    results = []
+    for bboxes, bbox_nums, image_id in zip(bboxes_list, bbox_nums_list,
+                                           image_id_list):
+        results += _get_det_res(bboxes, bbox_nums, image_id, clsid2catid)
+
+    output = "bbox.json"
+    with open(output, 'w') as f:
+        json.dump(results, f)
+
+    coco_dt = coco_gt.loadRes(output)
+    coco_eval = COCOeval(coco_gt, coco_dt, 'bbox')
+    coco_eval.evaluate()
+    coco_eval.accumulate()
+    coco_eval.summarize()
+    return coco_eval.stats
+
+
+def _get_det_res(bboxes, bbox_nums, image_id, label_to_cat_id_map):
+    det_res = []
+    k = 0
+    for i in range(len(bbox_nums)):
+        cur_image_id = int(image_id[i][0])
+        det_nums = bbox_nums[i]
+        for j in range(det_nums):
+            dt = bboxes[k]
+            k = k + 1
+            num_id, score, xmin, ymin, xmax, ymax = dt.tolist()
+            if int(num_id) < 0:
+                continue
+            category_id = label_to_cat_id_map[int(num_id)]
+            w = xmax - xmin
+            h = ymax - ymin
+            bbox = [xmin, ymin, w, h]
+            dt_res = {
+                'image_id': cur_image_id,
+                'category_id': category_id,
+                'bbox': bbox,
+                'score': score
+            }
+            det_res.append(dt_res)
+    return det_res