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)