forked from PulseFocusPlatform/PulseFocusPlatform
214 lines
7.5 KiB
Python
214 lines
7.5 KiB
Python
# coding: utf-8
|
|
# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
from __future__ import division
|
|
|
|
import cv2
|
|
import numpy as np
|
|
from PIL import Image, ImageDraw
|
|
|
|
|
|
def visualize_box_mask(im, results, labels, 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, im_h, im_w]
|
|
labels (list): labels:['class1', ..., 'classn']
|
|
threshold (float): Threshold of score.
|
|
Returns:
|
|
im (PIL.Image.Image): visualized image
|
|
"""
|
|
if isinstance(im, str):
|
|
im = Image.open(im).convert('RGB')
|
|
else:
|
|
im = Image.fromarray(im)
|
|
if 'masks' in results and 'boxes' in results and len(results['boxes']) > 0:
|
|
im = draw_mask(
|
|
im, results['boxes'], results['masks'], labels, threshold=threshold)
|
|
if 'boxes' in results and len(results['boxes']) > 0:
|
|
im = draw_box(im, results['boxes'], labels, threshold=threshold)
|
|
if 'segm' in results:
|
|
im = draw_segm(
|
|
im,
|
|
results['segm'],
|
|
results['label'],
|
|
results['score'],
|
|
labels,
|
|
threshold=threshold)
|
|
return im
|
|
|
|
|
|
def get_color_map_list(num_classes):
|
|
"""
|
|
Args:
|
|
num_classes (int): number of class
|
|
Returns:
|
|
color_map (list): RGB color list
|
|
"""
|
|
color_map = num_classes * [0, 0, 0]
|
|
for i in range(0, num_classes):
|
|
j = 0
|
|
lab = i
|
|
while lab:
|
|
color_map[i * 3] |= (((lab >> 0) & 1) << (7 - j))
|
|
color_map[i * 3 + 1] |= (((lab >> 1) & 1) << (7 - j))
|
|
color_map[i * 3 + 2] |= (((lab >> 2) & 1) << (7 - j))
|
|
j += 1
|
|
lab >>= 3
|
|
color_map = [color_map[i:i + 3] for i in range(0, len(color_map), 3)]
|
|
return color_map
|
|
|
|
|
|
def draw_mask(im, np_boxes, np_masks, labels, 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, im_h, im_w]
|
|
labels (list): labels:['class1', ..., 'classn']
|
|
threshold (float): threshold of mask
|
|
Returns:
|
|
im (PIL.Image.Image): visualized image
|
|
"""
|
|
color_list = get_color_map_list(len(labels))
|
|
w_ratio = 0.4
|
|
alpha = 0.7
|
|
im = np.array(im).astype('float32')
|
|
clsid2color = {}
|
|
expect_boxes = (np_boxes[:, 1] > threshold) & (np_boxes[:, 0] > -1)
|
|
np_boxes = np_boxes[expect_boxes, :]
|
|
np_masks = np_masks[expect_boxes, :, :]
|
|
for i in range(len(np_masks)):
|
|
clsid, score = int(np_boxes[i][0]), np_boxes[i][1]
|
|
mask = np_masks[i]
|
|
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(mask)
|
|
color_mask = np.array(color_mask)
|
|
im[idx[0], idx[1], :] *= 1.0 - alpha
|
|
im[idx[0], idx[1], :] += alpha * color_mask
|
|
return Image.fromarray(im.astype('uint8'))
|
|
|
|
|
|
def draw_box(im, np_boxes, labels, 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]
|
|
labels (list): labels:['class1', ..., 'classn']
|
|
threshold (float): threshold of box
|
|
Returns:
|
|
im (PIL.Image.Image): visualized image
|
|
"""
|
|
draw_thickness = min(im.size) // 320
|
|
draw = ImageDraw.Draw(im)
|
|
clsid2color = {}
|
|
color_list = get_color_map_list(len(labels))
|
|
expect_boxes = (np_boxes[:, 1] > threshold) & (np_boxes[:, 0] > -1)
|
|
np_boxes = np_boxes[expect_boxes, :]
|
|
|
|
for dt in np_boxes:
|
|
clsid, bbox, score = int(dt[0]), dt[2:], dt[1]
|
|
if clsid not in clsid2color:
|
|
clsid2color[clsid] = color_list[clsid]
|
|
color = tuple(clsid2color[clsid])
|
|
|
|
if len(bbox) == 4:
|
|
xmin, ymin, xmax, ymax = bbox
|
|
print('class_id:{:d}, confidence:{:.4f}, left_top:[{:.2f},{:.2f}],'
|
|
'right_bottom:[{:.2f},{:.2f}]'.format(
|
|
int(clsid), score, xmin, ymin, xmax, ymax))
|
|
# draw bbox
|
|
draw.line(
|
|
[(xmin, ymin), (xmin, ymax), (xmax, ymax), (xmax, ymin),
|
|
(xmin, ymin)],
|
|
width=draw_thickness,
|
|
fill=color)
|
|
elif len(bbox) == 8:
|
|
x1, y1, x2, y2, x3, y3, x4, y4 = bbox
|
|
draw.line(
|
|
[(x1, y1), (x2, y2), (x3, y3), (x4, y4), (x1, y1)],
|
|
width=2,
|
|
fill=color)
|
|
xmin = min(x1, x2, x3, x4)
|
|
ymin = min(y1, y2, y3, y4)
|
|
|
|
# draw label
|
|
text = "{} {:.4f}".format(labels[clsid], score)
|
|
tw, th = draw.textsize(text)
|
|
draw.rectangle(
|
|
[(xmin + 1, ymin - th), (xmin + tw + 1, ymin)], fill=color)
|
|
draw.text((xmin + 1, ymin - th), text, fill=(255, 255, 255))
|
|
return im
|
|
|
|
|
|
def draw_segm(im,
|
|
np_segms,
|
|
np_label,
|
|
np_score,
|
|
labels,
|
|
threshold=0.5,
|
|
alpha=0.7):
|
|
"""
|
|
Draw segmentation on image
|
|
"""
|
|
mask_color_id = 0
|
|
w_ratio = .4
|
|
color_list = get_color_map_list(len(labels))
|
|
im = np.array(im).astype('float32')
|
|
clsid2color = {}
|
|
np_segms = np_segms.astype(np.uint8)
|
|
for i in range(np_segms.shape[0]):
|
|
mask, score, clsid = np_segms[i], np_score[i], np_label[i]
|
|
if score < threshold:
|
|
continue
|
|
|
|
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(mask)
|
|
color_mask = np.array(color_mask)
|
|
im[idx[0], idx[1], :] *= 1.0 - alpha
|
|
im[idx[0], idx[1], :] += alpha * color_mask
|
|
sum_x = np.sum(mask, axis=0)
|
|
x = np.where(sum_x > 0.5)[0]
|
|
sum_y = np.sum(mask, axis=1)
|
|
y = np.where(sum_y > 0.5)[0]
|
|
x0, x1, y0, y1 = x[0], x[-1], y[0], y[-1]
|
|
cv2.rectangle(im, (x0, y0), (x1, y1),
|
|
tuple(color_mask.astype('int32').tolist()), 1)
|
|
bbox_text = '%s %.2f' % (labels[clsid], score)
|
|
t_size = cv2.getTextSize(bbox_text, 0, 0.3, thickness=1)[0]
|
|
cv2.rectangle(im, (x0, y0), (x0 + t_size[0], y0 - t_size[1] - 3),
|
|
tuple(color_mask.astype('int32').tolist()), -1)
|
|
cv2.putText(
|
|
im,
|
|
bbox_text, (x0, y0 - 2),
|
|
cv2.FONT_HERSHEY_SIMPLEX,
|
|
0.3, (0, 0, 0),
|
|
1,
|
|
lineType=cv2.LINE_AA)
|
|
return Image.fromarray(im.astype('uint8'))
|