Draw segmentation on image
(im,
np_segms,
np_label,
np_score,
labels,
threshold=0.5,
alpha=0.7)
| 193 | |
| 194 | |
| 195 | def draw_segm(im, |
| 196 | np_segms, |
| 197 | np_label, |
| 198 | np_score, |
| 199 | labels, |
| 200 | threshold=0.5, |
| 201 | alpha=0.7): |
| 202 | """ |
| 203 | Draw segmentation on image |
| 204 | """ |
| 205 | mask_color_id = 0 |
| 206 | w_ratio = .4 |
| 207 | color_list = get_color_map_list(len(labels)) |
| 208 | im = np.array(im).astype('float32') |
| 209 | clsid2color = {} |
| 210 | np_segms = np_segms.astype(np.uint8) |
| 211 | for i in range(np_segms.shape[0]): |
| 212 | mask, score, clsid = np_segms[i], np_score[i], np_label[i] |
| 213 | if score < threshold: |
| 214 | continue |
| 215 | |
| 216 | if clsid not in clsid2color: |
| 217 | clsid2color[clsid] = color_list[clsid] |
| 218 | color_mask = clsid2color[clsid] |
| 219 | for c in range(3): |
| 220 | color_mask[c] = color_mask[c] * (1 - w_ratio) + w_ratio * 255 |
| 221 | idx = np.nonzero(mask) |
| 222 | color_mask = np.array(color_mask) |
| 223 | idx0 = np.minimum(idx[0], im.shape[0] - 1) |
| 224 | idx1 = np.minimum(idx[1], im.shape[1] - 1) |
| 225 | im[idx0, idx1, :] *= 1.0 - alpha |
| 226 | im[idx0, idx1, :] += alpha * color_mask |
| 227 | sum_x = np.sum(mask, axis=0) |
| 228 | x = np.where(sum_x > 0.5)[0] |
| 229 | sum_y = np.sum(mask, axis=1) |
| 230 | y = np.where(sum_y > 0.5)[0] |
| 231 | x0, x1, y0, y1 = x[0], x[-1], y[0], y[-1] |
| 232 | cv2.rectangle(im, (x0, y0), (x1, y1), |
| 233 | tuple(color_mask.astype('int32').tolist()), 1) |
| 234 | bbox_text = '%s %.2f' % (labels[clsid], score) |
| 235 | t_size = cv2.getTextSize(bbox_text, 0, 0.3, thickness=1)[0] |
| 236 | cv2.rectangle(im, (x0, y0), (x0 + t_size[0], y0 - t_size[1] - 3), |
| 237 | tuple(color_mask.astype('int32').tolist()), -1) |
| 238 | cv2.putText( |
| 239 | im, |
| 240 | bbox_text, (x0, y0 - 2), |
| 241 | cv2.FONT_HERSHEY_SIMPLEX, |
| 242 | 0.3, (0, 0, 0), |
| 243 | 1, |
| 244 | lineType=cv2.LINE_AA) |
| 245 | return Image.fromarray(im.astype('uint8')) |
| 246 | |
| 247 | |
| 248 | def get_color(idx): |
no test coverage detected