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] labels (list): labels:['class1', ..., 'classn'] threshold (float): threshold of box
(im, np_boxes, labels, threshold=0.5)
| 124 | |
| 125 | |
| 126 | def draw_box(im, np_boxes, labels, threshold=0.5): |
| 127 | """ |
| 128 | Args: |
| 129 | im (PIL.Image.Image): PIL image |
| 130 | np_boxes (np.ndarray): shape:[N,6], N: number of box, |
| 131 | matix element:[class, score, x_min, y_min, x_max, y_max] |
| 132 | labels (list): labels:['class1', ..., 'classn'] |
| 133 | threshold (float): threshold of box |
| 134 | Returns: |
| 135 | im (PIL.Image.Image): visualized image |
| 136 | """ |
| 137 | draw_thickness = min(im.size) // 320 |
| 138 | draw = ImageDraw.Draw(im) |
| 139 | clsid2color = {} |
| 140 | color_list = get_color_map_list(len(labels)) |
| 141 | expect_boxes = (np_boxes[:, 1] > threshold) & (np_boxes[:, 0] > -1) |
| 142 | np_boxes = np_boxes[expect_boxes, :] |
| 143 | |
| 144 | vis_order = False |
| 145 | if len(np_boxes) > 0 and len(np_boxes[0]) == 7: |
| 146 | np_boxes = sorted(np_boxes, key=lambda x: x[6]) |
| 147 | vis_order = True |
| 148 | |
| 149 | centers = [] |
| 150 | for dt in np_boxes: |
| 151 | if len(dt) == 7: |
| 152 | clsid, bbox, score, read_order = int(dt[0]), dt[2:6], dt[1], int(dt[6]) |
| 153 | else: |
| 154 | clsid, bbox, score = int(dt[0]), dt[2:], dt[1] |
| 155 | if clsid not in clsid2color: |
| 156 | clsid2color[clsid] = color_list[clsid] |
| 157 | color = tuple(clsid2color[clsid]) |
| 158 | |
| 159 | if len(bbox) == 4: |
| 160 | xmin, ymin, xmax, ymax = bbox |
| 161 | print('class_id:{:d}, confidence:{:.4f}, left_top:[{:.2f},{:.2f}],' |
| 162 | 'right_bottom:[{:.2f},{:.2f}]'.format( |
| 163 | int(clsid), score, xmin, ymin, xmax, ymax)) |
| 164 | # draw bbox |
| 165 | draw.line( |
| 166 | [(xmin, ymin), (xmin, ymax), (xmax, ymax), (xmax, ymin), |
| 167 | (xmin, ymin)], |
| 168 | width=draw_thickness, |
| 169 | fill=color) |
| 170 | cx, cy = int((xmin + xmax)/2), int((ymin + ymax)/2) |
| 171 | centers.append((cx, cy)) |
| 172 | elif len(bbox) == 8: |
| 173 | x1, y1, x2, y2, x3, y3, x4, y4 = bbox |
| 174 | draw.line( |
| 175 | [(x1, y1), (x2, y2), (x3, y3), (x4, y4), (x1, y1)], |
| 176 | width=2, |
| 177 | fill=color) |
| 178 | xmin = min(x1, x2, x3, x4) |
| 179 | ymin = min(y1, y2, y3, y4) |
| 180 | |
| 181 | # draw label |
| 182 | text = "{} {:.4f}".format(labels[clsid], score) |
| 183 | tw, th = imagedraw_textsize_c(draw, text) |
no test coverage detected