Bladeren bron

Merge branch 'develop_qh' of https://github.com/FlyingQianMM/PaddleX into develop_qh

FlyingQianMM 4 jaren geleden
bovenliggende
commit
72954ad57c

+ 2 - 1
docs/apis/visualize.md

@@ -5,7 +5,7 @@ PaddleX提供了一系列模型预测和结果分析的可视化函数。
 ## paddlex.det.visualize
 > **目标检测/实例分割预测结果可视化**  
 ```
-paddlex.det.visualize(image, result, threshold=0.5, save_dir='./')
+paddlex.det.visualize(image, result, threshold=0.5, save_dir='./', color=None)
 ```
 将目标检测/实例分割模型预测得到的Box框和Mask在原图上进行可视化。
 
@@ -14,6 +14,7 @@ paddlex.det.visualize(image, result, threshold=0.5, save_dir='./')
 > * **result** (str): 模型预测结果。
 > * **threshold**(float): score阈值,将Box置信度低于该阈值的框过滤不进行可视化。默认0.5
 > * **save_dir**(str): 可视化结果保存路径。若为None,则表示不保存,该函数将可视化的结果以np.ndarray的形式返回;若设为目录路径,则将可视化结果保存至该目录下。默认值为'./'。
+> * **color**(list|tuple|np.array): 各类别的BGR颜色值组成的数组,形状为Nx3(N为类别数量),数值范围为[0, 255]。例如针对2个类别的[[255, 0, 0], [0, 255, 0]]。若为None,则自动生成各类别的颜色。默认值为None。
 
 ### 使用示例
 > 点击下载如下示例中的[模型](https://bj.bcebos.com/paddlex/models/xiaoduxiong_epoch_12.tar.gz)

+ 41 - 4
paddlex/cv/models/load_model.py

@@ -55,6 +55,18 @@ def load_model(model_dir, fixed_input_shape=None):
                              format(fixed_input_shape))
                 model.fixed_input_shape = fixed_input_shape
 
+    if info['Model'].count('RCNN') > 0:
+        if info['_init_params']['with_fpn']:
+            if model.fixed_input_shape is not None:
+                if model.fixed_input_shape[0] % 32 > 0:
+                    raise Exception(
+                        "The first value in fixed_input_shape must be a multiple of 32, but recieved {}.".
+                        format(model.fixed_input_shape[0]))
+                if model.fixed_input_shape[1] % 32 > 0:
+                    raise Exception(
+                        "The second value in fixed_input_shape must be a multiple of 32, but recieved {}.".
+                        format(model.fixed_input_shape[1]))
+
     with fluid.scope_guard(model_scope):
         if status == "Normal" or \
                 status == "Prune" or status == "fluid.save":
@@ -137,12 +149,37 @@ def fix_input_shape(info, fixed_input_shape=None):
                 if list(info['Transforms'][i].keys())[0] == 'Resize':
                     resize_op_index = i
             if resize_op_index is not None:
-                info['Transforms'][resize_op_index]['Resize']['target_size'] = fixed_input_shape[0]
-        else:
+                info['Transforms'][resize_op_index]['Resize'][
+                    'target_size'] = fixed_input_shape[0]
+        elif info['Model'].count('RCNN') > 0:
+            resize_op_index = None
+            for i in range(len(info['Transforms'])):
+                if list(info['Transforms'][i].keys())[0] == 'ResizeByShort':
+                    resize_op_index = i
+            if resize_op_index is not None:
+                info['Transforms'][resize_op_index]['ResizeByShort'][
+                    'short_size'] = min(fixed_input_shape)
+                info['Transforms'][resize_op_index]['ResizeByShort'][
+                    'max_size'] = max(fixed_input_shape)
+            else:
+                resize['ResizeByShort']['short_size'] = min(fixed_input_shape)
+                resize['ResizeByShort']['max_size'] = max(fixed_input_shape)
+                info['Transforms'].append(resize)
+
+            padding_op_index = None
+            for i in range(len(info['Transforms'])):
+                if list(info['Transforms'][i].keys())[0] == 'Padding':
+                    padding_op_index = i
+            if padding_op_index is not None:
+                info['Transforms'][padding_op_index]['Padding'][
+                    'target_size'] = list(fixed_input_shape)
+            else:
+                padding['Padding']['target_size'] = list(fixed_input_shape)
+                info['Transforms'].append(padding)
+        elif info['_Attributes']['model_type'] == 'segmenter':
             resize['ResizeByShort']['short_size'] = min(fixed_input_shape)
             resize['ResizeByShort']['max_size'] = max(fixed_input_shape)
             padding['Padding']['target_size'] = list(fixed_input_shape)
