import argparse
import cv2
import os
import json
import numpy as np
from PIL import Image as PILImage
import joblib
def mask_nms(masks, bbox_scores, instances_confidence_threshold=0.5, overlap_threshold=0.7):
"""
NMS-like procedure used in Panoptic Segmentation
Remove the overlap areas of different instances in Instance Segmentation
"""
panoptic_seg = np.zeros(masks.shape[:2], dtype=np.uint8)
sorted_inds = list(range(len(bbox_scores)))
current_segment_id = 0
segments_score = []
for inst_id in sorted_inds:
score = bbox_scores[inst_id]
if score < instances_confidence_threshold:
break
mask = masks[:, :, inst_id]
mask_area = mask.sum()
if mask_area == 0:
continue
intersect = (mask > 0) & (panoptic_seg > 0)
intersect_area = intersect.sum()
if intersect_area * 1.0 / mask_area > overlap_threshold:
continue
if intersect_area > 0:
mask = mask & (panoptic_seg == 0)
current_segment_id += 1
# panoptic_seg[np.where(mask==1)] = current_segment_id
# panoptic_seg = panoptic_seg + current_segment_id*mask
panoptic_seg = np.where(mask == 0, panoptic_seg, current_segment_id)
segments_score.append(score)
# print(np.unique(panoptic_seg))
return panoptic_seg, segments_score
def extend(si, sj, instance_label, global_label, panoptic_seg_mask, class_map):
"""
"""
directions = [[-1, 0], [0, 1], [1, 0], [0, -1],
[1, 1], [1, -1], [-1, 1], [-1, -1]]
inst_class = instance_label[si, sj]
human_class = panoptic_seg_mask[si, sj]
global_class = class_map[inst_class]
queue = [[si, sj]]
while len(queue) != 0:
cur = queue[0]
queue.pop(0)
for direction in directions:
ni = cur[0] + direction[0]
nj = cur[1] + direction[1]
if ni >= 0 and nj >= 0 and \
ni < instance_label.shape[0] and \
nj < instance_label.shape[1] and \
instance_label[ni, nj] == 0 and \
global_label[ni, nj] == global_class:
instance_label[ni, nj] = inst_class
# Using refined instance label to refine human label
panoptic_seg_mask[ni, nj] = human_class
queue.append([ni, nj])
def refine(instance_label, panoptic_seg_mask, global_label, class_map):
"""
Inputs:
[ instance_label ]
np.array() with shape [h, w]
[ global_label ] with shape [h, w]
np.array()
"""
for i in range(instance_label.shape[0]):
for j in range(instance_label.shape[1]):
if instance_label[i, j] != 0:
extend(i, j, instance_label, global_label, panoptic_seg_mask, class_map)
def get_palette(num_cls):
""" Returns the color map for visualizing the segmentation mask.
Inputs:
=num_cls=
Number of classes.
Returns:
The color map.
"""
n = num_cls
palette = [0] * (n * 3)
for j in range(0, n):
lab = j
palette[j * 3 + 0] = 0
palette[j * 3 + 1] = 0
palette[j * 3 + 2] = 0
i = 0
while lab:
palette[j * 3 + 0] |= (((lab >> 0) & 1) << (7 - i))
palette[j * 3 + 1] |= (((lab >> 1) & 1) << (7 - i))
palette[j * 3 + 2] |= (((lab >> 2) & 1) << (7 - i))
i += 1
lab >>= 3
return palette
def patch2img_output(patch_dir, img_name, img_height, img_width, bbox, bbox_type, num_class):
"""transform bbox patch outputs to image output"""
assert bbox_type == 'gt' or 'msrcnn'
output = np.zeros((img_height, img_width, num_class), dtype='float')
output[:, :, 0] = np.inf
count_predictions = np.zeros((img_height, img_width, num_class), dtype='int32')
for i in range(len(bbox)): # person index starts from 1
file_path = os.path.join(patch_dir, os.path.splitext(img_name)[0] + '_' + str(i + 1) + '_' + bbox_type + '.npy')
bbox_output = np.load(file_path)
output[bbox[i][1]:bbox[i][3] + 1, bbox[i][0]:bbox[i][2] + 1, 1:] += bbox_output[:, :, 1:]
count_predictions[bbox[i][1]:bbox[i][3] + 1, bbox[i][0]:bbox[i][2] + 1, 1:] += 1
output[bbox[i][1]:bbox[i][3] + 1, bbox[i][0]:bbox[i][2] + 1, 0] \
= np.minimum(output[bbox[i][1]:bbox[i][3] + 1, bbox[i][0]:bbox[i][2] + 1, 0], bbox_output[:, :, 0])
# Caution zero dividing.
count_predictions[count_predictions == 0] = 1
return output / count_predictions
def get_instance(cat_gt, panoptic_seg_mask):
"""
"""
instance_gt = np.zeros_like(cat_gt, dtype=np.uint8)
num_humans = len(np.unique(panoptic_seg_mask)) - 1
class_map = {}
total_part_num = 0
for id in range(1, num_humans + 1):
human_part_label = np.where(panoptic_seg_mask == id, cat_gt, 0).astype(np.uint8)
# human_part_label = (np.where(panoptic_seg_mask==id) * cat_gt).astype(np.uint8)
part_classes = np.unique(human_part_label)
exceed = False
for part_id in part_classes:
if part_id == 0: # background
continue
total_part_num += 1
if total_part_num > 255:
print("total_part_num exceed, return current instance map: {}".format(total_part_num))
exceed = True
break
class_map[total_part_num] = part_id
instance_gt[np.where(human_part_label == part_id)] = total_part_num
if exceed:
break
# Make instance id continous.
ori_cur_labels = np.unique(instance_gt)
total_num_label = len(ori_cur_labels)
if instance_gt.max() + 1 != total_num_label:
for label in range(1, total_num_label):
instance_gt[instance_gt == ori_cur_labels[label]] = label
final_class_map = {}
for label in range(1, total_num_label):
if label >= 1:
final_class_map[label] = class_map[ori_cur_labels[label]]
return instance_gt, final_class_map
def compute_confidence(im_name, feature_map, class_map,
instance_label, output_dir,
panoptic_seg_mask, seg_score_list):
"""
"""
conf_file = open(os.path.join(output_dir, os.path.splitext(im_name)[0] + '.txt'), 'w')
weighted_map = np.zeros_like(feature_map[:, :, 0])
for index, score in enumerate(seg_score_list):
weighted_map += (panoptic_seg_mask == index + 1) * score
for label in class_map.keys():
cls = class_map[label]
confidence = feature_map[:, :, cls].reshape(-1)[np.where(instance_label.reshape(-1) == label)]
confidence = (weighted_map * feature_map[:, :, cls].copy()).reshape(-1)[
np.where(instance_label.reshape(-1) == label)]
confidence = confidence.sum() / len(confidence)
conf_file.write('{} {}\n'.format(cls, confidence))
conf_file.close()
def result_saving(fused_output, img_name, img_height, img_width, output_dir, mask_output_path, bbox_score, msrcnn_bbox):
if not os.path.exists(output_dir):
os.makedirs(output_dir)
global_root = os.path.join(output_dir, 'global_parsing')
instance_root = os.path.join(output_dir, 'instance_parsing')
tag_dir = os.path.join(output_dir, 'global_tag')
if not os.path.exists(global_root):
os.makedirs(global_root)
if not os.path.exists(instance_root):
os.makedirs(instance_root)
if not os.path.exists(tag_dir):
os.makedirs(tag_dir)
# For visualizing indexed png image.
palette = get_palette(256)
fused_output = cv2.resize(fused_output, dsize=(img_width, img_height), interpolation=cv2.INTER_LINEAR)
seg_pred = np.asarray(np.argmax(fused_output, axis=2), dtype=np.uint8)
masks = np.load(mask_output_path)
masks[np.where(seg_pred == 0)] = 0
panoptic_seg_mask = masks
seg_score_list = bbox_score
instance_pred, class_map = get_instance(seg_pred, panoptic_seg_mask)
refine(instance_pred, panoptic_seg_mask, seg_pred, class_map)
compute_confidence(img_name, fused_output, class_map, instance_pred, instance_root,
panoptic_seg_mask, seg_score_list)
ins_seg_results = open(os.path.join(tag_dir, os.path.splitext(img_name)[0] + '.txt'), "a")
keep_human_id_list = list(np.unique(panoptic_seg_mask))
if 0 in keep_human_id_list:
keep_human_id_list.remove(0)
for i in keep_human_id_list:
ins_seg_results.write('{:.6f} {} {} {} {}\n'.format(seg_score_list[i - 1],
int(msrcnn_bbox[i - 1][1]), int(msrcnn_bbox[i - 1][0]),
int(msrcnn_bbox[i - 1][3]), int(msrcnn_bbox[i - 1][2])))
ins_seg_results.close()
output_im_global = PILImage.fromarray(seg_pred)
output_im_instance = PILImage.fromarray(instance_pred)
output_im_tag = PILImage.fromarray(panoptic_seg_mask)
output_im_global.putpalette(palette)
output_im_instance.putpalette(palette)
output_im_tag.putpalette(palette)
output_im_global.save(os.path.join(global_root, os.path.splitext(img_name)[0] + '.png'))
output_im_instance.save(os.path.join(instance_root, os.path.splitext(img_name)[0] + '.png'))
output_im_tag.save(os.path.join(tag_dir, os.path.splitext(img_name)[0] + '.png'))
def multi_process(a, args):
img_name = a['im_name']
img_height = a['img_height']
img_width = a['img_width']
msrcnn_bbox = a['person_bbox']
bbox_score = a['person_bbox_score']
######### loading outputs from gloabl and local models #########
global_output = np.load(os.path.join(args.global_output_dir, os.path.splitext(img_name)[0] + '.npy'))
msrcnn_output = patch2img_output(args.msrcnn_output_dir, img_name, img_height, img_width, msrcnn_bbox,
bbox_type='msrcnn', num_class=20)
gt_output = patch2img_output(args.gt_output_dir, img_name, img_height, img_width, msrcnn_bbox, bbox_type='msrcnn',
num_class=20)
#### global and local branch logits fusion #####
# fused_output = global_output + msrcnn_output + gt_output
fused_output = global_output + gt_output
mask_output_path = os.path.join(args.mask_output_dir, os.path.splitext(img_name)[0] + '_mask.npy')
result_saving(fused_output, img_name, img_height, img_width, args.save_dir, mask_output_path, bbox_score, msrcnn_bbox)
return
def main(args):
json_file = open(args.test_json_path)
anno = json.load(json_file)['root']
results = joblib.Parallel(n_jobs=24, verbose=10, pre_dispatch="all")(
[joblib.delayed(multi_process)(a, args) for i, a in enumerate(anno)]
)
def get_arguments():
parser = argparse.ArgumentParser(description="obtain final prediction by logits fusion")
parser.add_argument("--test_json_path", type=str, default='./data/CIHP/cascade_152_finetune/test.json')
parser.add_argument("--global_output_dir", type=str,
default='./data/CIHP/global/global_result-cihp-resnet101/global_output')
# parser.add_argument("--msrcnn_output_dir", type=str,
# default='./data/CIHP/cascade_152__finetune/msrcnn_result-cihp-resnet101/msrcnn_output')
parser.add_argument("--gt_output_dir", type=str,
default='./data/CIHP/cascade_152__finetune/gt_result-cihp-resnet101/gt_output')
parser.add_argument("--mask_output_dir", type=str, default='./data/CIHP/cascade_152_finetune/mask')
parser.add_argument("--save_dir", type=str, default='./data/CIHP/fusion_results/cihp-msrcnn_finetune')
return parser.parse_args()
if __name__ == '__main__':
args = get_arguments()
main(args)