Spaces:
Runtime error
Runtime error
fifa-tryon-demo
/
Self-Correction-Human-Parsing-for-ACGPN
/mhp_extension
/make_crop_and_mask_w_mask_nms.py
| import numpy as np | |
| import cv2, torch | |
| import os | |
| import json | |
| import argparse | |
| import pycocotools.mask as mask_util | |
| from tqdm import tqdm | |
| def bbox_expand(img_height, img_width, bbox, exp_ratio): | |
| x_min, y_min, x_max, y_max = bbox[:] | |
| exp_x = (x_max - x_min) * ((exp_ratio - 1) / 2) | |
| exp_y = (y_max - y_min) * ((exp_ratio - 1) / 2) | |
| new_x_min = 0 if x_min - exp_x < 0 else np.round(x_min - exp_x) | |
| new_y_min = 0 if y_min - exp_y < 0 else np.round(y_min - exp_y) | |
| new_x_max = img_width - 1 if x_max + exp_x > img_width - 1 else np.round(x_max + exp_x) | |
| new_y_max = img_height - 1 if y_max + exp_y > img_height - 1 else np.round(y_max + exp_y) | |
| return int(new_x_min), int(new_y_min), int(new_x_max), int(new_y_max) | |
| def make_crop_and_mask(img_info, pred, file_list, crop_save_dir, mask_save_dir, args): | |
| img_name = img_info['file_name'] | |
| img_id = img_info['id'] - 1 # img_info['id'] start form 1 | |
| img_w = img_info['width'] | |
| img_h = img_info['height'] | |
| img = cv2.imread(os.path.join(args.img_dir, img_name)) | |
| exp_bbox = [] | |
| ori_bbox = [] | |
| bbox_name_list = [] | |
| bbox_score_list = [] | |
| person_idx = 0 | |
| panoptic_seg = np.zeros((img_h, img_w), dtype=np.uint8) | |
| assert len(pred[img_id]['instances']) > 0, 'image without instance prediction' | |
| for instance in pred[img_id]['instances']: | |
| score = instance['score'] | |
| if score < args.conf_thres: | |
| break | |
| mask = mask_util.decode(instance['segmentation']) | |
| mask_area = mask.sum() | |
| if mask_area == 0: # if mask_area < img_w*img_h/1000: | |
| continue | |
| intersect = (mask > 0) & (panoptic_seg > 0) | |
| intersect_area = intersect.sum() | |
| if intersect_area * 1.0 / mask_area > args.overlap_threshold: # todo add args | |
| continue | |
| if intersect_area > 0: | |
| mask = mask & (panoptic_seg == 0) | |
| person_idx += 1 | |
| panoptic_seg = np.where(mask == 0, panoptic_seg, person_idx) | |
| bbox_score_list.append(score) | |
| ins_bbox = instance['bbox'] # [x,y,w,h] format | |
| x_min, y_min, box_w, box_h = ins_bbox | |
| x_max, y_max = x_min + box_w, y_min + box_h | |
| exp_x_min, exp_y_min, exp_x_max, exp_y_max = bbox_expand(img_h, img_w, [x_min, y_min, x_max, y_max], | |
| args.exp_ratio) | |
| crop_img = img[exp_y_min:exp_y_max + 1, exp_x_min:exp_x_max + 1, :] | |
| exp_bbox.append([exp_x_min, exp_y_min, exp_x_max, exp_y_max]) | |
| ori_bbox.append([x_min, y_min, x_max, y_max]) | |
| bbox_name = os.path.splitext(img_name)[0] + '_' + str(person_idx) + '_msrcnn.jpg' | |
| bbox_name_list.append(bbox_name) | |
| cv2.imwrite(os.path.join(crop_save_dir, bbox_name), crop_img) | |
| assert person_idx > 0, 'image without instance' | |
| mask_name = os.path.splitext(img_name)[0] + '_mask.npy' | |
| np.save(os.path.join(mask_save_dir, mask_name), panoptic_seg) | |
| ############## json writing ################## | |
| item = {} | |
| item['dataset'] = 'CIHP' | |
| item['im_name'] = img_name | |
| item['img_height'] = img_h | |
| item['img_width'] = img_w | |
| item['center'] = [img_h / 2, img_w / 2] | |
| item['person_num'] = person_idx | |
| item['person_bbox'] = exp_bbox | |
| item['real_person_bbox'] = ori_bbox | |
| item['person_bbox_score'] = bbox_score_list | |
| item['bbox_name'] = bbox_name_list | |
| item['mask_name'] = mask_name | |
| file_list.append(item) | |
| json_file = {'root': file_list} | |
| return json_file, file_list | |
| def get_arguments(): | |
| parser = argparse.ArgumentParser(description="crop person val/test demo for inference") | |
| parser.add_argument("--exp_ratio", type=float, default=1.2) | |
| parser.add_argument("--overlap_threshold", type=float, default=0.5) | |
| parser.add_argument("--conf_thres", type=float, default=0.5) | |
| parser.add_argument("--img_dir", type=str, | |
| default='/data03/v_xuyunqiu/data/instance-level_human_parsing/Testing/Images') | |
| parser.add_argument("--save_dir", type=str, | |
| default='/data03/v_xuyunqiu/Projects/experiment_data/testing/resnest_200_TTA_mask_nms_all_data') | |
| parser.add_argument("--img_list", type=str, | |
| default='/data03/v_xuyunqiu/Projects/pycococreator/annotations/CIHP_test.json') | |
| parser.add_argument("--det_res", type=str, | |
| default='/data02/v_xuyunqiu/detectron2-ResNeSt/tools/output_cihp_inference_resnest/inference_TTA/instances_predictions.pth') | |
| return parser.parse_args() | |
| def main(args): | |
| img_info_list = json.load(open(args.img_list, encoding='UTF-8')) | |
| pred = torch.load(args.det_res) | |
| crop_save_dir = os.path.join(args.save_dir, 'crop_pic') | |
| if not os.path.exists(crop_save_dir): | |
| os.makedirs(crop_save_dir) | |
| mask_save_dir = os.path.join(args.save_dir, 'crop_mask') | |
| if not os.path.exists(mask_save_dir): | |
| os.makedirs(mask_save_dir) | |
| file_list = [] | |
| for img_info in tqdm(img_info_list['images']): | |
| json_file, file_list = make_crop_and_mask(img_info, pred, file_list, crop_save_dir, mask_save_dir, args) | |
| with open(os.path.join(args.save_dir, 'crop.json'), 'w') as f: | |
| json.dump(json_file, f, indent=2) | |
| if __name__ == '__main__': | |
| args = get_arguments() | |
| main(args) | |