def vis_one_image_opencv(
im, boxes, segms=None, keypoints=None, thresh=0.9, kp_thresh=2,
show_box=False, dataset=None, show_class=False):
if isinstance(boxes, list):
boxes, segms, keypoints, classes = convert_from_cls_format(
boxes, segms, keypoints)
if boxes is None or boxes.shape[0] == 0 or max(boxes[:, 4]) < thresh:
return im
if segms is not None and len(segms) > 0:
masks = mask_util.decode(segms)
color_list = colormap()
mask_color_id = 0
areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
sorted_inds = np.argsort(-areas)
class_shuzu=[]
for i in sorted_inds:
bbox = boxes[i, :4]
score = boxes[i, -1]
if score < thresh:
continue
if show_box:
im = vis_bbox(
im, (bbox[0], bbox[1], bbox[2] - bbox[0], bbox[3] - bbox[1]))
if show_class:
class_str = get_class_string(classes[i], score, dataset)
print ("woca**************************class_str")
print (class_str)
im = vis_class(im, (bbox[0], bbox[1] - 2), class_str)
if segms is not None and len(segms) > i:
color_mask = color_list[mask_color_id % len(color_list), 0:3]
mask_color_id += 1
im = vis_mask(im, masks[..., i], color_mask)
if keypoints is not None and len(keypoints) > i:
im = vis_keypoints(im, keypoints[i], kp_thresh)
return (im)
def write_txt(fid, img_idx, boxes, segms=None, keypoints=None, thresh=0.7):
if isinstance(boxes, list):
boxes, segms, keypoints, classes = convert_from_cls_format(
boxes, segms, keypoints)
if boxes is None or boxes.shape[0] == 0 or max(boxes[:, 4]) < thresh:
return
if segms is not None and len(segms) > 0:
masks = mask_util.decode(segms)
color_list = colormap()
mask_color_id = 0
# Display in largest to smallest order to reduce occlusion
# (x1, y1, x2, y2) box format
areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
sorted_inds = np.argsort(-areas)
for i in sorted_inds:
# just show person
if classes[i] != 1:
continue
bbox = boxes[i, :4]
score = boxes[i, -1]
if score < thresh:
continue
txt_tmp = "%d,-1,%.1f,%.1f,%.1f,%.1f,%.3f\n" % (img_idx, bbox[0], bbox[1], bbox[2] - bbox[0], bbox[3] - bbox[1], score)
fid.write(txt_tmp)
def segmented_images_in_original_image_size(
im, im_name, output_dir, boxes, segms=None, keypoints=None, thresh=0.9,
kp_thresh=2, dpi=200, box_alpha=0.0, dataset=None, show_class=False,
ext='pdf'):
"""Extract segmented images."""
if not os.path.exists(output_dir):
os.makedirs(output_dir)
if isinstance(boxes, list):
boxes, segms, keypoints, classes = convert_from_cls_format(
boxes, segms, keypoints)
if boxes is None or boxes.shape[0] == 0 or max(boxes[:, 4]) < thresh:
return
if segms is not None:
masks = mask_util.decode(segms)
color_list = colormap(rgb=True) / 255
# Display in largest to smallest order to reduce occlusion
areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
sorted_inds = np.argsort(-areas)
mask_color_id = 0
segmented_images = []
segmented_classes = []
segmented_scores = []
segmented_binary_masks = []
contours = []
for i in sorted_inds:
bbox = boxes[i, :4]
score = boxes[i, -1]
if score < thresh:
continue
# show mask
if segms is not None and len(segms) > i:
img = np.zeros(im.shape)
color_mask = color_list[mask_color_id % len(color_list), 0:3]
mask_color_id += 1
w_ratio = .4
e = masks[:, :, i]
for channel in range(3):
img[:,:, channel] = im[:,:, channel] * e[:,:]
_, contour, hier = cv2.findContours(
e.copy(), cv2.RETR_CCOMP, cv2.CHAIN_APPROX_NONE)
segmented_images.insert(len(segmented_images), img)
segmented_binary_masks.insert(len(segmented_binary_masks), e)
segmented_classes.insert(len(segmented_classes), dataset.classes[classes[i]])
segmented_scores.insert(len(segmented_scores), score)
return segmented_images, segmented_classes, segmented_scores, segmented_binary_masks
def vis_one_image_opencv(im,
boxes,
segms=None,
keypoints=None,
thresh=0.9,
kp_thresh=2,
show_box=False,
dataset=None,
show_class=False):
"""Constructs a numpy array with the detections visualized."""
if isinstance(boxes, list):
boxes, segms, keypoints, classes = convert_from_cls_format(
boxes, segms, keypoints)
if boxes is None or boxes.shape[0] == 0 or max(boxes[:, 4]) < thresh:
return im
if segms is not None:
masks = mask_util.decode(segms)
color_list = colormap()
mask_color_id = 0
# Display in largest to smallest order to reduce occlusion
areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
sorted_inds = np.argsort(-areas)
for i in sorted_inds:
bbox = boxes[i, :4]
score = boxes[i, -1]
if score < thresh:
continue
# show box (off by default)
if show_box:
im = vis_bbox(
im, (bbox[0], bbox[1], bbox[2] - bbox[0], bbox[3] - bbox[1]))
# show class (off by default)
if show_class:
class_str = get_class_string(classes[i], score, dataset)
im = vis_class(im, (bbox[0], bbox[1] - 2), class_str)
# show mask
if segms is not None and len(segms) > i:
color_mask = color_list[mask_color_id % len(color_list), 0:3]
mask_color_id += 1
im = vis_mask(im, masks[..., i], color_mask)
# show keypoints
if keypoints is not None and len(keypoints) > i:
im = vis_keypoints(im, keypoints[i], kp_thresh)
return im
def vis_one_image_opencv(
im, boxes, segms=None, keypoints=None, thresh=0.9, kp_thresh=2,
show_box=False, dataset=None, show_class=False):
"""Constructs a numpy array with the detections visualized."""
if isinstance(boxes, list):
boxes, segms, keypoints, classes = convert_from_cls_format(
boxes, segms, keypoints)
if boxes is None or boxes.shape[0] == 0 or max(boxes[:, 4]) < thresh:
return im
if segms is not None:
masks = mask_util.decode(segms)
color_list = colormap()
mask_color_id = 0
# Display in largest to smallest order to reduce occlusion
areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
sorted_inds = np.argsort(-areas)
for i in sorted_inds:
bbox = boxes[i, :4]
score = boxes[i, -1]
if score < thresh:
continue
# show box (off by default)
if show_box:
im = vis_bbox(
im, (bbox[0], bbox[1], bbox[2] - bbox[0], bbox[3] - bbox[1]))
# show class (off by default)
if show_class:
class_str = get_class_string(classes[i], score, dataset)
im = vis_class(im, (bbox[0], bbox[1] - 2), class_str)
# show mask
if segms is not None and len(segms) > i:
color_mask = color_list[mask_color_id % len(color_list), 0:3]
mask_color_id += 1
im = vis_mask(im, masks[..., i], color_mask)
# show keypoints
if keypoints is not None and len(keypoints) > i:
im = vis_keypoints(im, keypoints[i], kp_thresh)
return im
def visualize_all_segmentations(img_path,
pkl_path,
dataset,
images_dir,
MIN=0,
MAX=1):
# Open image and pkl
im = imread(img_path)
cls_boxes, cls_segms, cls_keyps = pickle.load(open(pkl_path, 'rb'))
boxes, segms, keypoints, classes = convert_from_cls_format(
cls_boxes, cls_segms, cls_keyps)
if boxes is None or boxes.shape[0] == 0:
return None
if segms is not None:
masks = mask_util.decode(segms)
color_list = colormap()
mask_color_id = 0
# Display in largest to smallest order to reduce occlusion
areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
sorted_inds = np.argsort(-areas)
for i in sorted_inds:
bbox = boxes[i, :4]
score = boxes[i, -1]
cls = classes[i]
if score < MIN or score > MAX:
continue
# show bbox
im = vis_bbox(im,
(bbox[0], bbox[1], bbox[2] - bbox[0], bbox[3] - bbox[1]))
# show class
class_str = get_class_string(classes[i], score, dataset)
im = vis_class(im, (bbox[0], bbox[1] - 2), class_str)
# show mask
if segms is not None and len(segms) > i:
color_mask = color_list[mask_color_id % len(color_list), 0:3]
mask_color_id += 1
im = vis_mask(im, masks[..., i], color_mask)
# show keypoints
if keypoints is not None and len(keypoints) > i:
im = vis_keypoints(im, keypoints[i], kp_thresh)
vis_path = save(im, images_dir)
return vis_path
def vis_one_image_car(im,
im_name,
output_dir,
boxes,
segms=None,
keypoints=None,
thresh=0.9,
kp_thresh=2,
dpi=200,
box_alpha=0.0,
dataset=None,
show_class=False,
ext='jpg'):
"""Visual debugging of detections."""
if not os.path.exists(output_dir):
os.makedirs(output_dir)
if isinstance(boxes, list):
boxes, segms, keypoints, classes = convert_from_cls_format(
boxes, segms, keypoints)
if boxes is None or boxes.shape[0] == 0 or max(boxes[:, 4]) < thresh:
return
if segms is not None:
masks = mask_util.decode(segms)
color_list = colormap(rgb=True) / 255
dataset_keypoints, _ = keypoint_utils.get_keypoints()
kp_lines = kp_connections_car(dataset_keypoints)
cmap = plt.get_cmap('rainbow')
colors = [cmap(i) for i in np.linspace(0, 1, len(kp_lines) + 2)]
fig = plt.figure(frameon=False)
fig.set_size_inches(im.shape[1] / dpi, im.shape[0] / dpi)
ax = plt.Axes(fig, [0., 0., 1., 1.])
ax.axis('off')
fig.add_axes(ax)
ax.imshow(im)
# Display in largest to smallest order to reduce occlusion
areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
sorted_inds = np.argsort(-areas)
mask_color_id = 0
for i in sorted_inds:
bbox = boxes[i, :4]
score = boxes[i, -1]
if score < thresh:
continue
print('vehicle', score)
# show box (off by default, box_alpha=0.0)
ax.add_patch(
plt.Rectangle((bbox[0], bbox[1]),
bbox[2] - bbox[0],
bbox[3] - bbox[1],
fill=False,
edgecolor='g',
linewidth=0.5,
alpha=box_alpha))
if show_class:
ax.text(bbox[0],
bbox[1] - 2,
get_class_string_car(classes[i], score, dataset),
fontsize=3,
family='serif',
bbox=dict(facecolor='g',
alpha=0.4,
pad=0,
edgecolor='none'),
color='white')
# show keypoints
if keypoints is not None and len(keypoints) > i:
kps = keypoints[i]
plt.autoscale(False)
for l in range(len(kp_lines)):
i1 = kp_lines[l][0]
i2 = kp_lines[l][1]
if i1 > 13 or i2 > 13:
continue
#print(kps[2],i1,i2)
if kps[2, i1] > kp_thresh and kps[2, i2] > kp_thresh:
x = [kps[0, i1], kps[0, i2]]
y = [kps[1, i1], kps[1, i2]]
line = ax.plot(x, y)
plt.setp(line, color=colors[l], linewidth=1.0, alpha=0.7)
if kps[2, i1] > kp_thresh:
ax.plot(kps[0, i1],
kps[1, i1],
'.',
color=colors[l],
markersize=3.0,
alpha=0.7)
if kps[2, i2] > kp_thresh:
ax.plot(kps[0, i2],
kps[1, i2],
'.',
color=colors[l],
markersize=3.0,
alpha=0.7)
output_name = os.path.basename(im_name) + '.' + ext
fig.savefig(os.path.join(output_dir, '{}'.format(output_name)),
dpi=dpi)
plt.close('all')
def vis_one_image(
im, im_name, output_dir, classes, boxes, masks=None, keypoints=None, thresh=0.9,
kp_thresh=2, dpi=200, box_alpha=0.0, dataset=None, show_class=False,
ext='pdf', show=False, W=224, H=224):
"""Visual debugging of detections."""
if not os.path.exists(output_dir):
os.makedirs(output_dir)
color_list = colormap(rgb=True) / 255
fig = plt.figure(frameon=False)
fig.set_size_inches(im.shape[1] / dpi, im.shape[0] / dpi)
ax = plt.Axes(fig, [0., 0., 1., 1.])
ax.axis('off')
fig.add_axes(ax)
ax.imshow(im)
# Display in largest to smallest order to reduce occlusion
boxes[:,0] *= W
boxes[:, 2] *= W
boxes[:,1] *= H
boxes[:, 3] *= H
areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
sorted_inds = np.argsort(-areas)
# torch to numpy
# ipdb.set_trace()
if masks is not None:
# uint8
masks = masks.astype('uint8')
# rescale
w_masks, h_masks, _ = masks.shape
mask_color_id = 0
# ipdb.set_trace()
for i in sorted_inds:
bbox = boxes[i, :4]
score = boxes[i, -1]
if score < thresh:
continue
# show box (off by default)
ax.add_patch(
plt.Rectangle((bbox[0], bbox[1]),
bbox[2] - bbox[0],
bbox[3] - bbox[1],
fill=False, edgecolor='g',
linewidth=1, alpha=box_alpha))
if show_class:
# (x, y) = (bbox[0], bbox[3] + 2) if classes[i] == 1 else (bbox[3], bbox[1] - 2) # below for person or above for the rest
x, y = (bbox[0], bbox[1] - 2)
classes_i = np.argmax(classes[i])
# print(get_class_string(classes_i, score, dataset), classes_i, score)
ax.text(
x, y,
get_class_string(classes_i, score, dataset),
fontsize=4,
family='serif',
bbox=dict(
facecolor='g', alpha=0.4, pad=0, edgecolor='none'),
color='white')
# show mask
if masks is not None:
# ipdb.set_trace()
img = np.ones(im.shape)
color_mask = color_list[mask_color_id % len(color_list), 0:3]
mask_color_id += 1
w_ratio = .4
for c in range(3):
color_mask[c] = color_mask[c] * (1 - w_ratio) + w_ratio
for c in range(3):
img[:, :, c] = color_mask[c]
e_down = masks[i, :, :]
# Rescale mask
e_pil = Image.fromarray(e_down)
e_pil_up = e_pil.resize((H, W),Image.ANTIALIAS)
e = np.array(e_pil_up)
_, contour, hier = cv2.findContours(
e.copy(), cv2.RETR_CCOMP, cv2.CHAIN_APPROX_NONE)
for c in contour:
polygon = Polygon(
c.reshape((-1, 2)),
fill=True, facecolor=color_mask,
edgecolor='w', linewidth=1.2,
alpha=0.5)
ax.add_patch(polygon)
output_name = os.path.basename(im_name) + '.' + ext
fig.savefig(os.path.join(output_dir, '{}'.format(output_name)), dpi=dpi)
print('result saved to {}'.format(os.path.join(output_dir, '{}'.format(output_name))))
if show:
plt.show()
plt.close('all')
def main():
cfg_file = r'/home/twang/Documents/detectron/configs/12_2017_baselines/e2e_mask_rcnn_R-101-FPN_2x.yaml'
weights_file = r'/home/twang/Documents/detectron/model-weights/mask_rcnn_R-101-FPN_2x_model_final.pkl'
merge_cfg_from_file(cfg_file)
cfg.NUM_GPUS = 1
weights = cache_url(weights_file, cfg.DOWNLOAD_CACHE)
assert_and_infer_cfg()
model = infer_engine.initialize_model_from_cfg(weights)
dummy_coco_dataset = dummy_datasets.get_coco_dataset()
video_dir = '/media/network_shared_disk/WangTao/test_video/KLA_airport/Entrance_Peak_Hour.avi'
cap = cv2.VideoCapture(video_dir)
print(*'MJPG')
fourcc = cv2.VideoWriter_fourcc(*'MJPG')
video_output = cv2.VideoWriter("out.mp4", fourcc, 5, (1280, 720))
while cap.isOpened():
t1 = time.time()
ret, frame = cap.read()
if not ret:
break
frame = cv2.resize(frame, dsize=(1280, 720))
timers = defaultdict(Timer)
with c2_utils.NamedCudaScope(0):
cls_boxes, cls_segms, cls_keyps = infer_engine.im_detect_all(
model, frame, None, timers=timers)
thresh = 0.7
show_box = True
show_box = True
show_class = True
crop_person = True
dataset = dummy_coco_dataset
frame_for_person_crop = frame.copy()
frame_for_mask = frame.copy()
"""Constructs a numpy array with the detections visualized."""
if isinstance(cls_boxes, list):
boxes, segms, keypoints, classes = convert_from_cls_format(
cls_boxes, cls_segms, cls_keyps)
if boxes is None or boxes.shape[0] == 0 or max(boxes[:, 4]) < thresh:
return frame
if segms is not None and len(segms) > 0:
masks = mask_util.decode(segms)
color_list = colormap()
mask_color_id = 0
# Display in largest to smallest order to reduce occlusion
areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
sorted_inds = np.argsort(-areas)
for i in sorted_inds:
bbox = boxes[i, :4]
score = boxes[i, -1]
if score < thresh:
continue
# show class (person, backpack, handbag, suitcase)
class_default = ['person', 'backpack', 'handbag', 'suitcase']
if show_class:
class_str, class_text = get_class_string(
classes[i], score, dataset)
if class_text in class_default:
frame = vis_class(frame, (bbox[0], bbox[1] - 2), class_str)
#show bounding box
if show_box:
frame = vis_bbox(frame,
(bbox[0], bbox[1], bbox[2] - bbox[0],
bbox[3] - bbox[1]))
# show mask
if segms is not None and len(segms) > i:
color_mask = color_list[mask_color_id %
len(color_list), 0:3]
mask_color_id += 1
frame_for_mask = vis_mask(frame_for_mask,
masks[..., i], color_mask)
t2 = time.time()
durr = float(t2 - t1)
fps = 1.0 / durr
#cv2.putText(frame,"fps:%.3f"%fps,(20,20),4, 0.5, (0, 255, 0), 1, cv2.LINE_AA)
cv2.imshow('Detection using box', frame)
cv2.imshow('Detection using mask', frame_for_mask)
video_output.write(frame)
#video_mask.write(frame_for_mask)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cap.release()
video_output.release()
cv2.destroyAllWindows()
def vis_one_image(im,
im_name,
output_dir,
boxes,
segms=None,
keypoints=None,
thresh=0.9,
kp_thresh=2,
dpi=200,
box_alpha=0.8,
dataset=None,
show_class=False,
ext='pdf',
labels=None):
"""Visual debugging of detections."""
# if not os.path.exists ( output_dir ) :
# os.makedirs ( output_dir )
print("Processing image: {}".format(im_name))
if isinstance(boxes, list):
boxes, segms, keypoints, classes = convert_from_cls_format(
boxes, segms, keypoints)
if boxes is None or boxes.shape[0] == 0 or max(boxes[:, 4]) < thresh:
return
# 这里的mask是一个三维数组,1,2维分别代表原图的横纵坐标，3维一共有当前预测出的instanecs个数的大小。
if segms is not None:
masks = mask_util.decode(segms)
color_list = colormap(rgb=True) / 255
# print ( np.unique ( classes ) )
dataset_keypoints, _ = keypoint_utils.get_keypoints()
kp_lines = kp_connections(dataset_keypoints)
cmap = plt.get_cmap('rainbow')
colors = [cmap(i) for i in np.linspace(0, 1, len(kp_lines) + 2)]
fig = plt.figure(frameon=False)
fig.set_size_inches(im.shape[1] / dpi, im.shape[0] / dpi)
ax = plt.Axes(fig, [0., 0., 1., 1.])
ax.axis('off')
fig.add_axes(ax)
ax.imshow(im)
# Display in largest to smallest order to reduce occlusion
areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
# 默认从大到小排序，先画大的，小的会在后面把大的覆盖，这样大的和小的都能看到。
sorted_inds = np.argsort(-areas)
instance_cnt = defaultdict(int)
labels_graph = np.zeros((im.shape[0], im.shape[1]))
for item in enumerate(["instances_text", "instances", "labels"]):
if not os.path.exists(os.path.join(output_dir, item[1])):
os.makedirs(os.path.join(output_dir, item[1]), 0o777)
mask_file = open("{}/instances_text/{}.txt".format(output_dir, im_name),
"w")
# output labels prediction
for i in sorted_inds:
bbox = boxes[i, :4]
score = boxes[i, -1]
# 如果置信度大于0.5，执行
if score > 0.5:
single_mask = masks[:, :, i]
instances_graph = np.zeros((im.shape[0], im.shape[1]))
label_id = classes[i]
instance_cnt[dataset.classes[label_id]] += 1
instance_id = instance_cnt[dataset.classes[label_id]]
labels_graph[single_mask == 1] = label_id
instances_graph[single_mask == 1] = 255
scipy.misc.imsave(
'{}/instances/{}_{}.png'.format(output_dir, im_name,
label_id * 256 + instance_id),
instances_graph)
# 写入格式为mask_file_point class_name score
mask_file.write("{}/instances/{}_{}.png {} {}\n".format(
output_dir, im_name, label_id * 256 + instance_id, classes[i],
score))
# # current_graph存的是比较raw的值
# scipy.misc.imsave ( '/nfs/project/libo_i/mask-rcnn.pytorch/map_evaluation_format/raw/{}.jpg'.format ( im_name ) ,
# labels_graph )
# print ( np.unique ( labels_graph ) )
# colored_graph里存的是经过config中的cmap赋值之后的值。
# colored_graph = apply_color_map ( labels_graph , labels )
# gray_graph = rgb2gray ( colored_graph )
scipy.misc.imsave('{}/labels/{}.png'.format(output_dir, im_name),
labels_graph)
mask_file.close()
def vis_one_image(im,
im_name,
output_dir,
boxes=None,
segms=None,
keypoints=None,
score_info=None,
thresh=0.9,
kp_thresh=2,
dpi=200,
box_alpha=0.0,
dataset=None,
show_class=False,
ext='png'):
"""Visual debugging of detections."""
if not os.path.exists(output_dir):
os.makedirs(output_dir)
boxes, segms, keypoints, classes = convert_from_cls_format(
boxes, segms, keypoints)
if segms is not None:
masks = mask_util.decode(segms)
color_list = colormap(rgb=True) / 255
dataset_keypoints, _ = keypoint_utils.get_keypoints()
kp_lines = kp_connections(dataset_keypoints)
cmap = plt.get_cmap('rainbow')
colors = [cmap(i) for i in np.linspace(0, 1, len(kp_lines) + 2)]
fig = plt.figure(frameon=False)
fig.set_size_inches(im.shape[1] / dpi, im.shape[0] / dpi)
ax = plt.Axes(fig, [0., 0., 1., 1.])
ax.axis('off')
fig.add_axes(ax)
ax.imshow(im)
# Display in largest to smallest order to reduce occlusion
# areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
# sorted_inds = np.argsort(-areas)
score_list = [sl for slist in score_info for sl in slist]
# segms = [s for slist in cls_segms for s in slist]
mask_color_id = 0
for i in range(len(classes)):
score = score_list[i]
if score < thresh:
continue
print(dataset.classes[classes[i]], score)
# show box (off by default, box_alpha=0.0)
# ax.add_patch(
# plt.Rectangle((bbox[0], bbox[1]),
# bbox[2] - bbox[0],
# bbox[3] - bbox[1],
# fill=False, edgecolor='g',
# linewidth=0.5, alpha=box_alpha))
# if show_class:
# ax.text(
# bbox[0], bbox[1] - 2,
# get_class_string(classes[i], score, dataset),
# fontsize=3,
# family='serif',
# bbox=dict(
# facecolor='g', alpha=0.4, pad=0, edgecolor='none'),
# color='white')
# show mask
if segms is not None and len(segms) > i:
img = np.ones(im.shape)
color_mask = color_list[mask_color_id % len(color_list), 0:3]
mask_color_id += 1
w_ratio = .4
for c in range(3):
color_mask[c] = color_mask[c] * (1 - w_ratio) + w_ratio
for c in range(3):
img[:, :, c] = color_mask[c]
e = masks[:, :, i]
_, contour, hier = cv2.findContours(e.copy(), cv2.RETR_CCOMP,
cv2.CHAIN_APPROX_NONE)
for c in contour:
polygon = Polygon(c.reshape((-1, 2)),
fill=True,
facecolor=color_mask,
edgecolor='w',
linewidth=1.2,
alpha=0.5)
ax.add_patch(polygon)
output_name = os.path.basename(im_name) + '.' + ext
fig.savefig(os.path.join(output_dir, '{}'.format(output_name)),
dpi=dpi)
plt.close('all')
print("Saved pdfs in {}/{}".format(output_dir, im_name))
def vis_one_image(im,
im_name,
output_dir,
boxes,
segms=None,
keypoints=None,
thresh=0.9,
kp_thresh=2,
dpi=200,
box_alpha=0.0,
dataset=None,
show_class=False,
ext='pdf'):
# 图像从OpenCV格式转换成PIL格式
img_PIL = Image.fromarray(cv2.cvtColor(im, cv2.COLOR_BGR2RGB))
# 字体 字体*.ttc的存放路径一般是： /usr/share/fonts/opentype/noto/ 查找指令locate *.ttc
font = ImageFont.truetype('NotoSansCJK-Black.ttc', 40)
# 字体颜色
fillColor = (255, 0, 0)
# 文字输出位置
position = (100, 100)
# 输出内容
str = '在图片上输出中文'
draw = ImageDraw.Draw(img_PIL)
draw.text(position, str, font=font, fill=fillColor)
# 使用PIL中的save方法保存图片到本地
# img_PIL.save('02.jpg', 'jpeg')
# 转换回OpenCV格式
im = cv2.cvtColor(np.asarray(img_PIL), cv2.COLOR_RGB2BGR)
"""Visual debugging of detections."""
if not os.path.exists(output_dir):
os.makedirs(output_dir)
if isinstance(boxes, list):
boxes, segms, keypoints, classes = convert_from_cls_format( # -> n x 5, _, _, [c1, c2, ...]
boxes, segms, keypoints)
if boxes is None or boxes.shape[0] == 0 or max(boxes[:, 4]) < thresh:
return
if segms is not None:
masks = mask_util.decode(segms)
color_list = colormap(rgb=True) / 255
dataset_keypoints, _ = keypoint_utils.get_keypoints()
kp_lines = kp_connections(dataset_keypoints)
cmap = plt.get_cmap('rainbow')
colors = [cmap(i) for i in np.linspace(0, 1, len(kp_lines) + 2)]
fig = plt.figure(frameon=False)
fig.set_size_inches(im.shape[1] / dpi, im.shape[0] / dpi)
ax = plt.Axes(fig, [0., 0., 1., 1.])
ax.axis('off')
fig.add_axes(ax)
ax.imshow(im)
# Display in largest to smallest order to reduce occlusion
areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
sorted_inds = np.argsort(-areas)
mask_color_id = 0
for i in sorted_inds:
bbox = boxes[i, :4]
score = boxes[i, -1]
if score < thresh:
continue
print(dataset.classes[classes[i]], score)
# show box (off by default, box_alpha=0.0)
ax.add_patch(
plt.Rectangle((bbox[0], bbox[1]),
bbox[2] - bbox[0],
bbox[3] - bbox[1],
fill=False,
edgecolor='g',
linewidth=0.5,
alpha=box_alpha))
text = get_class_string(classes[i], score, dataset)
if show_class: # show class name
ax.text(bbox[0],
bbox[1] - 2,
text,
fontsize=3,
family='serif',
bbox=dict(facecolor='g',
alpha=0.4,
pad=0,
edgecolor='none'),
color='white')
# show mask
if segms is not None and len(segms) > i:
img = np.ones(im.shape)
color_mask = color_list[mask_color_id % len(color_list), 0:3]
mask_color_id += 1
w_ratio = .4
for c in range(3):
color_mask[c] = color_mask[c] * (1 - w_ratio) + w_ratio
for c in range(3):
img[:, :, c] = color_mask[c]
e = masks[:, :, i]
_, contour, hier = cv2.findContours(e.copy(), cv2.RETR_CCOMP,
cv2.CHAIN_APPROX_NONE)
for c in contour:
polygon = Polygon(c.reshape((-1, 2)),
fill=True,
facecolor=color_mask,
edgecolor='w',
linewidth=1.2,
alpha=0.5)
ax.add_patch(polygon)
# show keypoints
if keypoints is not None and len(keypoints) > i:
kps = keypoints[i]
plt.autoscale(False)
for l in range(len(kp_lines)):
i1 = kp_lines[l][0]
i2 = kp_lines[l][1]
if kps[2, i1] > kp_thresh and kps[2, i2] > kp_thresh:
x = [kps[0, i1], kps[0, i2]]
y = [kps[1, i1], kps[1, i2]]
line = ax.plot(x, y)
plt.setp(line, color=colors[l], linewidth=1.0, alpha=0.7)
if kps[2, i1] > kp_thresh:
ax.plot(kps[0, i1],
kps[1, i1],
'.',
color=colors[l],
markersize=3.0,
alpha=0.7)
if kps[2, i2] > kp_thresh:
ax.plot(kps[0, i2],
kps[1, i2],
'.',
color=colors[l],
markersize=3.0,
alpha=0.7)
# add mid shoulder / mid hip for better visualization
mid_shoulder = (
kps[:2, dataset_keypoints.index('right_shoulder')] +
kps[:2, dataset_keypoints.index('left_shoulder')]) / 2.0
sc_mid_shoulder = np.minimum(
kps[2, dataset_keypoints.index('right_shoulder')],
kps[2, dataset_keypoints.index('left_shoulder')])
mid_hip = (kps[:2, dataset_keypoints.index('right_hip')] +
kps[:2, dataset_keypoints.index('left_hip')]) / 2.0
sc_mid_hip = np.minimum(
kps[2, dataset_keypoints.index('right_hip')],
kps[2, dataset_keypoints.index('left_hip')])
if (sc_mid_shoulder > kp_thresh
and kps[2, dataset_keypoints.index('nose')] > kp_thresh):
x = [mid_shoulder[0], kps[0, dataset_keypoints.index('nose')]]
y = [mid_shoulder[1], kps[1, dataset_keypoints.index('nose')]]
line = ax.plot(x, y)
plt.setp(line,
color=colors[len(kp_lines)],
linewidth=1.0,
alpha=0.7)
if sc_mid_shoulder > kp_thresh and sc_mid_hip > kp_thresh:
x = [mid_shoulder[0], mid_hip[0]]
y = [mid_shoulder[1], mid_hip[1]]
line = ax.plot(x, y)
plt.setp(line,
color=colors[len(kp_lines) + 1],
linewidth=1.0,
alpha=0.7)
output_name = os.path.basename(im_name) + '.' + ext
fig.savefig(os.path.join(output_dir, '{}'.format(output_name)),
dpi=dpi)
plt.close('all')
def vis_one_image_eccv2018_car_3d(im,
im_name,
output_dir,
boxes,
car_cls_prob=None,
euler_angle=None,
trans_pred=None,
car_models=None,
intrinsic=None,
segms=None,
keypoints=None,
thresh=0.9,
dpi=200,
box_alpha=0.0,
dataset=None,
ext='jpg'):
"""Visual debugging of detections."""
if not os.path.exists(output_dir):
os.makedirs(output_dir)
if isinstance(boxes, list):
boxes, segms, keypoints, classes = convert_from_cls_format(
boxes, segms, keypoints)
car_cls = np.argmax(car_cls_prob, axis=1)
if boxes is None or boxes.shape[0] == 0 or max(boxes[:, 4]) < thresh:
return
if segms is not None:
masks = mask_util.decode(segms)
color_list = colormap(rgb=True) / 255
fig = plt.figure(frameon=False)
fig.set_size_inches(im.shape[1] / dpi, im.shape[0] / dpi)
ax = plt.Axes(fig, [0., 0., 1., 1.])
ax.axis('off')
fig.add_axes(ax)
ax.imshow(im)
# Display in largest to smallest order to reduce occlusion
areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
sorted_inds = np.argsort(-areas)
mask_color_id = 0
# [35, 38, 36, 39, 40, 34, 33]
for i in sorted_inds:
bbox = boxes[i, :4]
score = boxes[i, -1]
if score < thresh:
continue
json_id = dataset.contiguous_category_id_to_json_id[classes[i]]
class_string = dataset.id_map_to_cat[json_id]
# car_cls
car_cls_i = car_cls[i]
euler_angle_i = euler_angle[i]
trans_pred_i = trans_pred[i]
car_model_i = dataset.unique_car_models[car_cls_i]
car_model_name_i = dataset.car_id2name[car_model_i]
if class_string == 'car':
print("%s: %.4f, car model: %s" %
(class_string, score, car_model_name_i.name) +
'. quaternion: ' +
", ".join(['{:.3f}'.format(i)
for i in euler_angle_i]) + '. Trans: ' +
", ".join(['{:.3f}'.format(i) for i in trans_pred_i]))
else:
print(class_string, score)
# show box (off by default, box_alpha=0.0)
ax.add_patch(
plt.Rectangle((bbox[0], bbox[1]),
bbox[2] - bbox[0],
bbox[3] - bbox[1],
fill=False,
edgecolor='g',
linewidth=0.5,
alpha=box_alpha))
if class_string == 'car':
#vis_string = car_model_name_i.name + '. quaternion: ' + ", ".join(['{:.3f}'.format(i) for i in rot_pred_i]) + '. Trans: ' + ", ".join(['{:.3f}'.format(i) for i in trans_pred_i])
vis_string = car_model_name_i.name
ax.text(bbox[0],
bbox[1] - 2,
vis_string,
fontsize=10,
family='serif',
bbox=dict(facecolor='g',
alpha=0.4,
pad=0,
edgecolor='none'),
color='white')
# Show predicted pos mesh here:
if trans_pred_i is not None and euler_angle_i is not None:
#euler_angle = quaternion_to_euler_angle(rot_pred_i)
pose = np.concatenate((euler_angle_i, trans_pred_i))
car_name = car_model_name_i.name
car = car_models[car_name]
pose = np.array(pose)
# project 3D points to 2d image plane
rmat = euler_angles_to_rotation_matrix(pose[:3])
rvect, _ = cv2.Rodrigues(rmat)
imgpts, jac = cv2.projectPoints(np.float32(car['vertices']),
rvect,
pose[3:],
intrinsic,
distCoeffs=None)
triangles = np.array(car['faces'] - 1).astype('int64')
x = np.squeeze(imgpts[:, :, 1])
y = np.squeeze(imgpts[:, :, 0])
triangles = triangles
color_mask = color_list[mask_color_id % len(color_list), 0:3]
ax.triplot(y,
x,
triangles,
alpha=0.8,
linewidth=1.2,
color=color_mask)
# show mask
if segms is not None and len(segms) > i:
img = np.ones(im.shape)
color_mask = color_list[mask_color_id % len(color_list), 0:3]
mask_color_id += 1
w_ratio = .4
for c in range(3):
color_mask[c] = color_mask[c] * (1 - w_ratio) + w_ratio
for c in range(3):
img[:, :, c] = color_mask[c]
e = masks[:, :, i]
_, contour, hier = cv2.findContours(e.copy(), cv2.RETR_CCOMP,
cv2.CHAIN_APPROX_NONE)
for c in contour:
polygon = Polygon(c.reshape((-1, 2)),
fill=True,
facecolor=color_mask,
edgecolor='w',
linewidth=1.2,
alpha=0.2)
ax.add_patch(polygon)
output_name = os.path.basename(im_name) + '.' + ext
fig.savefig(os.path.join(output_dir, '{}'.format(output_name)), dpi=dpi)
plt.close('all')
def vis_one_image_opencv(im,
config,
boxes,
classes,
segms=None,
keypoints=None,
parsing=None,
uv=None,
dataset=None):
"""Constructs a numpy array with the detections visualized."""
timers = defaultdict(Timer)
timers['bbox_prproc'].tic()
global cfg
cfg = config
if boxes is None or boxes.shape[0] == 0 or max(boxes[:,
4]) < cfg.VIS.VIS_TH:
return im
if segms is not None and len(segms) > 0:
masks = mask_util.decode(segms)
# get color map
ins_colormap, parss_colormap = get_instance_parsing_colormap()
# Display in largest to smallest order to reduce occlusion
areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
sorted_inds = np.argsort(-areas)
timers['bbox_prproc'].toc()
instance_id = 1
for i in sorted_inds:
bbox = boxes[i, :4]
score = boxes[i, -1]
if score < cfg.VIS.VIS_TH:
continue
# get instance color (box, class_bg)
if cfg.VIS.SHOW_BOX.COLOR_SCHEME == 'category':
ins_color = ins_colormap[classes[i]]
elif cfg.VIS.SHOW_BOX.COLOR_SCHEME == 'instance':
instance_id = instance_id % len(ins_colormap.keys())
ins_color = ins_colormap[instance_id]
else:
ins_color = _GREEN
instance_id += 1
# show box (off by default)
if cfg.VIS.SHOW_BOX.ENABLED:
timers['show_box'].tic()
im = vis_bbox(
im, (bbox[0], bbox[1], bbox[2] - bbox[0], bbox[3] - bbox[1]),
ins_color)
timers['show_box'].toc()
# show class (off by default)
if cfg.VIS.SHOW_CLASS.ENABLED:
timers['show_class'].tic()
class_str = get_class_string(classes[i], score, dataset)
im = vis_class(im, (bbox[0], bbox[1] - 2), class_str, ins_color)
timers['show_class'].toc()
show_segms = True if cfg.VIS.SHOW_SEGMS.ENABLED and segms is not None and len(
segms) > i else False
show_kpts = True if cfg.VIS.SHOW_KPS.ENABLED and keypoints is not None and len(
keypoints) > i else False
show_parss = True if cfg.VIS.SHOW_PARSS.ENABLED and parsing is not None and len(
parsing) > i else False
show_uv = True if cfg.VIS.SHOW_UV.ENABLED and uv is not None and len(
uv) > i else False
# show mask
if show_segms:
timers['show_segms'].tic()
color_list = colormap_utils.colormap()
im = vis_mask(im, masks[..., i], ins_color, show_parss=show_parss)
timers['show_segms'].toc()
# show keypoints
if show_kpts:
timers['show_kpts'].tic()
im = vis_keypoints(im, keypoints[i], show_parss=show_parss)
timers['show_kpts'].toc()
# show parsing
if show_parss:
timers['show_parss'].tic()
im = vis_parsing(im,
parsing[i],
parss_colormap,
show_segms=show_segms)
timers['show_parss'].toc()
# show uv
if show_uv:
timers['show_uv'].tic()
im = vis_uv(im, uv[i], bbox)
timers['show_uv'].toc()
# for k, v in timers.items():
# print(' | {}: {:.3f}s'.format(k, v.total_time))
return im
def vis_one_image2(
im, im_name, output_dir, boxes, segms=None, keypoints=None, thresh=0.9,
kp_thresh=2, dpi=200, box_alpha=0.0, dataset=None, show_class=False,
ext='jpg'):
"""Visual debugging of detections."""
if not os.path.exists(output_dir):
os.makedirs(output_dir)
if isinstance(boxes, list):
boxes, segms, keypoints, classes = convert_from_cls_format(
boxes, segms, keypoints)
if boxes is None or boxes.shape[0] == 0 or max(boxes[:, 4]) < thresh:
return
dataset_keypoints, _ = keypoint_utils.get_keypoints()
if segms is not None:
masks = mask_util.decode(segms)
color_list = colormap(rgb=True) / 255
kp_lines = kp_connections(dataset_keypoints)
cmap = plt.get_cmap('rainbow')
colors = [cmap(i) for i in np.linspace(0, 1, len(kp_lines) + 2)]
fig = plt.figure(frameon=False)
fig.set_size_inches(im.shape[1] / dpi, im.shape[0] / dpi)
ax = plt.Axes(fig, [0., 0., 1., 1.])
ax.axis('off')
fig.add_axes(ax)
ax.imshow(im)
# Display in largest to smallest order to reduce occlusion
areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
sorted_inds = np.argsort(-areas)
output_name = os.path.basename(im_name) + '.txt'
file = open(os.path.join(output_dir, '{}'.format(output_name)), 'w')
class_list =[]
mask_color_id = 0
for i in sorted_inds:
bbox = boxes[i, :4]
score = boxes[i, -1]
if score < thresh:
continue
# show box (off by default)
ax.add_patch(
plt.Rectangle((bbox[0], bbox[1]),
bbox[2] - bbox[0],
bbox[3] - bbox[1],
fill=False, edgecolor='g',
linewidth=0.5, alpha=box_alpha))
if show_class:
ax.text(
bbox[0], bbox[1] - 2,
get_class_string(classes[i], score, dataset),
fontsize=3,
family='serif',
bbox=dict(
facecolor='g', alpha=0.4, pad=0, edgecolor='none'),
color='white')
# show mask
# show keypoints
# for kp in classes[i]:
# print (type(classes[i]))
class_list.append((get_class_string(classes[i], score, dataset)))
file.write(get_class_string(classes[i], score, dataset)+'\n')
file.close()
return class_list
def vis_one_image(
im, other_im, im_name, output_dir, boxes, gt_boxes, segms=None, keypoints=None, thresh=0.9,
kp_thresh=2, dpi=300, box_alpha=0.0, dataset=None, show_class=False,
ext='png', gt_classes=None):
"""Visual debugging of detections."""
if not os.path.exists(output_dir):
os.makedirs(output_dir)
if isinstance(boxes, list):
boxes, segms, keypoints, classes = convert_from_cls_format(
boxes, segms, keypoints)
if isinstance(gt_boxes[0], list):
box_list = [b for b in gt_boxes[0]]
if len(box_list) > 0:
gt_boxes[0] = np.vstack(box_list)
else:
gt_boxes[0] = None
if isinstance(gt_boxes[1], list):
box_list = [b for b in gt_boxes[1]]
if len(box_list) > 0:
gt_boxes[1] = np.vstack(box_list)
else:
gt_boxes[1] = None
if boxes is None or boxes.shape[0] == 0 or max(boxes[:, 4]) < thresh:
return
#if segms is not None:
# masks = mask_util.decode(segms)
color_list = colormap(rgb=True) / 255
dataset_keypoints, _ = keypoint_utils.get_keypoints()
kp_lines = kp_connections(dataset_keypoints)
cmap = plt.get_cmap('rainbow')
colors = [cmap(i) for i in np.linspace(0, 1, len(kp_lines) + 2)]
if cfg.DATA_SOURCE == 'mammo':
im = np.hstack([im,im,other_im])
else:
im = np.hstack([im,im])
fig = plt.figure(frameon=False)
fig.set_size_inches(im.shape[1] / dpi, im.shape[0] / dpi)
ax = plt.Axes(fig, [0., 0., 1., 1.])
ax.axis('off')
fig.add_axes(ax)
ax.imshow(im)
# Display in largest to smallest order to reduce occlusion
areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
sorted_inds = np.argsort(-areas)
if cfg.DATA_SOURCE == 'mammo':
fsize = 20
lwidth = 2
else:
fsize = 3
lwidth = 0.5
mask_color_id = 0
for i in sorted_inds:
bbox = boxes[i, :4]
score = boxes[i, -1]
if score < thresh:
continue
# show box (off by default, box_alpha=0.0)
ax.add_patch(
plt.Rectangle((bbox[0], bbox[1]),
bbox[2] - bbox[0],
bbox[3] - bbox[1],
fill=False, edgecolor='g',
linewidth=lwidth, alpha=box_alpha))
if True:
if show_class:
show_text = get_class_string(classes[i], score, dataset)
else:
show_text = '%.2f' % score
ax.text(
bbox[0], bbox[1] - 2,
show_text,
fontsize=fsize,
family='sans-serif',
bbox=dict(
facecolor='g', alpha=0.4, pad=0, edgecolor='none'),
color='white')
# Draw GTs
if gt_boxes[0] is not None:
for i in range(gt_boxes[0].shape[0]):
bbox = gt_boxes[0][i, :4]
# show box (off by default, box_alpha=0.0)
ax.add_patch(
plt.Rectangle((bbox[0], bbox[1]),
bbox[2] - bbox[0],
bbox[3] - bbox[1],
fill=False, edgecolor='r',
linewidth=lwidth, alpha=box_alpha))
if gt_classes is not None:
show_text = get_gt_class_string(gt_classes[0][i], dataset)
ax.text(
bbox[0], bbox[3] + 10,
show_text,
fontsize=fsize,
family='sans-serif',
bbox=dict(
facecolor='r', alpha=0.4, pad=0, edgecolor='none'),
color='white')
if gt_boxes[1] is not None:
for i in range(gt_boxes[1].shape[0]):
bbox = gt_boxes[1][i, :4]
# show box (off by default, box_alpha=0.0)
ax.add_patch(
plt.Rectangle((bbox[0], bbox[1]),
bbox[2] - bbox[0],
bbox[3] - bbox[1],
fill=False, edgecolor='y',
linewidth=lwidth, alpha=box_alpha))
if gt_classes is not None:
show_text = get_gt_class_string(gt_classes[1][i], dataset)
ax.text(
bbox[0], bbox[3] + 10,
show_text,
fontsize=fsize,
family='sans-serif',
bbox=dict(
facecolor='y', alpha=0.4, pad=0, edgecolor='none'),
color='white')
output_name = os.path.basename(im_name) + '.' + ext
fig.savefig(os.path.join(output_dir, '{}'.format(output_name)), dpi=dpi)
plt.close('all')
def vis_one_image_cvpr2018_wad(
im, im_name, output_dir, boxes, segms=None, keypoints=None, thresh=0.9,
kp_thresh=2, dpi=200, box_alpha=0.0, dataset=None, show_class=False,
ext='jpg'):
"""Visual debugging of detections."""
if not os.path.exists(output_dir):
os.makedirs(output_dir)
if isinstance(boxes, list):
boxes, segms, keypoints, classes = convert_from_cls_format(boxes, segms, keypoints)
if boxes is None or boxes.shape[0] == 0 or max(boxes[:, 4]) < thresh:
return
if segms is not None:
masks = mask_util.decode(segms)
color_list = colormap(rgb=True) / 255
dataset_keypoints, _ = keypoint_utils.get_keypoints()
kp_lines = kp_connections(dataset_keypoints)
cmap = plt.get_cmap('rainbow')
colors = [cmap(i) for i in np.linspace(0, 1, len(kp_lines) + 2)]
fig = plt.figure(frameon=False)
fig.set_size_inches(im.shape[1] / dpi, im.shape[0] / dpi)
ax = plt.Axes(fig, [0., 0., 1., 1.])
ax.axis('off')
fig.add_axes(ax)
ax.imshow(im)
# Display in largest to smallest order to reduce occlusion
areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
sorted_inds = np.argsort(-areas)
mask_color_id = 0
# [35, 38, 36, 39, 40, 34, 33]
for i in sorted_inds:
bbox = boxes[i, :4]
score = boxes[i, -1]
if score < thresh:
continue
json_id = dataset.contiguous_category_id_to_json_id[classes[i]]
class_string = dataset.id_map_to_cat[json_id]
print(class_string, score)
# show box (off by default, box_alpha=0.0)
ax.add_patch(plt.Rectangle((bbox[0], bbox[1]),
bbox[2] - bbox[0],
bbox[3] - bbox[1],
fill=False, edgecolor='g',
linewidth=0.5, alpha=box_alpha))
if show_class:
ax.text(bbox[0], bbox[1] - 2,
class_string + ' {:0.2f}'.format(score).lstrip('0'),
fontsize=10,
family='serif',
bbox=dict(facecolor='g', alpha=0.4, pad=0, edgecolor='none'),
color='white')
# show mask
if segms is not None and len(segms) > i:
img = np.ones(im.shape)
color_mask = color_list[mask_color_id % len(color_list), 0:3]
mask_color_id += 1
w_ratio = .4
for c in range(3):
color_mask[c] = color_mask[c] * (1 - w_ratio) + w_ratio
for c in range(3):
img[:, :, c] = color_mask[c]
e = masks[:, :, i]
_, contour, hier = cv2.findContours(
e.copy(), cv2.RETR_CCOMP, cv2.CHAIN_APPROX_NONE)
for c in contour:
polygon = Polygon(
c.reshape((-1, 2)),
fill=True, facecolor=color_mask,
edgecolor='w', linewidth=1.2,
alpha=0.5)
ax.add_patch(polygon)
output_name = os.path.basename(im_name) + '.' + ext
fig.savefig(os.path.join(output_dir, '{}'.format(output_name)), dpi=dpi)
plt.close('all')
def vis_one_image(im,
im_name,
output_dir,
boxes,
segms=None,
keypoints=None,
thresh=0.9,
kp_thresh=2,
dpi=200,
box_alpha=0.0,
dataset=None,
show_class=False,
ext='pdf',
gt_entry=None):
"""Visual debugging of detections."""
if not os.path.exists(output_dir):
os.makedirs(output_dir)
if isinstance(boxes, list):
boxes, segms, keypoints, classes = convert_from_cls_format(
boxes, segms, keypoints)
if boxes is None or boxes.shape[0] == 0 or max(boxes[:, 4]) < thresh:
return
if segms is not None and len(segms) > 0:
masks = mask_util.decode(segms)
color_list = colormap(rgb=True) / 255
dataset_keypoints, _ = keypoint_utils.get_keypoints()
kp_lines = kp_connections(dataset_keypoints)
cmap = plt.get_cmap('rainbow')
colors = [cmap(i) for i in np.linspace(0, 1, len(kp_lines) + 2)]
fig = plt.figure(frameon=False)
fig.set_size_inches(im.shape[1] / dpi, im.shape[0] / dpi)
ax = plt.Axes(fig, [0., 0., 1., 1.])
ax.axis('off')
fig.add_axes(ax)
ax.imshow(im)
# Display in largest to smallest order to reduce occlusion
areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
sorted_inds = np.argsort(-areas)
matches = []
wrong_classes = []
if gt_entry is not None:
gt_boxes = gt_entry['boxes']
gt_classes = gt_entry['gt_classes']
areas_gt = (gt_boxes[:, 2] - gt_boxes[:, 0]) * (gt_boxes[:, 3] -
gt_boxes[:, 1])
sorted_inds_gt = np.argsort(-areas_gt)
matches, wrong_classes, matches_gt = match_gt_dt(
boxes, sorted_inds, gt_boxes, sorted_inds_gt, classes, gt_classes,
thresh)
for i in sorted_inds_gt:
bbox = gt_boxes[i, :]
# only add ground-truth box if not matched
if matches_gt[i] == 0 and not cfg.VIS.ONLY_DETS:
ax.add_patch(
plt.Rectangle((bbox[0], bbox[1]),
bbox[2] - bbox[0],
bbox[3] - bbox[1],
fill=False,
edgecolor=cfg.VIS.GT_COLOR,
linewidth=cfg.VIS.BOX.LINEWIDTH,
alpha=box_alpha))
if show_class or cfg.VIS.GT_SHOW_CLASS:
ax.text(bbox[0] + 1,
bbox[1] - 6,
get_class_string(gt_classes[i], 1.0, dataset),
fontsize=cfg.VIS.LABEL.FONTSIZE,
family=cfg.VIS.LABEL.FAMILY,
weight=cfg.VIS.LABEL.WEIGHT,
bbox=dict(facecolor=cfg.VIS.LABEL.GT_TEXTCOLOR,
alpha=cfg.VIS.LABEL.ALPHA,
pad=cfg.VIS.LABEL.PAD,
edgecolor='none'),
color=cfg.VIS.LABEL.GT_TEXTCOLOR)
mask_color_id = 0
for i in sorted_inds:
bbox = boxes[i, :4]
score = boxes[i, -1]
if score < thresh:
continue
edge_color = 'b'
text_color = 'white'
if gt_entry is not None and not cfg.VIS.ONLY_DETS:
if matches[i] == -1:
edge_color = cfg.VIS.FP_COLOR
text_color = cfg.VIS.LABEL.FP_TEXTCOLOR
elif matches[i] == 0:
edge_color = cfg.VIS.FP_COLOR
text_color = cfg.VIS.LABEL.FP_TEXTCOLOR
elif matches[i] == 1:
edge_color = cfg.VIS.DT_COLOR
text_color = cfg.VIS.LABEL.DT_TEXTCOLOR
ax.add_patch(
plt.Rectangle((bbox[0], bbox[1]),
bbox[2] - bbox[0],
bbox[3] - bbox[1],
fill=False,
edgecolor=edge_color,
linewidth=cfg.VIS.BOX.LINEWIDTH,
alpha=cfg.VIS.BOX.ALPHA))
# do not show label of not matched detections
# if gt-boxes drawn: show_classes always for wrong (red) detections
if not cfg.VIS.ONLY_DETS and (gt_entry is not None and matches[i] == 0
and cfg.VIS.FP_SHOW_CLASS):
if cfg.VIS.FP_SHOW_CORRECT_CLASS:
ax.text(bbox[0] + 11,
bbox[1] - 6,
get_class_string(classes[i], score, dataset) +
'\n({})'.format(
get_class_string(wrong_classes[i], 1.0, dataset)),
fontsize=cfg.VIS.LABEL.FONTSIZE,
family=cfg.VIS.LABEL.FAMILY,
weight=cfg.VIS.LABEL.WEIGHT,
bbox=dict(facecolor=edge_color,
alpha=cfg.VIS.LABEL.ALPHA,
pad=cfg.VIS.LABEL.PAD,
edgecolor='none'),
color=text_color)
else:
ax.text(bbox[0] + 1,
bbox[1] - 6,
get_class_string(classes[i], score, dataset),
fontsize=cfg.VIS.LABEL.FONTSIZE,
family=cfg.VIS.LABEL.FAMILY,
weight=cfg.VIS.LABEL.WEIGHT,
bbox=dict(facecolor=edge_color,
alpha=cfg.VIS.LABEL.ALPHA,
pad=cfg.VIS.LABEL.PAD,
edgecolor='none'),
color=text_color)
elif not cfg.VIS.ONLY_DETS and (gt_entry is not None and matches[i]
== 1 and cfg.VIS.DT_SHOW_CLASS):
ax.text(bbox[0] + 1,
bbox[1] - 6,
get_class_string(classes[i], score, dataset),
fontsize=cfg.VIS.LABEL.FONTSIZE,
family=cfg.VIS.LABEL.FAMILY,
weight=cfg.VIS.LABEL.WEIGHT,
bbox=dict(facecolor=edge_color,
alpha=cfg.VIS.LABEL.ALPHA,
pad=cfg.VIS.LABEL.PAD,
edgecolor='none'),
color=text_color)
elif show_class and cfg.VIS.ONLY_DETS:
ax.text(bbox[0] + 1,
bbox[1] - 6,
get_class_string(classes[i], score, dataset),
fontsize=cfg.VIS.LABEL.FONTSIZE,
family=cfg.VIS.LABEL.FAMILY,
weight=cfg.VIS.LABEL.WEIGHT,
bbox=dict(facecolor=edge_color,
alpha=cfg.VIS.LABEL.ALPHA,
pad=cfg.VIS.LABEL.PAD,
edgecolor='none'),
color=text_color)
# show mask
if segms is not None and len(segms) > i:
img = np.ones(im.shape)
color_mask = color_list[mask_color_id % len(color_list), 0:3]
mask_color_id += 1
w_ratio = .4
for c in range(3):
color_mask[c] = color_mask[c] * (1 - w_ratio) + w_ratio
for c in range(3):
img[:, :, c] = color_mask[c]
e = masks[:, :, i]
_, contour, hier = cv2.findContours(e.copy(), cv2.RETR_CCOMP,
cv2.CHAIN_APPROX_NONE)
for c in contour:
polygon = Polygon(c.reshape((-1, 2)),
fill=True,
facecolor=color_mask,
edgecolor='w',
linewidth=1.2,
alpha=0.5)
ax.add_patch(polygon)
# show keypoints
if keypoints is not None and len(keypoints) > i:
kps = keypoints[i]
plt.autoscale(False)
for l in range(len(kp_lines)):
i1 = kp_lines[l][0]
i2 = kp_lines[l][1]
if kps[2, i1] > kp_thresh and kps[2, i2] > kp_thresh:
x = [kps[0, i1], kps[0, i2]]
y = [kps[1, i1], kps[1, i2]]
line = ax.plot(x, y)
plt.setp(line, color=colors[l], linewidth=1.0, alpha=0.7)
if kps[2, i1] > kp_thresh:
ax.plot(kps[0, i1],
kps[1, i1],
'.',
color=colors[l],
markersize=3.0,
alpha=0.7)
if kps[2, i2] > kp_thresh:
ax.plot(kps[0, i2],
kps[1, i2],
'.',
color=colors[l],
markersize=3.0,
alpha=0.7)
# add mid shoulder / mid hip for better visualization
mid_shoulder = (
kps[:2, dataset_keypoints.index('right_shoulder')] +
kps[:2, dataset_keypoints.index('left_shoulder')]) / 2.0
sc_mid_shoulder = np.minimum(
kps[2, dataset_keypoints.index('right_shoulder')],
kps[2, dataset_keypoints.index('left_shoulder')])
mid_hip = (kps[:2, dataset_keypoints.index('right_hip')] +
kps[:2, dataset_keypoints.index('left_hip')]) / 2.0
sc_mid_hip = np.minimum(
kps[2, dataset_keypoints.index('right_hip')],
kps[2, dataset_keypoints.index('left_hip')])
if (sc_mid_shoulder > kp_thresh
and kps[2, dataset_keypoints.index('nose')] > kp_thresh):
x = [mid_shoulder[0], kps[0, dataset_keypoints.index('nose')]]
y = [mid_shoulder[1], kps[1, dataset_keypoints.index('nose')]]
line = ax.plot(x, y)
plt.setp(line,
color=colors[len(kp_lines)],
linewidth=1.0,
alpha=0.7)
if sc_mid_shoulder > kp_thresh and sc_mid_hip > kp_thresh:
x = [mid_shoulder[0], mid_hip[0]]
y = [mid_shoulder[1], mid_hip[1]]
line = ax.plot(x, y)
plt.setp(line,
color=colors[len(kp_lines) + 1],
linewidth=1.0,
alpha=0.7)
output_name = os.path.basename(im_name) + '.' + ext
fig.savefig(os.path.join(output_dir, '{}'.format(output_name)),
dpi=dpi)
plt.close('all')
parser = argparse.ArgumentParser(
prog='Code to visualize the annotations in TFRecords.')
parser.add_argument('--tfrecord_dir',
type=str,
required=True,
help='Folder containing TFRecord files.')
parser.add_argument('--save_dir',
type=str,
required=True,
help='Folder where to save the visualizations.')
parser.add_argument('--labelmap_file',
type=str,
required=False,
help='Full path to the labelmap file.')
colorlist = colormap()
def plot_bboxes(image, bboxes, category_names):
for bbox, category_name in zip(bboxes, category_names):
image = cv2.rectangle(image, (bbox[1], bbox[0]), (bbox[3], bbox[2]),
(0, 255, 0), 1)
image = cv2.putText(image, category_name, (bbox[1], bbox[0]),
cv2.FONT_HERSHEY_TRIPLEX, 0.4, (255, 255, 255), 1,
cv2.LINE_AA)
return image
def plot_instances(image,
instances,
def extract_bbox(im,
im_name,
output_dir,
boxes,
segms=None,
keypoints=None,
thresh=0.9,
kp_thresh=2,
dpi=200,
box_alpha=0.0,
dataset=None,
show_class=False,
ext='pdf'):
"""Visual debugging of detections."""
if not os.path.exists(output_dir):
os.makedirs(output_dir)
if isinstance(boxes, list):
boxes, segms, keypoints, classes = convert_from_cls_format(
boxes, segms, keypoints)
if boxes is None or boxes.shape[0] == 0 or max(boxes[:, 4]) < thresh:
return
dataset_keypoints, _ = keypoint_utils.get_keypoints()
if segms is not None and len(segms) > 0:
masks = mask_util.decode(segms)
color_list = colormap(rgb=True) / 255
kp_lines = kp_connections(dataset_keypoints)
cmap = plt.get_cmap('rainbow')
colors = [cmap(i) for i in np.linspace(0, 1, len(kp_lines) + 2)]
"""
fig = plt.figure(frameon=False)
fig.set_size_inches(im.shape[1] / dpi, im.shape[0] / dpi)
ax = plt.Axes(fig, [0., 0., 1., 1.])
ax.axis('off')
fig.add_axes(ax)
ax.imshow(im)
"""
# Display in largest to smallest order to reduce occlusion
areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
sorted_inds = np.argsort(-areas)
results = {}
det_results = []
mask_color_id = 0
for i in sorted_inds:
bbox = boxes[i, :4]
score = boxes[i, -1]
if score < thresh:
continue
result = {}
result['location'] = bbox.tolist()
result['label'] = classes[i]
result['score'] = score.tolist()
det_results.append(result)
# show box (off by default)
"""
ax.add_patch(
plt.Rectangle((bbox[0], bbox[1]),
bbox[2] - bbox[0],
bbox[3] - bbox[1],
fill=False, edgecolor='g',
linewidth=0.5, alpha=box_alpha))
if show_class:
ax.text(
bbox[0], bbox[1] - 2,
get_class_string(classes[i], score, dataset),
fontsize=3,
family='serif',
bbox=dict(
facecolor='g', alpha=0.4, pad=0, edgecolor='none'),
color='white')
"""
print('# of bbox for img {} is {}'.format(im_name, len(det_results)))
results['detection'] = det_results
img_id = im_name.split('/')[-1][:-4]
results['image_id'] = img_id
# output_name = os.path.basename(im_name) + '.' + ext
# fig.savefig(os.path.join(output_dir, '{}'.format(output_name)), dpi=dpi)
# plt.close('all')
return results
def visualize_segmentations(img_path, pkl_path, dataset, images_dir):
# Open image and pkl
im = imread(img_path)
cls_boxes, cls_segms, cls_keyps = pickle.load(open(pkl_path, 'rb'))
boxes, segms, keypoints, classes = convert_from_cls_format(
cls_boxes, cls_segms, cls_keyps)
segmentations = []
if boxes is None:
return segmentations
masks = mask_util.decode(segms)
color_list = colormap()
mask_color_id = 0
scores = boxes[:, -1]
sorted_inds = np.argsort(scores)
for i in sorted_inds[::-1]:
bbox = boxes[i, :4]
score = boxes[i, -1]
cls = classes[i]
class_name = dataset.classes[cls]
class_str = class_name + ' {:0.2f}'.format(score).lstrip('0')
vis_im = im.copy()
# show mask
if segms is not None and len(segms) > i:
color_mask = color_list[mask_color_id % len(color_list), 0:3]
mask_color_id += 1
vis_im = vis_mask(vis_im, masks[..., i], color_mask, alpha=1.)
# Crop with context
x0, y0, x1, y1 = bbox
bbox_h = y1 - y0
bbox_w = x1 - x0
vis_im = vis_bbox(vis_im, (x0, y0, bbox_w, bbox_h)) # show bbox
h, w = vis_im.shape[:2]
y0_c = max(0, int(y0 - .5 * bbox_h))
y1_c = min(h, int(y1 + .5 * bbox_h))
x0_c = max(0, int(x0 - .5 * bbox_w))
x1_c = min(w, int(x1 + .5 * bbox_w))
vis_im = vis_im[y0_c:y1_c, x0_c:x1_c, :]
im_crop = im[y0_c:y1_c, x0_c:x1_c, :]
vis_im = np.concatenate((im_crop, vis_im), axis=1)
# resize
r = 256. / vis_im.shape[0]
vis_im = cv2.resize(vis_im,
dsize=(0, 0),
fx=r,
fy=r,
interpolation=cv2.INTER_NEAREST)
# show class
vis_im = vis_class(vis_im, (5, 15), class_str)
vis_path = save(vis_im, images_dir)
segmentations.append([vis_path, class_name, score])
return segmentations
def vis_one_image(im,
im_name,
output_dir,
boxes,
segms=None,
keypoints=None,
thresh=0.9,
kp_thresh=2,
dpi=200,
box_alpha=0.0,
dataset=None,
show_class=False,
ext='pdf'):
"""Visual debugging of detections."""
if output_dir is not None:
if not os.path.exists(output_dir):
os.makedirs(output_dir)
if isinstance(boxes, list):
boxes, segms, keypoints, classes = convert_from_cls_format(
boxes, segms, keypoints)
if boxes is None or boxes.shape[0] == 0 or max(boxes[:, 4]) < thresh:
return
if segms is not None:
masks = mask_util.decode(segms)
color_list = colormap(rgb=True) / 255
dataset_keypoints, _ = keypoint_utils.get_keypoints()
kp_lines = kp_connections(dataset_keypoints)
cmap = plt.get_cmap('rainbow')
colors = [cmap(i) for i in np.linspace(0, 1, len(kp_lines) + 2)]
fig = plt.figure(frameon=False)
if output_dir is not None:
fig.set_size_inches(im.shape[1] / dpi, im.shape[0] / dpi)
ax = plt.Axes(fig, [0., 0., 1., 1.])
ax.axis('off')
fig.add_axes(ax)
ax.imshow(im)
# preprocess the boxes
if thresh < 0:
# When VIS_TH less than zero, it means take the highest -thresh score boxes
sorted_inds = np.argsort(-boxes[:, -1])
boxes = boxes[sorted_inds[:-int(thresh)]]
classes = [classes[_] for _ in sorted_inds[:-int(thresh)]]
# Display in largest to smallest order to reduce occlusion
areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
sorted_inds = np.argsort(-areas)
_boxes = []
texts = []
for i in sorted_inds:
bbox = boxes[i, :4]
score = boxes[i, -1]
if score < thresh:
continue
if len(_boxes) > 0 and (bbox == _boxes[-1][:4]).all():
# Same box, merge prediction
texts[-1] += '/' + get_class_string(classes[i], score, dataset)
else:
_boxes.append(boxes[i])
texts.append(get_class_string(classes[i], score, dataset))
boxes = np.stack(_boxes)
mask_color_id = 0
for i in range(len(boxes)):
bbox = boxes[i, :4]
score = boxes[i, -1]
if score < thresh:
continue
# print(dataset.classes[classes[i]], score)
print(texts[i])
# show box (off by default, box_alpha=0.0)
ax.add_patch(
plt.Rectangle(
(bbox[0], bbox[1]),
bbox[2] - bbox[0],
bbox[3] - bbox[1],
fill=False,
edgecolor='r', #'##66FF66' if '@'in texts[i] else '#0099FF' ,
linewidth=5,
alpha=box_alpha))
if show_class:
ax.text(
bbox[0],
bbox[1] + 6,
texts[i].split(' ')
[0], #get_class_string(classes[i], score, dataset),
fontsize=3,
family='serif',
bbox=dict(
facecolor='#66FF66' if '@' in texts[i] else '#0099FF',
alpha=0.4,
pad=0,
edgecolor='none'),
color='white')
# show mask
if segms is not None and len(segms) > i:
img = np.ones(im.shape)
color_mask = color_list[mask_color_id % len(color_list), 0:3]
mask_color_id += 1
w_ratio = .4
for c in range(3):
color_mask[c] = color_mask[c] * (1 - w_ratio) + w_ratio
for c in range(3):
img[:, :, c] = color_mask[c]
e = masks[:, :, i]
_, contour, hier = cv2.findContours(e.copy(), cv2.RETR_CCOMP,
cv2.CHAIN_APPROX_NONE)
for c in contour:
polygon = Polygon(c.reshape((-1, 2)),
fill=True,
facecolor=color_mask,
edgecolor='w',
linewidth=1.2,
alpha=0.5)
ax.add_patch(polygon)
# show keypoints
if keypoints is not None and len(keypoints) > i:
kps = keypoints[i]
plt.autoscale(False)
for l in range(len(kp_lines)):
i1 = kp_lines[l][0]
i2 = kp_lines[l][1]
if kps[2, i1] > kp_thresh and kps[2, i2] > kp_thresh:
x = [kps[0, i1], kps[0, i2]]
y = [kps[1, i1], kps[1, i2]]
line = ax.plot(x, y)
plt.setp(line, color=colors[l], linewidth=1.0, alpha=0.7)
if kps[2, i1] > kp_thresh:
ax.plot(kps[0, i1],
kps[1, i1],
'.',
color=colors[l],
markersize=3.0,
alpha=0.7)
if kps[2, i2] > kp_thresh:
ax.plot(kps[0, i2],
kps[1, i2],
'.',
color=colors[l],
markersize=3.0,
alpha=0.7)
# add mid shoulder / mid hip for better visualization
mid_shoulder = (
kps[:2, dataset_keypoints.index('right_shoulder')] +
kps[:2, dataset_keypoints.index('left_shoulder')]) / 2.0
sc_mid_shoulder = np.minimum(
kps[2, dataset_keypoints.index('right_shoulder')],
kps[2, dataset_keypoints.index('left_shoulder')])
mid_hip = (kps[:2, dataset_keypoints.index('right_hip')] +
kps[:2, dataset_keypoints.index('left_hip')]) / 2.0
sc_mid_hip = np.minimum(
kps[2, dataset_keypoints.index('right_hip')],
kps[2, dataset_keypoints.index('left_hip')])
if (sc_mid_shoulder > kp_thresh
and kps[2, dataset_keypoints.index('nose')] > kp_thresh):
x = [mid_shoulder[0], kps[0, dataset_keypoints.index('nose')]]
y = [mid_shoulder[1], kps[1, dataset_keypoints.index('nose')]]
line = ax.plot(x, y)
plt.setp(line,
color=colors[len(kp_lines)],
linewidth=1.0,
alpha=0.7)
if sc_mid_shoulder > kp_thresh and sc_mid_hip > kp_thresh:
x = [mid_shoulder[0], mid_hip[0]]
y = [mid_shoulder[1], mid_hip[1]]
line = ax.plot(x, y)
plt.setp(line,
color=colors[len(kp_lines) + 1],
linewidth=1.0,
alpha=0.7)
if output_dir is not None:
output_name = os.path.basename(im_name) + '.' + ext
fig.savefig(os.path.join(output_dir, '{}'.format(output_name)),
dpi=dpi)
plt.close('all')
else:
plt.plot()
def draw_all(im,
all_boxes=None,
all_classes=None,
box_color=None,
do_draw_box=True,
all_segs=None,
all_seg_boxes=None,
seg_color=None,
all_masks=None,
all_mask_boxes=None,
mask_color=None,
all_kps=None,
skeleton=None,
point_color=None,
skeleton_color=None,
kps_thresh=0,
kps_show_num=False,
all_densepose_iuv=None,
all_densepose_iuv_point=None,
all_track_ids=None,
track_id_color=None):
# im: (h, w, 3)
# all_boxes: (num_boxes, 4)
# all_kps: (num_boxes, num_kps*3)
# all_masks: (num_boxes, h, w)
# all_segs: (num_boxes, h, w)
from utils.colormap import colormap
color_list = colormap()
im_draw = copy.deepcopy(im)
# draw boxes
if all_boxes is not None and do_draw_box:
all_boxes = np.round(all_boxes).astype(int)
for i in range(len(all_boxes)):
if all_classes is not None:
color_idx = int(abs(all_classes[i])) - 1
else:
color_idx = i
if box_color is None:
box_color_i = color_list[color_idx % len(color_list)]
elif isinstance(box_color, list):
box_color_i = box_color[color_idx % len(box_color)]
else:
box_color_i = box_color
im_draw = draw_box(im_draw, all_boxes[i], box_color=box_color_i)
# draw segs
if all_segs is not None:
for i in range(len(all_segs)):
if seg_color is None:
seg_color_i = color_list
else:
seg_color_i = seg_color
seg_box_i = all_seg_boxes[i].astype(
int) if all_seg_boxes is not None else None
im_draw = draw_seg(im_draw,
all_segs[i],
seg_box_i,
seg_color=seg_color_i)
# draw masks
if all_masks is not None:
for i in range(len(all_masks)):
if mask_color is None:
mask_color_i = color_list[i % len(color_list)]
elif isinstance(mask_color, list):
mask_color_i = mask_color[i % len(mask_color)]
else:
mask_color_i = mask_color
mask_box_i = all_mask_boxes[i].astype(
int) if all_mask_boxes is not None else None
im_draw = draw_mask(im_draw,
all_masks[i],
mask_box_i,
mask_color=mask_color_i)
# draw kps
if all_kps is not None:
for i in range(len(all_kps)):
if point_color is None:
point_color_i = color_list[i % len(color_list)]
elif isinstance(point_color, list):
point_color_i = point_color[i % len(point_color)]
else:
point_color_i = point_color
if skeleton_color is None:
skeleton_color_i = color_list
else:
skeleton_color_i = skeleton_color
im_draw = draw_kps(im_draw,
all_kps[i],
skeleton=skeleton,
point_color=point_color_i,
skeleton_color=skeleton_color_i,
kps_thresh=kps_thresh,
kps_show_num=kps_show_num)
# draw densepose IUV
if all_densepose_iuv_point is not None:
if len(all_densepose_iuv_point[1]) > 0:
im_draw = draw_densepose_iuv_point(im_draw,
all_densepose_iuv_point)
if all_densepose_iuv is not None:
for i in range(len(all_densepose_iuv)):
seg_densepose_box_i = all_seg_boxes[i].astype(
int) if all_seg_boxes is not None else None
im_draw = draw_densepose_iuv(im_draw, all_densepose_iuv[i],
seg_densepose_box_i)
# draw track ids
if all_track_ids is not None:
assert all_boxes is not None
for i in range(len(all_track_ids)):
if track_id_color is None:
track_id_color_i = color_list[i % len(color_list)]
elif isinstance(track_id_color, list):
track_id_color_i = track_id_color[i]
else:
track_id_color_i = track_id_color
x = all_boxes[i, 0] + (all_boxes[i, 2] - all_boxes[i, 0]) / 3
im_draw = draw_track_id(im_draw,
all_track_ids[i],
pos=(x, all_boxes[i, 1]),
track_id_color=track_id_color_i)
return im_draw
def vis_one_image_orig(im,
im_name,
output_dir,
boxes,
segms=None,
keypoints=None,
thresh=0.9,
kp_thresh=2,
dpi=200,
box_alpha=0.0,
dataset=None,
show_class=False,
ext='pdf',
entry=None):
"""Visual debugging of detections."""
if not os.path.exists(output_dir):
os.makedirs(output_dir)
# pdb.set_trace()
if isinstance(boxes, list):
boxes, segms, keypoints, classes = convert_from_cls_format( # -> n x 5, _, _, [c1, c2, ...]
boxes, segms, keypoints)
if boxes is None or boxes.shape[0] == 0 or max(boxes[:, 4]) < thresh:
return
if segms is not None:
masks = mask_util.decode(segms)
color_list = colormap(rgb=True) / 255
dataset_keypoints, _ = keypoint_utils.get_keypoints()
kp_lines = kp_connections(dataset_keypoints)
cmap = plt.get_cmap('rainbow')
colors = [cmap(i) for i in np.linspace(0, 1, len(kp_lines) + 2)]
fig = plt.figure(frameon=False)
fig.set_size_inches(im.shape[1] / dpi, im.shape[0] / dpi)
ax = plt.Axes(fig, [0., 0., 1., 1.])
ax.axis('off')
fig.add_axes(ax)
ax.imshow(im)
#LJ
imdir = '/data1/liujingyu/DR/fold_all'
img_path = os.path.join(imdir, im_name.split('_')[1] + '.png')
# pdb.set_trace()
assert os.path.exists(img_path)
img_ori = cv2.imread(img_path, 1)
# Display in largest to smallest order to reduce occlusion
areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
sorted_inds = np.argsort(-areas)
mask_color_id = 0
Flag = False
for i in sorted_inds:
bbox = boxes[i, :4]
score = boxes[i, -1]
if score < thresh:
continue
# LJ 注释
# print(dataset.eng_classes[classes[i]], score)
# show box (off by default, box_alpha=0.0)
# LJ 不显示 box
if False:
ax.add_patch(
plt.Rectangle((bbox[0], bbox[1]),
bbox[2] - bbox[0],
bbox[3] - bbox[1],
fill=False,
edgecolor='g',
linewidth=0.5,
alpha=box_alpha))
if False:
# if show_class:
ax.text(
bbox[0],
bbox[1] - 2,
# LJ 注释
get_class_eng_string(classes[i], score, dataset),
fontsize=12,
family='serif',
bbox=dict(facecolor='g', alpha=0.4, pad=0, edgecolor='none'),
# LJ
#bbox = None,
color='white')
# show mask
# pdb.set_trace()
if segms is not None and len(segms) > i:
img = np.ones(im.shape)
color_mask = color_list[mask_color_id % len(color_list), 0:3]
mask_color_id += 1
w_ratio = .4
for c in range(3):
color_mask[c] = color_mask[c] * (1 - w_ratio) + w_ratio
for c in range(3):
img[:, :, c] = color_mask[c]
e = masks[:, :, i]
# LJ opencv 版本的原因 返回三值
bin, contour, hier = cv2.findContours(e.copy(), cv2.RETR_CCOMP,
cv2.CHAIN_APPROX_NONE)
# 生成512图
color_edge = [0, 0, 0]
for cm in range(3):
color_edge[cm] = int(color_mask[cm] * 255)
# pdb.set_trace()
point0 = (int(bbox[0]), int(bbox[1]))
text = dataset.eng_classes[classes[i]]
if classes[i] in entry['cls_list']:
Flag = True
cv2.putText(img_ori, text, point0, cv2.FONT_HERSHEY_SIMPLEX, 2,
(250, 250, 250), 3)
img_vis = cv2.drawContours(
img_ori, contour, -1,
(color_edge[2], color_edge[1], color_edge[0]), 8)
# 生成半透明mask
# for c in contour:
# polygon = Polygon(
# c.reshape((-1, 2)),
# fill=True, facecolor=color_mask,
# edgecolor='w', linewidth=1.2,
# alpha=0.5)
# ax.add_patch(polygon)
# show keypoints
if keypoints is not None and len(keypoints) > i:
kps = keypoints[i]
plt.autoscale(False)
for l in range(len(kp_lines)):
i1 = kp_lines[l][0]
i2 = kp_lines[l][1]
if kps[2, i1] > kp_thresh and kps[2, i2] > kp_thresh:
x = [kps[0, i1], kps[0, i2]]
y = [kps[1, i1], kps[1, i2]]
line = ax.plot(x, y)
plt.setp(line, color=colors[l], linewidth=1.0, alpha=0.7)
if kps[2, i1] > kp_thresh:
ax.plot(kps[0, i1],
kps[1, i1],
'.',
color=colors[l],
markersize=3.0,
alpha=0.7)
if kps[2, i2] > kp_thresh:
ax.plot(kps[0, i2],
kps[1, i2],
'.',
color=colors[l],
markersize=3.0,
alpha=0.7)
# add mid shoulder / mid hip for better visualization
mid_shoulder = (
kps[:2, dataset_keypoints.index('right_shoulder')] +
kps[:2, dataset_keypoints.index('left_shoulder')]) / 2.0
sc_mid_shoulder = np.minimum(
kps[2, dataset_keypoints.index('right_shoulder')],
kps[2, dataset_keypoints.index('left_shoulder')])
mid_hip = (kps[:2, dataset_keypoints.index('right_hip')] +
kps[:2, dataset_keypoints.index('left_hip')]) / 2.0
sc_mid_hip = np.minimum(
kps[2, dataset_keypoints.index('right_hip')],
kps[2, dataset_keypoints.index('left_hip')])
if (sc_mid_shoulder > kp_thresh
and kps[2, dataset_keypoints.index('nose')] > kp_thresh):
x = [mid_shoulder[0], kps[0, dataset_keypoints.index('nose')]]
y = [mid_shoulder[1], kps[1, dataset_keypoints.index('nose')]]
line = ax.plot(x, y)
plt.setp(line,
color=colors[len(kp_lines)],
linewidth=1.0,
alpha=0.7)
if sc_mid_shoulder > kp_thresh and sc_mid_hip > kp_thresh:
x = [mid_shoulder[0], mid_hip[0]]
y = [mid_shoulder[1], mid_hip[1]]
line = ax.plot(x, y)
plt.setp(line,
color=colors[len(kp_lines) + 1],
linewidth=1.0,
alpha=0.7)
# pdb.set_trace()
# output_name = os.path.basename(im_name) + '.' + ext
# fig.savefig(os.path.join(output_dir, '{}'.format(output_name)), dpi=dpi)
# plt.close('all')
# LJ
if Flag:
output_name = os.path.basename(im_name) + '.jpg'
img_vis = cv2.resize(img_vis, (512, 512),
interpolation=cv2.INTER_CUBIC)
cv2.imwrite(os.path.join(output_dir, '{}'.format(output_name)),
img_vis)
def vis_one_image(im,
im_name,
boxes,
segms=None,
keypoints=None,
thresh=0.9,
kp_thresh=2,
dpi=200,
box_alpha=0.0,
dataset=None,
show_class=False,
ext=None):
if isinstance(boxes, list):
boxes, segms, keypoints, classes = convert_from_cls_format(
boxes, segms, keypoints)
if boxes is None or boxes.shape[0] == 0 or max(boxes[:, 4]) < thresh:
return
dataset_keypoints, _ = keypoint_utils.get_keypoints()
if segms is not None and len(segms) > 0:
masks = mask_util.decode(segms)
color_list = colormap(rgb=True) / 255
kp_lines = kp_connections(dataset_keypoints)
cmap = plt.get_cmap('rainbow')
colors = [cmap(i) for i in np.linspace(0, 1, len(kp_lines) + 2)]
fig = plt.figure(frameon=False)
fig.set_size_inches(im.shape[1] / dpi, im.shape[0] / dpi)
ax = plt.Axes(fig, [0., 0., 1., 1.])
ax.axis('off')
fig.add_axes(ax)
ax.imshow(im)
areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
print("areas%%%%%%%%%%%%%%%%%")
print(areas)
sorted_inds = np.argsort(-areas)
print("sorted_inds&&&&&&&&&&&&&&&&&&")
print(sorted_inds)
mask_color_id = 0
for i in sorted_inds:
bbox = boxes[i, :4]
score = boxes[i, -1]
print("zhege IIIIIIIIIII")
print(i)
if score < thresh:
continue
ax.add_patch(
plt.Rectangle((bbox[0], bbox[1]),
bbox[2] - bbox[0],
bbox[3] - bbox[1],
fill=False,
edgecolor='g',
linewidth=0.5,
alpha=box_alpha))
if show_class:
ax.text(bbox[0],
bbox[1] - 2,
get_class_string(classes[i], score, dataset),
fontsize=3,
family='serif',
bbox=dict(facecolor='g',
alpha=0.4,
pad=0,
edgecolor='none'),
color='white')
if segms is not None and len(segms) > i:
print("duoshaogeIIIIIIIIIIIIIIIIIIIIIIIIIII")
print(i)
img = np.ones(im.shape)
color_mask = color_list[mask_color_id % len(color_list), 0:3]
mask_color_id += 1
w_ratio = .4
for c in range(3):
color_mask[c] = color_mask[c] * (1 - w_ratio) + w_ratio
for c in range(3):
img[:, :, c] = color_mask[c]
e = masks[:, :, i]
_, contour, hier = cv2.findContours(e.copy(), cv2.RETR_CCOMP,
cv2.CHAIN_APPROX_NONE)
for c in contour:
polygon = Polygon(c.reshape((-1, 2)),
fill=True,
facecolor=color_mask,
edgecolor='w',
linewidth=0.2,
alpha=0.5)
ax.add_patch(polygon)
woca = c.reshape((-1, 2))
arr = np.array(woca)
key = np.unique(woca)
result = {}
k_qu = []
arr_end = []
arr_end2 = []
arr_endlast = []
A = []
B = []
Bplus = []
jjj = []
Bwanmei = []
Bcao = []
Bao = []
for k in key:
mask = (arr == k)
arr_new = arr[mask]
v = arr_new.size
result[k] = v
x = np.argwhere(arr == k)
if v > 1:
x = np.argwhere(arr == k)
x = np.array(x)
x0 = arr[:, 0]
y0 = arr[:, 1]
y0lie = []
for i in range(0, len(x0)):
if x0[i] != k:
pass
if x0[i] == k:
y0lie.append(y0[i])
y0lienew = []
arr_new = []
arr_new_2 = []
if y0lie == []:
pass
else:
miny0 = np.min(y0lie)
maxy0 = np.max(y0lie)
for i in range(miny0, maxy0 + 1):
y0lienew.append(i)
y0liezuizhong = []
if y0lienew == []:
pass
else:
miny0lienew = np.min(y0lienew)
maxy0lienew = np.max(y0lienew)
for i in range(miny0lienew, maxy0lienew + 1):
y0liezuizhong.append(i)
for i in range(0, len(y0liezuizhong)):
arr_temp = [k, y0liezuizhong[i]]
arr_new.append(arr_temp)
arr_end.append(arr_new)
x0lie = []
for i in range(0, len(y0)):
if y0[i] != k:
pass
if y0[i] == k:
x0lie.append(x0[i])
x0lienew = []
arr_new2 = []
if x0lie == []:
pass
else:
minx0 = np.min(x0lie)
maxx0 = np.max(x0lie)
for i in range(minx0, maxx0 + 1):
x0lienew.append(i)
x0liezuizhong = []
if x0lienew == []:
pass
else:
minx0lienew = np.min(x0lienew)
maxx0lienew = np.max(x0lienew)
for i in range(minx0lienew, maxx0lienew + 1):
x0liezuizhong.append(i)
for i in range(0, len(x0liezuizhong)):
arr_temp = [x0liezuizhong[i], k]
arr_new2.append(arr_temp)
arr_end2.append(arr_new2)
arr_endlast = arr_end + arr_end2
A = list(chain(*arr_endlast))
B = np.array(list(set([tuple(t) for t in A])))
if len(B) > 4:
Bplus = random.sample(B, 2)
if len(B) < 4:
jjj = arr_endlast
Bwanmei = Bplus + jjj
Bcaotmp = list(chain(*Bwanmei))
Bcao = np.array(Bcaotmp)
Bao = Bcao.reshape(-1, 2)
print("#####################################")
print(Bao)
if keypoints is not None and len(keypoints) > i:
kps = keypoints[i]
plt.autoscale(False)
for l in range(len(kp_lines)):
i1 = kp_lines[l][0]
i2 = kp_lines[l][1]
if kps[2, i1] > kp_thresh and kps[2, i2] > kp_thresh:
x = [kps[0, i1], kps[0, i2]]
y = [kps[1, i1], kps[1, i2]]
line = plt.plot(x, y)
plt.setp(line, color=colors[l], linewidth=1.0, alpha=0.7)
if kps[2, i1] > kp_thresh:
plt.plot(kps[0, i1],
kps[1, i1],
'.',
color=colors[l],
markersize=3.0,
alpha=0.7)
if kps[2, i2] > kp_thresh:
plt.plot(kps[0, i2],
kps[1, i2],
'.',
color=colors[l],
markersize=3.0,
alpha=0.7)
mid_shoulder = (
kps[:2, dataset_keypoints.index('right_shoulder')] +
kps[:2, dataset_keypoints.index('left_shoulder')]) / 2.0
sc_mid_shoulder = np.minimum(
kps[2, dataset_keypoints.index('right_shoulder')],
kps[2, dataset_keypoints.index('left_shoulder')])
mid_hip = (kps[:2, dataset_keypoints.index('right_hip')] +
kps[:2, dataset_keypoints.index('left_hip')]) / 2.0
sc_mid_hip = np.minimum(
kps[2, dataset_keypoints.index('right_hip')],
kps[2, dataset_keypoints.index('left_hip')])
if (sc_mid_shoulder > kp_thresh
and kps[2, dataset_keypoints.index('nose')] > kp_thresh):
x = [mid_shoulder[0], kps[0, dataset_keypoints.index('nose')]]
y = [mid_shoulder[1], kps[1, dataset_keypoints.index('nose')]]
line = plt.plot(x, y)
plt.setp(line,
color=colors[len(kp_lines)],
linewidth=1.0,
alpha=0.7)
if sc_mid_shoulder > kp_thresh and sc_mid_hip > kp_thresh:
x = [mid_shoulder[0], mid_hip[0]]
y = [mid_shoulder[1], mid_hip[1]]
line = plt.plot(x, y)
plt.setp(line,
color=colors[len(kp_lines) + 1],
linewidth=1.0,
alpha=0.7)
plt.close('all')
def vis_one_image(
im, im_name, output_dir, boxes, segms=None, keypoints=None, thresh=0.9,
kp_thresh=2, dpi=200, box_alpha=0.0, dataset=None, show_class=False,
ext=None):
"""Visual debugging of detections."""
if not os.path.exists(output_dir):
os.makedirs(output_dir)
if isinstance(boxes, list):
boxes, segms, keypoints, classes = convert_from_cls_format(
boxes, segms, keypoints)
if boxes is None or boxes.shape[0] == 0 or max(boxes[:, 4]) < thresh:
return
dataset_keypoints, _ = keypoint_utils.get_keypoints()
if segms is not None and len(segms) > 0:
masks = mask_util.decode(segms)
color_list = colormap(rgb=True) / 255
kp_lines = kp_connections(dataset_keypoints)
cmap = plt.get_cmap('rainbow')
colors = [cmap(i) for i in np.linspace(0, 1, len(kp_lines) + 2)]
fig = plt.figure(frameon=False)
fig.set_size_inches(im.shape[1] / dpi, im.shape[0] / dpi)
ax = plt.Axes(fig, [0., 0., 1., 1.])
ax.axis('off')
fig.add_axes(ax)
ax.imshow(im)
# Display in largest to smallest order to reduce occlusion
areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
sorted_inds = np.argsort(-areas)
mask_color_id = 0
print ("arv9")
print('canshu%s:%s' % (9, sys.argv[9]))
C=[]
for i in sorted_inds:
bbox = boxes[i, :4]
score = boxes[i, -1]
if score < thresh:
continue
# show box (off by default)
ax.add_patch(
plt.Rectangle((bbox[0], bbox[1]),
bbox[2] - bbox[0],
bbox[3] - bbox[1],
fill=False, edgecolor='g',
linewidth=0.5, alpha=box_alpha))
if show_class:
ax.text(
bbox[0], bbox[1] - 2,
get_class_string(classes[i], score, dataset),
fontsize=3,
family='serif',
bbox=dict(
facecolor='g', alpha=0.4, pad=0, edgecolor='none'),
color='white')
# show mask
if segms is not None and len(segms) > i:
img = np.ones(im.shape)
color_mask = color_list[mask_color_id % len(color_list), 0:3]
mask_color_id += 1
w_ratio = .4
for c in range(3):
color_mask[c] = color_mask[c] * (1 - w_ratio) + w_ratio
for c in range(3):
img[:, :, c] = color_mask[c]
e = masks[:, :, i]
_, contour, hier = cv2.findContours(
e.copy(), cv2.RETR_CCOMP, cv2.CHAIN_APPROX_NONE)
for c in contour:
polygon = Polygon(
c.reshape((-1, 2)),
fill=True, facecolor=color_mask,
edgecolor='w', linewidth=0.2,
alpha=0.5)
ax.add_patch(polygon)
woca= c.reshape((-1, 2))
arr = np.array(woca)
key = np.unique(woca)
result = {}
k_qu=[]
arr_end=[]
arr_end2=[]
arr_endlast=[]
A=[]
B=[]
Bplus=[]
jjj=[]
Bwanmei=[]
Bcao=[]
Bao=[]
for k in key:
mask = (arr == k)
arr_new = arr[mask]
v = arr_new.size
result[k] = v
x=np.argwhere(arr== k)
if v>1:
x=np.argwhere(arr== k)
x=np.array(x)
x0=arr[:,0]
y0=arr[:,1]
y0lie=[]
for i in range(0,len(x0)):
if x0[i]!=k:
a=[]
if x0[i]==k:
y0lie.append(y0[i])
y0lienew=[]
arr_new=[]
arr_new_2=[]
if y0lie==[]:
a=[]
else:
miny0=np.min(y0lie)
maxy0=np.max(y0lie)
for i in range(miny0,maxy0+1):
y0lienew.append(i)
y0liezuizhong=[]
if y0lienew==[]:
a=[]
else:
miny0lienew=np.min(y0lienew)
maxy0lienew=np.max(y0lienew)
for i in range(miny0lienew,maxy0lienew+1):
y0liezuizhong.append(i)
for i in range(0,len(y0liezuizhong)):
arr_temp=[k,y0liezuizhong[i]]
arr_new.append(arr_temp)
arr_end.append(arr_new)
x0lie=[]
for i in range(0,len(y0)):
if y0[i]!=k:
a=[]
if y0[i]==k:
x0lie.append(x0[i])
x0lienew=[]
arr_new2=[]
if x0lie==[]:
a=[]
else:
minx0=np.min(x0lie)
maxx0=np.max(x0lie)
for i in range(minx0,maxx0+1):
x0lienew.append(i)
x0liezuizhong=[]
if x0lienew==[]:
a=[]
else:
minx0lienew=np.min(x0lienew)
maxx0lienew=np.max(x0lienew)
for i in range(minx0lienew,maxx0lienew+1):
x0liezuizhong.append(i)
for i in range(0,len(x0liezuizhong)):
arr_temp=[x0liezuizhong[i],k]
arr_new2.append(arr_temp)
arr_end2.append(arr_new2)
arr_endlast=arr_end+arr_end2
A=list(chain(*arr_endlast))
B=np.array(list(set([tuple(t) for t in A])))
print ("********************************")
# show keypoints
if keypoints is not None and len(keypoints) > i:
kps = keypoints[i]
plt.autoscale(False)
for l in range(len(kp_lines)):
i1 = kp_lines[l][0]
i2 = kp_lines[l][1]
if kps[2, i1] > kp_thresh and kps[2, i2] > kp_thresh:
x = [kps[0, i1], kps[0, i2]]
y = [kps[1, i1], kps[1, i2]]
line = plt.plot(x, y)
plt.setp(line, color=colors[l], linewidth=1.0, alpha=0.7)
if kps[2, i1] > kp_thresh:
plt.plot(
kps[0, i1], kps[1, i1], '.', color=colors[l],
markersize=3.0, alpha=0.7)
if kps[2, i2] > kp_thresh:
plt.plot(
kps[0, i2], kps[1, i2], '.', color=colors[l],
markersize=3.0, alpha=0.7)
# add mid shoulder / mid hip for better visualization
mid_shoulder = (
kps[:2, dataset_keypoints.index('right_shoulder')] +
kps[:2, dataset_keypoints.index('left_shoulder')]) / 2.0
sc_mid_shoulder = np.minimum(
kps[2, dataset_keypoints.index('right_shoulder')],
kps[2, dataset_keypoints.index('left_shoulder')])
mid_hip = (
kps[:2, dataset_keypoints.index('right_hip')] +
kps[:2, dataset_keypoints.index('left_hip')]) / 2.0
sc_mid_hip = np.minimum(
kps[2, dataset_keypoints.index('right_hip')],
kps[2, dataset_keypoints.index('left_hip')])
if (sc_mid_shoulder > kp_thresh and
kps[2, dataset_keypoints.index('nose')] > kp_thresh):
x = [mid_shoulder[0], kps[0, dataset_keypoints.index('nose')]]
y = [mid_shoulder[1], kps[1, dataset_keypoints.index('nose')]]
line = plt.plot(x, y)
plt.setp(
line, color=colors[len(kp_lines)], linewidth=1.0, alpha=0.7)
if sc_mid_shoulder > kp_thresh and sc_mid_hip > kp_thresh:
x = [mid_shoulder[0], mid_hip[0]]
y = [mid_shoulder[1], mid_hip[1]]
line = plt.plot(x, y)
plt.setp(
line, color=colors[len(kp_lines) + 1], linewidth=1.0,
alpha=0.7)
#print ("CCCCCCCCCCCCCCCCCCC*********************")
#C=np.array(C)
#print(C)
output_name = os.path.basename(im_name)
fig.savefig(os.path.join(output_dir, '{}'.format(output_name)), dpi=dpi)
plt.close('all')
def vis_one_image(im,
im_name,
output_dir,
boxes,
segms=None,
keypoints=None,
body_uv=None,
thresh=0.9,
kp_thresh=2,
dpi=200,
box_alpha=0.0,
dataset=None,
show_class=False,
ext='pdf'):
"""Visual debugging of detections."""
if not os.path.exists(output_dir):
os.makedirs(output_dir)
if isinstance(boxes, list):
boxes, segms, keypoints, classes = convert_from_cls_format(
boxes, segms, keypoints)
if boxes is None or boxes.shape[0] == 0 or max(boxes[:, 4]) < thresh:
return
if segms is not None:
masks = mask_util.decode(segms)
color_list = colormap(rgb=True) / 255
dataset_keypoints, _ = keypoint_utils.get_keypoints()
kp_lines = kp_connections(dataset_keypoints)
cmap = plt.get_cmap('rainbow')
colors = [cmap(i) for i in np.linspace(0, 1, len(kp_lines) + 2)]
fig = plt.figure(frameon=False)
fig.set_size_inches(im.shape[1] / dpi, im.shape[0] / dpi)
ax = plt.Axes(fig, [0., 0., 1., 1.])
ax.axis('off')
fig.add_axes(ax)
ax.imshow(im)
# Display in largest to smallest order to reduce occlusion
areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
sorted_inds = np.argsort(-areas)
mask_color_id = 0
for i in sorted_inds:
bbox = boxes[i, :4]
score = boxes[i, -1]
if score < thresh:
continue
print(dataset.classes[classes[i]], score)
# show box (off by default, box_alpha=0.0)
ax.add_patch(
plt.Rectangle((bbox[0], bbox[1]),
bbox[2] - bbox[0],
bbox[3] - bbox[1],
fill=False,
edgecolor='g',
linewidth=0.5,
alpha=box_alpha))
if show_class:
ax.text(bbox[0],
bbox[1] - 2,
get_class_string(classes[i], score, dataset),
fontsize=3,
family='serif',
bbox=dict(facecolor='g',
alpha=0.4,
pad=0,
edgecolor='none'),
color='white')
# show mask
if segms is not None and len(segms) > i:
img = np.ones(im.shape)
color_mask = color_list[mask_color_id % len(color_list), 0:3]
mask_color_id += 1
w_ratio = .4
for c in range(3):
color_mask[c] = color_mask[c] * (1 - w_ratio) + w_ratio
for c in range(3):
img[:, :, c] = color_mask[c]
e = masks[:, :, i]
_, contour, hier = cv2.findContours(e.copy(), cv2.RETR_CCOMP,
cv2.CHAIN_APPROX_NONE)
for c in contour:
polygon = Polygon(c.reshape((-1, 2)),
fill=True,
facecolor=color_mask,
edgecolor='w',
linewidth=1.2,
alpha=0.5)
ax.add_patch(polygon)
# show keypoints
if keypoints is not None and len(keypoints) > i:
kps = keypoints[i]
plt.autoscale(False)
for l in range(len(kp_lines)):
i1 = kp_lines[l][0]
i2 = kp_lines[l][1]
if kps[2, i1] > kp_thresh and kps[2, i2] > kp_thresh:
x = [kps[0, i1], kps[0, i2]]
y = [kps[1, i1], kps[1, i2]]
line = ax.plot(x, y)
plt.setp(line, color=colors[l], linewidth=1.0, alpha=0.7)
if kps[2, i1] > kp_thresh:
ax.plot(kps[0, i1],
kps[1, i1],
'.',
color=colors[l],
markersize=3.0,
alpha=0.7)
if kps[2, i2] > kp_thresh:
ax.plot(kps[0, i2],
kps[1, i2],
'.',
color=colors[l],
markersize=3.0,
alpha=0.7)
# add mid shoulder / mid hip for better visualization
mid_shoulder = (
kps[:2, dataset_keypoints.index('right_shoulder')] +
kps[:2, dataset_keypoints.index('left_shoulder')]) / 2.0
sc_mid_shoulder = np.minimum(
kps[2, dataset_keypoints.index('right_shoulder')],
kps[2, dataset_keypoints.index('left_shoulder')])
mid_hip = (kps[:2, dataset_keypoints.index('right_hip')] +
kps[:2, dataset_keypoints.index('left_hip')]) / 2.0
sc_mid_hip = np.minimum(
kps[2, dataset_keypoints.index('right_hip')],
kps[2, dataset_keypoints.index('left_hip')])
if (sc_mid_shoulder > kp_thresh
and kps[2, dataset_keypoints.index('nose')] > kp_thresh):
x = [mid_shoulder[0], kps[0, dataset_keypoints.index('nose')]]
y = [mid_shoulder[1], kps[1, dataset_keypoints.index('nose')]]
line = ax.plot(x, y)
plt.setp(line,
color=colors[len(kp_lines)],
linewidth=1.0,
alpha=0.7)
if sc_mid_shoulder > kp_thresh and sc_mid_hip > kp_thresh:
x = [mid_shoulder[0], mid_hip[0]]
y = [mid_shoulder[1], mid_hip[1]]
line = ax.plot(x, y)
plt.setp(line,
color=colors[len(kp_lines) + 1],
linewidth=1.0,
alpha=0.7)
# DensePose Visualization Starts!!
## Get full IUV image out
if body_uv is not None and len(body_uv) > 1:
IUV_fields = body_uv[1]
#
All_Coords = np.zeros(im.shape)
All_inds = np.zeros([im.shape[0], im.shape[1]])
K = 26
##
inds = np.argsort(boxes[:, 4])
##
for i, ind in enumerate(inds):
entry = boxes[ind, :]
if entry[4] > 0.65:
entry = entry[0:4].astype(int)
####
output = IUV_fields[ind]
####
All_Coords_Old = All_Coords[entry[1]:entry[1] +
output.shape[1],
entry[0]:entry[0] +
output.shape[2], :]
All_Coords_Old[All_Coords_Old == 0] = output.transpose(
[1, 2, 0])[All_Coords_Old == 0]
All_Coords[entry[1]:entry[1] + output.shape[1],
entry[0]:entry[0] +
output.shape[2], :] = All_Coords_Old
###
CurrentMask = (output[0, :, :] > 0).astype(np.float32)
All_inds_old = All_inds[entry[1]:entry[1] + output.shape[1],
entry[0]:entry[0] + output.shape[2]]
All_inds_old[All_inds_old ==
0] = CurrentMask[All_inds_old == 0] * i
All_inds[entry[1]:entry[1] + output.shape[1],
entry[0]:entry[0] + output.shape[2]] = All_inds_old
#
All_Coords[:, :, 1:3] = 255. * All_Coords[:, :, 1:3]
All_Coords[All_Coords > 255] = 255.
All_Coords = All_Coords.astype(np.uint8)
All_inds = All_inds.astype(np.uint8)
#
IUV_SaveName = os.path.basename(im_name).split('.')[0] + '_IUV.png'
INDS_SaveName = os.path.basename(im_name).split('.')[0] + '_INDS.png'
cv2.imwrite(os.path.join(output_dir, '{}'.format(IUV_SaveName)),
All_Coords)
cv2.imwrite(os.path.join(output_dir, '{}'.format(INDS_SaveName)),
All_inds)
print('IUV written to: ',
os.path.join(output_dir, '{}'.format(IUV_SaveName)))
###
### DensePose Visualization Done!!
#
output_name = os.path.basename(im_name) + '.' + ext
fig.savefig(os.path.join(output_dir, '{}'.format(output_name)), dpi=dpi)
plt.close('all')
# SMPL Visualization
if body_uv is not None and len(body_uv) > 2:
smpl_fields = body_uv[2]
#
All_Coords = np.zeros(im.shape)
# All_inds = np.zeros([im.shape[0], im.shape[1]])
K = 26
##
inds = np.argsort(boxes[:, 4])
##
for i, ind in enumerate(inds):
entry = boxes[ind, :]
if entry[4] > 0.75:
entry = entry[0:4].astype(int)
center_roi = [(entry[2] + entry[0]) / 2.,
(entry[3] + entry[1]) / 2.]
####
output, center_out = smpl_fields[ind]
####
x1_img = max(int(center_roi[0] - center_out[0]), 0)
y1_img = max(int(center_roi[1] - center_out[1]), 0)
x2_img = min(
int(center_roi[0] - center_out[0]) + output.shape[2],
im.shape[1])
y2_img = min(
int(center_roi[1] - center_out[1]) + output.shape[1],
im.shape[0])
All_Coords_Old = All_Coords[y1_img:y2_img, x1_img:x2_img, :]
x1_out = max(int(center_out[0] - center_roi[0]), 0)
y1_out = max(int(center_out[1] - center_roi[1]), 0)
x2_out = x1_out + (x2_img - x1_img)
y2_out = y1_out + (y2_img - y1_img)
output = output[:, y1_out:y2_out, x1_out:x2_out]
# All_Coords_Old = All_Coords[entry[1]: entry[1] + output.shape[1], entry[0]:entry[0] + output.shape[2],
# :]
All_Coords_Old[All_Coords_Old == 0] = output.transpose(
[1, 2, 0])[All_Coords_Old == 0]
All_Coords[y1_img:y2_img, x1_img:x2_img, :] = All_Coords_Old
###
# CurrentMask = (output[0, :, :] > 0).astype(np.float32)
# All_inds_old = All_inds[entry[1]: entry[1] + output.shape[1], entry[0]:entry[0] + output.shape[2]]
# All_inds_old[All_inds_old == 0] = CurrentMask[All_inds_old == 0] * i
# All_inds[entry[1]: entry[1] + output.shape[1], entry[0]:entry[0] + output.shape[2]] = All_inds_old
#
All_Coords = 255. * All_Coords
All_Coords[All_Coords > 255] = 255.
All_Coords = All_Coords.astype(np.uint8)
image_stacked = im[:, :, ::-1]
image_stacked[All_Coords > 20] = All_Coords[All_Coords > 20]
# All_inds = All_inds.astype(np.uint8)
#
SMPL_SaveName = os.path.basename(im_name).split('.')[0] + '_SMPL.png'
smpl_image_SaveName = os.path.basename(im_name).split(
'.')[0] + '_SMPLimg.png'
# INDS_SaveName = os.path.basename(im_name).split('.')[0] + '_INDS.png'
cv2.imwrite(os.path.join(output_dir, '{}'.format(SMPL_SaveName)),
All_Coords)
cv2.imwrite(os.path.join(output_dir, '{}'.format(smpl_image_SaveName)),
image_stacked)
# cv2.imwrite(os.path.join(output_dir, '{}'.format(INDS_SaveName)), All_inds)
print('SMPL written to: ',
os.path.join(output_dir, '{}'.format(SMPL_SaveName)))
###
### SMPL Visualization Done!!
#
output_name = os.path.basename(im_name) + '.' + ext
fig.savefig(os.path.join(output_dir, '{}'.format(output_name)), dpi=dpi)
plt.close('all')
def main(args):
cfg_file = r'/home/twang/Documents/detectron/configs/12_2017_baselines/e2e_mask_rcnn_R-101-FPN_2x.yaml'
weights_file = r'/home/twang/Documents/detectron/model-weights/mask_rcnn_R-101-FPN_2x_model_final.pkl'
video_dir = args.video_dir
print("video_dir", video_dir)
video_name = os.path.basename(video_dir)
video_name = os.path.splitext(video_name)[0]
print("video_name", video_name)
directory_box = os.path.join(
os.path.join(r"/home/twang/Documents/HK-person", video_name), 'box')
print("directory_box", directory_box)
os.makedirs(directory_box)
directory_mask = os.path.join(
os.path.join(r"/home/twang/Documents/HK-person", video_name), 'mask')
print("directory_mask", directory_mask)
os.makedirs(directory_mask)
merge_cfg_from_file(cfg_file)
cfg.NUM_GPUS = 1
weights = cache_url(weights_file, cfg.DOWNLOAD_CACHE)
assert_and_infer_cfg()
model = infer_engine.initialize_model_from_cfg(weights)
dummy_coco_dataset = dummy_datasets.get_coco_dataset()
cap = cv2.VideoCapture(video_dir)
count = 0
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
frame = cv2.resize(frame, dsize=(1280, 720))
total_frame = cap.get(cv2.CAP_PROP_FRAME_COUNT)
current_frame = int(cap.get(cv2.CAP_PROP_POS_FRAMES))
Frame_step = 5
if current_frame + Frame_step < total_frame:
cap.set(cv2.CAP_PROP_POS_FRAMES, current_frame + Frame_step)
timers = defaultdict(Timer)
with c2_utils.NamedCudaScope(0):
cls_boxes, cls_segms, cls_keyps = infer_engine.im_detect_all(
model, frame, None, timers=timers)
thresh = 0.9
crop_box = True
dataset = dummy_coco_dataset
frame_for_box_crop = frame.copy()
frame_for_mask = frame.copy()
"""Constructs a numpy array with the detections visualized."""
if isinstance(cls_boxes, list):
boxes, segms, keypoints, classes = convert_from_cls_format(
cls_boxes, cls_segms, cls_keyps)
if boxes is None or boxes.shape[0] == 0 or max(boxes[:,
4]) < thresh:
return frame
if segms is not None and len(segms) > 0:
masks = mask_util.decode(segms)
color_list = colormap()
mask_color_id = 0
# Display in largest to smallest order to reduce occlusion
areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
sorted_inds = np.argsort(-areas)
for i in sorted_inds:
bbox = boxes[i, :4]
score = boxes[i, -1]
if score < thresh:
continue
# crop each box
if crop_box:
#frame = vis_bbox(frame, (bbox[0], bbox[1], bbox[2] - bbox[0], bbox[3] - bbox[1]))
(x1, y1, w, h) = (int(bbox[0]), int(bbox[1]),
int(bbox[2] - bbox[0]),
int(bbox[3] - bbox[1]))
x2 = x1 + w
y2 = y1 + h
cropped = frame_for_box_crop[y1:y2, x1:x2]
cv2.imwrite(
"%s/person_Frame%i_%i.png" %
(directory_box, current_frame, i), cropped)
# crop each mask
if segms is not None and len(segms) > i:
color_mask = color_list[mask_color_id % len(color_list),
0:3]
mask_color_id += 1
#frame = vis_mask(frame, masks[..., i], color_mask)
(x1, y1, w, h) = (int(bbox[0]), int(bbox[1]),
int(bbox[2] - bbox[0]),
int(bbox[3] - bbox[1]))
x2 = x1 + w
y2 = y1 + h
cropped_mask = masks[..., i]
cropped_mask = cropped_mask[y1:y2, x1:x2]
cropped_img = frame_for_mask[y1:y2, x1:x2]
cropped_img = vis_mask(cropped_img, cropped_mask,
color_mask)
cv2.imwrite(
"%s/person_Mask_Frame%i_%i.png" %
(directory_mask, current_frame, i), cropped_img)
count += 1
print("done:%i" % count)
else:
pass
cap.release()
cv2.destroyAllWindows()
def vis_one_image(
im, im_name, output_dir, boxes, segms=None, keypoints=None, thresh=0.9,
kp_thresh=2, dpi=200, box_alpha=0.0, dataset=None, show_class=False,
ext='pdf'):
"""Visual debugging of detections."""
if not os.path.exists(output_dir):
os.makedirs(output_dir)
if isinstance(boxes, list):
boxes, segms, keypoints, classes = convert_from_cls_format(
boxes, segms, keypoints)
if boxes is None or boxes.shape[0] == 0 or max(boxes[:, 4]) < thresh:
return
dataset_keypoints, _ = keypoint_utils.get_keypoints()
if segms is not None:
masks = mask_util.decode(segms)
color_list = colormap(rgb=True) / 255
kp_lines = kp_connections(dataset_keypoints)
cmap = plt.get_cmap('rainbow')
colors = [cmap(i) for i in np.linspace(0, 1, len(kp_lines) + 2)]
fig = plt.figure(frameon=False)
fig.set_size_inches(im.shape[1] / dpi, im.shape[0] / dpi)
ax = plt.Axes(fig, [0., 0., 1., 1.])
ax.axis('off')
fig.add_axes(ax)
ax.imshow(im)
# Display in largest to smallest order to reduce occlusion
areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
sorted_inds = np.argsort(-areas)
mask_color_id = 0
for i in sorted_inds:
bbox = boxes[i, :4]
score = boxes[i, -1]
if score < thresh:
continue
# show box (off by default)
ax.add_patch(
plt.Rectangle((bbox[0], bbox[1]),
bbox[2] - bbox[0],
bbox[3] - bbox[1],
fill=False, edgecolor='g',
linewidth=0.5, alpha=box_alpha))
if show_class:
ax.text(
bbox[0], bbox[1] - 2,
get_class_string(classes[i], score, dataset),
fontsize=3,
family='serif',
bbox=dict(
facecolor='g', alpha=0.4, pad=0, edgecolor='none'),
color='white')
# show mask
if segms is not None and len(segms) > i:
img = np.ones(im.shape)
color_mask = color_list[mask_color_id % len(color_list), 0:3]
mask_color_id += 1
w_ratio = .4
for c in range(3):
color_mask[c] = color_mask[c] * (1 - w_ratio) + w_ratio
for c in range(3):
img[:, :, c] = color_mask[c]
e = masks[:, :, i]
_, contour, hier = cv2.findContours(
e.copy(), cv2.RETR_CCOMP, cv2.CHAIN_APPROX_NONE)
for c in contour:
polygon = Polygon(
c.reshape((-1, 2)),
fill=True, facecolor=color_mask,
edgecolor='w', linewidth=1.2,
alpha=0.5)
ax.add_patch(polygon)
# show keypoints
if keypoints is not None and len(keypoints) > i:
kps = keypoints[i]
plt.autoscale(False)
for l in range(len(kp_lines)):
i1 = kp_lines[l][0]
i2 = kp_lines[l][1]
if kps[2, i1] > kp_thresh and kps[2, i2] > kp_thresh:
x = [kps[0, i1], kps[0, i2]]
y = [kps[1, i1], kps[1, i2]]
line = plt.plot(x, y)
plt.setp(line, color=colors[l], linewidth=1.0, alpha=0.7)
if kps[2, i1] > kp_thresh:
plt.plot(
kps[0, i1], kps[1, i1], '.', color=colors[l],
markersize=3.0, alpha=0.7)
if kps[2, i2] > kp_thresh:
plt.plot(
kps[0, i2], kps[1, i2], '.', color=colors[l],
markersize=3.0, alpha=0.7)
# add mid shoulder / mid hip for better visualization
mid_shoulder = (
kps[:2, dataset_keypoints.index('right_shoulder')] +
kps[:2, dataset_keypoints.index('left_shoulder')]) / 2.0
sc_mid_shoulder = np.minimum(
kps[2, dataset_keypoints.index('right_shoulder')],
kps[2, dataset_keypoints.index('left_shoulder')])
mid_hip = (
kps[:2, dataset_keypoints.index('right_hip')] +
kps[:2, dataset_keypoints.index('left_hip')]) / 2.0
sc_mid_hip = np.minimum(
kps[2, dataset_keypoints.index('right_hip')],
kps[2, dataset_keypoints.index('left_hip')])
if (sc_mid_shoulder > kp_thresh and
kps[2, dataset_keypoints.index('nose')] > kp_thresh):
x = [mid_shoulder[0], kps[0, dataset_keypoints.index('nose')]]
y = [mid_shoulder[1], kps[1, dataset_keypoints.index('nose')]]
line = plt.plot(x, y)
plt.setp(
line, color=colors[len(kp_lines)], linewidth=1.0, alpha=0.7)
if sc_mid_shoulder > kp_thresh and sc_mid_hip > kp_thresh:
x = [mid_shoulder[0], mid_hip[0]]
y = [mid_shoulder[1], mid_hip[1]]
line = plt.plot(x, y)
plt.setp(
line, color=colors[len(kp_lines) + 1], linewidth=1.0,
alpha=0.7)
output_name = os.path.basename(im_name) + '.' + ext
fig.savefig(os.path.join(output_dir, '{}'.format(output_name)), dpi=dpi)
plt.close('all')
def vis_one_image_srishti(
im, im_name, output_dir, boxes, segms=None, keypoints=None, thresh=0.9,
kp_thresh=2, dpi=200, box_alpha=0.0, dataset=None, show_class=False,
ext='pdf'):
"""Visual debugging of detections."""
if not os.path.exists(output_dir):
os.makedirs(output_dir)
if isinstance(boxes, list):
boxes, segms, keypoints, classes = convert_from_cls_format(
boxes, segms, keypoints)
if boxes is None or boxes.shape[0] == 0 or max(boxes[:, 4]) < thresh:
return
dataset_keypoints, _ = keypoint_utils.get_keypoints()
if segms is not None:
masks = mask_util.decode(segms)
color_list = colormap(rgb=True) / 255
kp_lines = kp_connections(dataset_keypoints)
cmap = plt.get_cmap('rainbow')
colors = [cmap(i) for i in np.linspace(0, 1, len(kp_lines) + 2)]
fig = plt.figure(frameon=False)
fig.set_size_inches(im.shape[1] / dpi, im.shape[0] / dpi)
ax = plt.Axes(fig, [0., 0., 1., 1.])
ax.axis('off')
fig.add_axes(ax)
ax.imshow(im)
# Display in largest to smallest order to reduce occlusion
areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
sorted_inds = np.argsort(-areas)
mask_color_id = 0
for i in sorted_inds:
bbox = boxes[i, :4]
score = boxes[i, -1]
if score < thresh:
continue
# show mask
if segms is not None and len(segms) > i:
img = np.ones(im.shape)
color_mask = color_list[mask_color_id % len(color_list), 0:3]
mask_color_id += 1
w_ratio = .4
for c in range(3):
color_mask[c] = color_mask[c] * (1 - w_ratio) + w_ratio
for c in range(3):
img[:, :, c] = color_mask[c]
e = masks[:, :, i]
_, contour, hier = cv2.findContours(
e.copy(), cv2.RETR_CCOMP, cv2.CHAIN_APPROX_NONE)
iterator = 0
for c in contour:
iterator = iterator +1
if iterator > 2:
break
polygon = Polygon(
c.reshape((-1, 2)),
fill=True, facecolor=color_mask,
edgecolor='w', linewidth=1.2,
alpha=0.5)
ax.add_patch(polygon)
output_name_png = os.path.basename(im_name) + '.' + 'png'
fig.savefig(os.path.join(output_dir, '{}'.format(output_name_png)), dpi=dpi)
plt.close('all')
def vis_one_image(im,
im_name,
output_dir,
boxes,
segms=None,
keypoints=None,
thresh=0.9,
kp_thresh=2,
dpi=200,
box_alpha=0.0,
dataset=None,
show_class=False,
ext='jpg'):
"""Visual debugging of detections."""
if not os.path.exists(output_dir):
os.makedirs(output_dir)
if isinstance(boxes, list):
boxes, segms, keypoints, classes = convert_from_cls_format(
boxes, segms, keypoints)
if boxes is None or boxes.shape[0] == 0 or max(boxes[:, 4]) < thresh:
return
dataset_keypoints, _ = keypoint_utils.get_keypoints()
if segms is not None:
masks = mask_util.decode(segms)
color_list = colormap(rgb=True) / 255
kp_lines = kp_connections(dataset_keypoints)
cmap = plt.get_cmap('rainbow')
colors = [cmap(i) for i in np.linspace(0, 1, len(kp_lines) + 2)]
fig = plt.figure(frameon=False)
fig.set_size_inches(im.shape[1] / dpi, im.shape[0] / dpi)
ax = plt.Axes(fig, [0., 0., 1., 1.])
ax.axis('off')
fig.add_axes(ax)
ax.imshow(im)
# Display in largest to smallest order to reduce occlusion
areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
sorted_inds = np.argsort(-areas)
mask_color_id = 0
for i in sorted_inds:
bbox = boxes[i, :4]
score = boxes[i, -1]
if score < thresh:
continue
# show box (off by default)
ax.add_patch(
plt.Rectangle((bbox[0], bbox[1]),
bbox[2] - bbox[0],
bbox[3] - bbox[1],
fill=False,
edgecolor='g',
linewidth=0.5,
alpha=box_alpha))
if show_class:
ax.text(bbox[0],
bbox[1] - 2,
get_class_string(classes[i], score, dataset),
fontsize=10,
family='serif',
bbox=dict(facecolor='g',
alpha=0.4,
pad=0,
edgecolor='none'),
color='white')
# show mask
if segms is not None and len(segms) > i:
img = np.ones(im.shape)
color_mask = color_list[mask_color_id % len(color_list), 0:3]
mask_color_id += 1
w_ratio = .4
for c in range(3):
color_mask[c] = color_mask[c] * (1 - w_ratio) + w_ratio
for c in range(3):
img[:, :, c] = color_mask[c]
e = masks[:, :, i]
_, contour, hier = cv2.findContours(e.copy(), cv2.RETR_CCOMP,
cv2.CHAIN_APPROX_NONE)
for c in contour:
polygon = Polygon(c.reshape((-1, 2)),
fill=True,
facecolor=color_mask,
edgecolor='w',
linewidth=1.2,
alpha=0.5)
ax.add_patch(polygon)
# show keypoints
if keypoints is not None and len(keypoints) > i:
kps = keypoints[i]
plt.autoscale(False)
for l in range(len(kp_lines)):
i1 = kp_lines[l][0]
i2 = kp_lines[l][1]
if kps[2, i1] > kp_thresh and kps[2, i2] > kp_thresh:
x = [kps[0, i1], kps[0, i2]]
y = [kps[1, i1], kps[1, i2]]
line = plt.plot(x, y)
plt.setp(line, color=colors[l], linewidth=1.0, alpha=0.7)
if kps[2, i1] > kp_thresh:
plt.plot(kps[0, i1],
kps[1, i1],
'.',
color=colors[l],
markersize=3.0,
alpha=0.7)
if kps[2, i2] > kp_thresh:
plt.plot(kps[0, i2],
kps[1, i2],
'.',
color=colors[l],
markersize=3.0,
alpha=0.7)
# add mid shoulder / mid hip for better visualization
mid_shoulder = (
kps[:2, dataset_keypoints.index('right_shoulder')] +
kps[:2, dataset_keypoints.index('left_shoulder')]) / 2.0
sc_mid_shoulder = np.minimum(
kps[2, dataset_keypoints.index('right_shoulder')],
kps[2, dataset_keypoints.index('left_shoulder')])
mid_hip = (kps[:2, dataset_keypoints.index('right_hip')] +
kps[:2, dataset_keypoints.index('left_hip')]) / 2.0
sc_mid_hip = np.minimum(
kps[2, dataset_keypoints.index('right_hip')],
kps[2, dataset_keypoints.index('left_hip')])
if (sc_mid_shoulder > kp_thresh
and kps[2, dataset_keypoints.index('nose')] > kp_thresh):
x = [mid_shoulder[0], kps[0, dataset_keypoints.index('nose')]]
y = [mid_shoulder[1], kps[1, dataset_keypoints.index('nose')]]
line = plt.plot(x, y)
plt.setp(line,
color=colors[len(kp_lines)],
linewidth=1.0,
alpha=0.7)
if sc_mid_shoulder > kp_thresh and sc_mid_hip > kp_thresh:
x = [mid_shoulder[0], mid_hip[0]]
y = [mid_shoulder[1], mid_hip[1]]
line = plt.plot(x, y)
plt.setp(line,
color=colors[len(kp_lines) + 1],
linewidth=1.0,
alpha=0.7)
image = fig2data(fig)
image = image[..., ::-1]
cv2.imshow("img", image)
cv2.waitKey(1)
# output_name = os.path.basename(im_name) + '.' + ext
# fig.savefig(os.path.join(output_dir, '{}'.format(output_name)), dpi=dpi)
plt.close('all')
def segmented_images(
im, im_name, output_dir, boxes, segms=None, keypoints=None, thresh=0.9,
kp_thresh=2, dpi=200, box_alpha=0.0, dataset=None, show_class=False,
ext='pdf'):
"""Extract segmented images."""
if not os.path.exists(output_dir):
os.makedirs(output_dir)
if isinstance(boxes, list):
boxes, segms, keypoints, classes = convert_from_cls_format(
boxes, segms, keypoints)
if boxes is None or boxes.shape[0] == 0 or max(boxes[:, 4]) < thresh:
return
if segms is not None:
masks = mask_util.decode(segms)
color_list = colormap(rgb=True) / 255
# Display in largest to smallest order to reduce occlusion
areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
sorted_inds = np.argsort(-areas)
mask_color_id = 0
segmented_images = []
segmented_classes = []
segmented_scores = []
segmented_binary_masks = []
contours = []
for i in sorted_inds:
bbox = boxes[i, :4]
score = boxes[i, -1]
if score < thresh:
continue
# show mask
if segms is not None and len(segms) > i:
img = np.zeros(im.shape)
color_mask = color_list[mask_color_id % len(color_list), 0:3]
mask_color_id += 1
w_ratio = .4
e = masks[:, :, i]
# testing code
# im_seg_mask = vis_binary_mask(im, masks[..., i])
# cv2.imwrite('/home/srishti/testMaskIMG2_' + str(i) + '.png', im_seg_mask)
for channel in range(3):
img[:,:, channel] = im[:,:, channel] * e[:,:]
_, contour, hier = cv2.findContours(
e.copy(), cv2.RETR_CCOMP, cv2.CHAIN_APPROX_NONE)
for c in contour:
x, y, w, h = cv2.boundingRect(c)
img_countour = np.zeros([h, w, 3])
img_countour_bin = np.zeros([h, w])
for channel in range(3):
img_countour[:, :, channel] = img[y:h + y, x:w + x, channel]
img_countour_bin[:, :] = e[y:h + y, x:w + x]
segmented_images.insert(len(segmented_images), img_countour)
segmented_binary_masks.insert(len(segmented_binary_masks), img_countour_bin)
segmented_classes.insert(len(segmented_classes), dataset.classes[classes[i]])
segmented_scores.insert(len(segmented_scores), score)
# contours.insert(len(contours), c)
return segmented_images, segmented_classes, segmented_scores, segmented_binary_masks
