PaddleDetection在python端的半自动标注

PaddleDetection静态版本的预测代码主要放置在该文件夹目录下，半自动标注原理很简单，就是要利用一个之前训练好的模型去预测一批图片，在预测图片的同时运用opencv的找外轮廓把mask的外轮廓点描绘出来，储存在一个json文件里面。
下面是具体修改的代码实例：

1. visualize.py

def draw_mask(im, np_boxes, np_masks, labels, resolution=14, threshold=0.5):
   """
   Args:
       im (PIL.Image.Image): PIL image
       np_boxes (np.ndarray): shape:[N,6], N: number of box,
                              matix element:[class, score, x_min, y_min, x_max, y_max]
       np_masks (np.ndarray): shape:[N, class_num, resolution, resolution]
       labels (list): labels:['class1', ..., 'classn']   模型标签列表
       resolution (int): shape of a mask is:[resolution, resolution]
       threshold (float): threshold of mask
   Returns:
       im (PIL.Image.Image): visualized image
   """
   color_list = get_color_map_list(len(labels))
   scale = (resolution + 2.0) / resolution
   im_w, im_h = im.size
   w_ratio = 0.4
   alpha = 0.7
   im = np.array(im).astype('float32')
   rects = np_boxes[:, 2:]
   expand_rects = expand_boxes(rects, scale)
   expand_rects = expand_rects.astype(np.int32)
   clsid_scores = np_boxes[:, 0:2]
   padded_mask = np.zeros((resolution + 2, resolution + 2), dtype=np.float32)
   clsid2color = {}
   **json_data = {}  # 搞一个字典  存jsondata
   label_name = {'Color_PP': 0, 'White_PP': 0, 'Black_PP': 0, 'Trans_PP': 0,
                 'White_HDPE': 0, 'Color_HDPE': 0,'Bottle_PET': 0, 'Trans_PET': 0, 'Hollow_PVC': 0, 'Hollow_PP': 0, 'Trans_HDPE': 0, 'Others': 0}**

   for idx in range(len(np_boxes)):              #以矩形框来遍历
       clsid, score = clsid_scores[idx].tolist()
       clsid = int(clsid)                       #第几个矩形框的类别
       label=labels[clsid]                      #将数字标签转换为文字
       xmin, ymin, xmax, ymax = expand_rects[idx].tolist()
       w = xmax - xmin + 1
       h = ymax - ymin + 1
       w = np.maximum(w, 1)
       h = np.maximum(h, 1)
       padded_mask[1:-1, 1:-1] = np_masks[idx, int(clsid), :, :]       #取出该类的mask
       resized_mask = cv2.resize(padded_mask, (w, h))             #mask还原成原图的长宽
       resized_mask = np.array(resized_mask > threshold, dtype=np.uint8)
       x0 = min(max(xmin, 0), im_w)
       x1 = min(max(xmax + 1, 0), im_w)
       y0 = min(max(ymin, 0), im_h)
       y1 = min(max(ymax + 1, 0), im_h)
       im_mask = np.zeros((im_h, im_w), dtype=np.uint8)             #创建一张0值图
       im_mask[y0:y1, x0:x1] = resized_mask[(y0 - ymin):(y1 - ymin), (
           x0 - xmin):(x1 - xmin)]                                    #将mask移动上去

       if clsid not in clsid2color:
           clsid2color[clsid] = color_list[clsid]
       color_mask = clsid2color[clsid]
       for c in range(3):
           color_mask[c] = color_mask[c] * (1 - w_ratio) + w_ratio * 255
       idx = np.nonzero(im_mask)
       color_mask = np.array(color_mask)
       im[idx[0], idx[1], :] *= 1.0 - alpha
       im[idx[0], idx[1], :] += alpha * color_mask
   **#找轮廓
       point = 20      #隔20个点取一次，数字越大轮廓点越稀疏
       contours, hierarchy = cv2.findContours(im_mask,cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)   #寻找外轮廓，模式为所有点
       for verts in contours:
           vert=[]
           for m in range(len(verts)):
               if m%point==0 :
                   vert.append(verts[m][0].tolist())
           if len(vert)<4:
               continue
           label_name[label]+=1
           json_data[label+str(label_name[label])]=vert         #将mask写入字典
   #print(json_data)
   return Image.fromarray(im.astype('uint8')), json_data     #将json_data返回**

2. #draw_mask函数 在此处被调用 def visualize_box_mask

def visualize_box_mask(im, results, labels, mask_resolution=14, threshold=0.5):
    """
    Args:
        im (str/np.ndarray): path of image/np.ndarray read by cv2
        results (dict): include 'boxes': np.ndarray: shape:[N,6], N: number of box,
                        matix element:[class, score, x_min, y_min, x_max, y_max]
                        MaskRCNN's results include 'masks': np.ndarray:
                        shape:[N, class_num, mask_resolution, mask_resolution]
        labels (list): labels:['class1', ..., 'classn']
        mask_resolution (int): shape of a mask is:[mask_resolution, mask_resolution]
        threshold (float): Threshold of score.
    Returns:
        im (PIL.Image.Image): visualized image
    """
    img_path=im             #保存原始的im（图像路径），后面解码就找不到了。
    if isinstance(im, str):
        im = Image.open(im).convert('RGB')
    else:
        im = Image.fromarray(im)
    if 'masks' and 'boxes' in results :      #有mask和boxes进入这个循环
        im,json_data = draw_mask(            #两个返回值
            im,
            results['boxes'],
            results['masks'],
            labels,
            resolution=mask_resolution)
     #将每张图片的json文件写入
        json_path =img_path[0: img_path.rindex('.')] +".json" #json文件写入的路径，
       与图片放在一块方便打开。
        with open(json_path, 'w') as fw:
            data = {}
            with open(json_path, 'r') as f:
                data.update({'version': '4.5.6'})
                data.update({'flags': {}})
                shapes = []
                for key, value in json_data.items():
                    shape = {}
                    key = key[:-1]
                    shape.update({'label': key})
                    shape.update({'points': value})
                    shape.update({"group_id": None})
                    shape.update({"shape_type": "polygon"})
                    shape.update({"flags": {}})
                    shapes.append(shape)
                data.update({'shapes': shapes})
                jie = img_path[img_path.rindex('/') + 1: len(img_path)]
                data.update({'imagePath':jie})
                data.update({'imageData': None})
               #图片的长宽，根据自己的图片大小相应设置
                data.update({'imageHeight': 1400}) 
                data.update({'imageWidth': 1400})
            json.dump(data, fw,indent=2)

    if 'boxes' in results:
        im = draw_box(im, results['boxes'], labels)
    if 'segm' in results:
        im = draw_segm(
            im,
            results['segm'],
            results['label'],
            results['score'],
            labels,
            threshold=threshold)
    if 'landmark' in results:
        im = draw_lmk(im, results['landmark'])
    return im

半自动标注的质量是由初始预测模型的质量来决定的，我这个模型比较垃圾，预测的实际掩膜就会是这个样子。正常模型会比这个好很多啦。