-            if info['_Attributes']['model_type'] == 'segmenter':
-                padding['Padding']['im_padding_value'] = [0.] * input_channel
+            padding['Padding']['im_padding_value'] = [0.] * input_channel
             info['Transforms'].append(resize)
             info['Transforms'].append(padding)

+ 6 - 2
paddlex/cv/models/utils/pretrain_weights.py

@@ -160,7 +160,8 @@ def get_pretrain_weights(flag, class_name, backbone, save_dir):
             logging.warning(warning_info.format(class_name, flag, 'IMAGENET'))
             flag = 'IMAGENET'
         elif class_name == 'FastSCNN':
-            logging.warning(warning_info.format(class_name, flag, 'CITYSCAPES'))
+            logging.warning(
+                warning_info.format(class_name, flag, 'CITYSCAPES'))
             flag = 'CITYSCAPES'
     elif flag == 'CITYSCAPES':
         model_name = '{}_{}'.format(class_name, backbone)
@@ -183,7 +184,8 @@ def get_pretrain_weights(flag, class_name, backbone, save_dir):
             logging.warning(warning_info.format(class_name, flag, 'COCO'))
             flag = 'COCO'
         elif class_name == 'FastSCNN':
-            logging.warning(warning_info.format(class_name, flag, 'CITYSCAPES'))
+            logging.warning(
+                warning_info.format(class_name, flag, 'CITYSCAPES'))
             flag = 'CITYSCAPES'
     elif flag == 'BAIDU10W':
         if class_name not in ['ResNet50_vd']:
@@ -254,6 +256,8 @@ def get_pretrain_weights(flag, class_name, backbone, save_dir):
         if getattr(paddlex, 'gui_mode', False):
             paddlex.utils.download_and_decompress(url, path=new_save_dir)
             return osp.join(new_save_dir, fname)
+
+        import paddlehub as hub
         try:
             logging.info(
                 "Connecting PaddleHub server to get pretrain weights...")

+ 20 - 5
paddlex/cv/models/utils/visualize.py

@@ -23,7 +23,11 @@ from .detection_eval import fixed_linspace, backup_linspace, loadRes
 from paddlex.cv.datasets.dataset import is_pic
 
 
-def visualize_detection(image, result, threshold=0.5, save_dir='./'):
+def visualize_detection(image,
+                        result,
+                        threshold=0.5,
+                        save_dir='./',
+                        color=None):
     """
         Visualize bbox and mask results
     """
@@ -34,7 +38,7 @@ def visualize_detection(image, result, threshold=0.5, save_dir='./'):
         image_name = os.path.split(image)[-1]
         image = cv2.imread(image)
 
-    image = draw_bbox_mask(image, result, threshold=threshold)
+    image = draw_bbox_mask(image, result, threshold=threshold, color_map=color)
     if save_dir is not None:
         if not os.path.exists(save_dir):
             os.makedirs(save_dir)
@@ -159,7 +163,7 @@ def clip_bbox(bbox):
     return xmin, ymin, xmax, ymax
 
 
-def draw_bbox_mask(image, results, threshold=0.5):
+def draw_bbox_mask(image, results, threshold=0.5, color_map=None):
     import matplotlib
     matplotlib.use('Agg')
     import matplotlib as mpl
@@ -201,7 +205,18 @@ def draw_bbox_mask(image, results, threshold=0.5):
     for dt in np.array(results):
         if dt['category'] not in labels:
             labels.append(dt['category'])
-    color_map = get_color_map_list(256)
+
+    if color_map is None:
+        color_map = get_color_map_list(len(labels) + 2)[2:]
+    else:
+        color_map = np.asarray(color_map)
+        if color_map.shape[0] != len(labels) or color_map.shape[1] != 3:
+            raise Exception(
+                "The shape for color_map is required to be {}x3, but recieved shape is {}x{}.".
+                format(len(labels), color_map.shape))
+        if np.max(color_map) > 255 or np.min(color_map) < 0:
+            raise ValueError(
+                " The values in color_map should be within 0-255 range.")
 
     keep_results = []
     areas = []
@@ -222,7 +237,7 @@ def draw_bbox_mask(image, results, threshold=0.5):
         xmax = xmin + w
         ymax = ymin + h
 
-        color = tuple(color_map[labels.index(cname) + 2])
+        color = tuple(color_map[labels.index(cname)])
         color = [c / 255. for c in color]
         # draw bbox
         ax.add_patch(

+ 1 - 0
requirements.txt

@@ -10,3 +10,4 @@ paddle2onnx
 paddlepaddle-gpu
 opencv-python
 psutil
+flask_cors