From 07f7a3720b978ea4bd656b5318ce19655b607b9a Mon Sep 17 00:00:00 2001 From: yinwanqi Date: Thu, 20 Jul 2023 17:21:31 +0800 Subject: [PATCH] add inference scripts and tools --- common/base.py | 18 ++- common/utils/inference_utils.py | 153 ++++++++++++++++++++++++++ main/config.py | 2 +- main/inference.py | 189 ++++++++++++++++++++++++++++++++ main/slurm_inference.sh | 58 ++++++++++ 5 files changed, 418 insertions(+), 2 deletions(-) create mode 100644 common/utils/inference_utils.py create mode 100644 main/inference.py create mode 100644 main/slurm_inference.sh diff --git a/common/base.py b/common/base.py index 82b394d..9c1f157 100644 --- a/common/base.py +++ b/common/base.py @@ -317,6 +317,17 @@ class Demoer(Base): self.test_epoch = int(test_epoch) super(Demoer, self).__init__(log_name='test_logs.txt') + def _make_batch_generator(self, demo_scene): + # data load and construct batch generator + self.logger.info("Creating dataset...") + from data.UBody.UBody import UBody + testset_loader = UBody(transforms.ToTensor(), "demo", demo_scene) # eval(demoset)(transforms.ToTensor(), "demo") + batch_generator = DataLoader(dataset=testset_loader, batch_size=cfg.num_gpus * cfg.test_batch_size, + shuffle=False, num_workers=cfg.num_thread, pin_memory=True) + + self.testset = testset_loader + self.batch_generator = batch_generator + def _make_model(self): self.logger.info('Load checkpoint from {}'.format(cfg.pretrained_model_path)) @@ -324,7 +335,7 @@ class Demoer(Base): self.logger.info("Creating graph...") model = get_model('test') model = DataParallel(model).cuda() - ckpt = torch.load(cfg.pretrained_model_path, map_location=torch.device('cpu')) + ckpt = torch.load(cfg.pretrained_model_path) from collections import OrderedDict new_state_dict = OrderedDict() @@ -338,3 +349,8 @@ class Demoer(Base): model.eval() self.model = model + + def _evaluate(self, outs, cur_sample_idx): + eval_result = self.testset.evaluate(outs, cur_sample_idx) + return eval_result + diff --git a/common/utils/inference_utils.py b/common/utils/inference_utils.py new file mode 100644 index 0000000..3ccc515 --- /dev/null +++ b/common/utils/inference_utils.py @@ -0,0 +1,153 @@ +from typing import Literal, Union + +def process_mmdet_results(mmdet_results: list, + cat_id: int = 0, + multi_person: bool = True) -> list: + """Process mmdet results, sort bboxes by area in descending order. + + Args: + mmdet_results (list): + Result of mmdet.apis.inference_detector + when the input is a batch. + Shape of the nested lists is + (n_frame, n_category, n_human, 5). + cat_id (int, optional): + Category ID. This function will only select + the selected category, and drop the others. + Defaults to 0, ID of human category. + multi_person (bool, optional): + Whether to allow multi-person detection, which is + slower than single-person. If false, the function + only assure that the first person of each frame + has the biggest bbox. + Defaults to True. + + Returns: + list: + A list of detected bounding boxes. + Shape of the nested lists is + (n_frame, n_human, 5) + and each bbox is (x, y, x, y, score). + """ + ret_list = [] + only_max_arg = not multi_person + # for _, frame_results in enumerate(mmdet_results): + cat_bboxes = mmdet_results[cat_id] + # import pdb; pdb.set_trace() + sorted_bbox = qsort_bbox_list(cat_bboxes, only_max_arg) + + if only_max_arg: + ret_list.append(sorted_bbox[0:1]) + else: + ret_list.append(sorted_bbox) + return ret_list + + +def qsort_bbox_list(bbox_list: list, + only_max: bool = False, + bbox_convention: Literal['xyxy', 'xywh'] = 'xyxy'): + """Sort a list of bboxes, by their area in pixel(W*H). + + Args: + input_list (list): + A list of bboxes. Each item is a list of (x1, y1, x2, y2) + only_max (bool, optional): + If True, only assure the max element at first place, + others may not be well sorted. + If False, return a well sorted descending list. + Defaults to False. + bbox_convention (str, optional): + Bbox type, xyxy or xywh. Defaults to 'xyxy'. + + Returns: + list: + A sorted(maybe not so well) descending list. + """ + # import pdb; pdb.set_trace() + if len(bbox_list) <= 1: + return bbox_list + else: + bigger_list = [] + less_list = [] + anchor_index = int(len(bbox_list) / 2) + anchor_bbox = bbox_list[anchor_index] + anchor_area = get_area_of_bbox(anchor_bbox, bbox_convention) + for i in range(len(bbox_list)): + if i == anchor_index: + continue + tmp_bbox = bbox_list[i] + tmp_area = get_area_of_bbox(tmp_bbox, bbox_convention) + if tmp_area >= anchor_area: + bigger_list.append(tmp_bbox) + else: + less_list.append(tmp_bbox) + if only_max: + return qsort_bbox_list(bigger_list) + \ + [anchor_bbox, ] + less_list + else: + return qsort_bbox_list(bigger_list) + \ + [anchor_bbox, ] + qsort_bbox_list(less_list) + +def get_area_of_bbox( + bbox: Union[list, tuple], + bbox_convention: Literal['xyxy', 'xywh'] = 'xyxy') -> float: + """Get the area of a bbox_xyxy. + + Args: + (Union[list, tuple]): + A list of [x1, y1, x2, y2]. + bbox_convention (str, optional): + Bbox type, xyxy or xywh. Defaults to 'xyxy'. + + Returns: + float: + Area of the bbox(|y2-y1|*|x2-x1|). + """ + # import pdb;pdb.set_trace() + if bbox_convention == 'xyxy': + return abs(bbox[2] - bbox[0]) * abs(bbox[3] - bbox[1]) + elif bbox_convention == 'xywh': + return abs(bbox[2] * bbox[3]) + else: + raise TypeError(f'Wrong bbox convention: {bbox_convention}') + +def calculate_iou(bbox1, bbox2): + # Calculate the Intersection over Union (IoU) between two bounding boxes + x1 = max(bbox1[0], bbox2[0]) + y1 = max(bbox1[1], bbox2[1]) + x2 = min(bbox1[2], bbox2[2]) + y2 = min(bbox1[3], bbox2[3]) + + intersection_area = max(0, x2 - x1 + 1) * max(0, y2 - y1 + 1) + + bbox1_area = (bbox1[2] - bbox1[0] + 1) * (bbox1[3] - bbox1[1] + 1) + bbox2_area = (bbox2[2] - bbox2[0] + 1) * (bbox2[3] - bbox2[1] + 1) + + union_area = bbox1_area + bbox2_area - intersection_area + + iou = intersection_area / union_area + return iou + + +def non_max_suppression(bboxes, iou_threshold): + # Sort the bounding boxes by their confidence scores (e.g., the probability of containing an object) + bboxes = sorted(bboxes, key=lambda x: x[4], reverse=True) + + # Initialize a list to store the selected bounding boxes + selected_bboxes = [] + + # Perform non-maximum suppression + while len(bboxes) > 0: + current_bbox = bboxes[0] + selected_bboxes.append(current_bbox) + bboxes = bboxes[1:] + + remaining_bboxes = [] + for bbox in bboxes: + iou = calculate_iou(current_bbox, bbox) + if iou < iou_threshold: + remaining_bboxes.append(bbox) + + bboxes = remaining_bboxes + + return selected_bboxes \ No newline at end of file diff --git a/main/config.py b/main/config.py index c696fda..29b2228 100644 --- a/main/config.py +++ b/main/config.py @@ -50,7 +50,7 @@ class Config: os.system(f'cp -r {self.root_dir}/{file} {self.code_dir}') def update_test_config(self, testset, agora_benchmark, shapy_eval_split, pretrained_model_path, use_cache, - eval_on_train, vis): + eval_on_train=False, vis=False): self.testset = testset self.agora_benchmark = agora_benchmark self.pretrained_model_path = pretrained_model_path diff --git a/main/inference.py b/main/inference.py new file mode 100644 index 0000000..3c68d1e --- /dev/null +++ b/main/inference.py @@ -0,0 +1,189 @@ +import os +import sys +import os.path as osp +import argparse +import numpy as np +import torchvision.transforms as transforms +import torch.backends.cudnn as cudnn +import torch +sys.path.insert(0, osp.join('..', 'main')) +sys.path.insert(0, osp.join('..', 'data')) +sys.path.insert(0, osp.join('..', 'common')) +from config import cfg +import cv2 +from tqdm import tqdm +import json +from typing import Literal, Union +from mmdet.apis import init_detector, inference_detector +from utils.inference_utils import process_mmdet_results, non_max_suppression + +def parse_args(): + parser = argparse.ArgumentParser() + parser.add_argument('--num_gpus', type=int, dest='num_gpus') + parser.add_argument('--exp_name', type=str, default='output/test') + parser.add_argument('--result_path', type=str, default='output/test') + parser.add_argument('--ckpt_idx', type=int, default=0) + parser.add_argument('--testset', type=str, default='EHF') + parser.add_argument('--agora_benchmark', type=str, default='na') + parser.add_argument('--img_path', type=str, default='input.png') + parser.add_argument('--start', type=str, default=1) + parser.add_argument('--end', type=str, default=1) + parser.add_argument('--output_folder', type=str, default='output') + parser.add_argument('--demo_dataset', type=str, default='na') + parser.add_argument('--demo_scene', type=str, default='all') + parser.add_argument('--show_verts', action="store_true") + parser.add_argument('--show_bbox', action="store_true") + parser.add_argument('--save_mesh', action="store_true") + parser.add_argument('--multi_person', action="store_true") + parser.add_argument('--iou_thr', type=float, default=0.5) + parser.add_argument('--bbox_thr', type=int, default=50) + args = parser.parse_args() + return args + +def main(): + + args = parse_args() + config_path = osp.join('../output',args.result_path, 'code', 'config_base.py') + ckpt_path = osp.join('../output', args.result_path, 'model_dump', f'snapshot_{int(args.ckpt_idx)}.pth.tar') + + cfg.get_config_fromfile(config_path) + cfg.update_test_config(args.testset, args.agora_benchmark, shapy_eval_split=None, + pretrained_model_path=ckpt_path, use_cache=False) + cfg.update_config(args.num_gpus, args.exp_name) + cudnn.benchmark = True + + # load model + from base import Demoer + from utils.preprocessing import load_img, process_bbox, generate_patch_image + from utils.vis import render_mesh, save_obj + from utils.human_models import smpl_x + demoer = Demoer() + demoer._make_model() + demoer.model.eval() + + start = int(args.start) + end = start + int(args.end) + multi_person = args.multi_person + + + ### mmdet init + checkpoint_file = '../pretrained_models/mmdet/faster_rcnn_r50_fpn_1x_coco_20200130-047c8118.pth' + config_file= '../pretrained_models/mmdet/mmdet_faster_rcnn_r50_fpn_coco.py' + model = init_detector(config_file, checkpoint_file, device='cuda:0') # or device='cuda:0' + + for frame in tqdm(range(start, end)): + img_path = os.path.join(args.img_path, f'{int(frame):06d}.jpg') + + # prepare input image + transform = transforms.ToTensor() + original_img = load_img(img_path) + vis_img = original_img.copy() + original_img_height, original_img_width = original_img.shape[:2] + os.makedirs(args.output_folder, exist_ok=True) + + ## mmdet inference + mmdet_results = inference_detector(model, img_path) + mmdet_box = process_mmdet_results(mmdet_results, cat_id=0, multi_person=True) + + # save original image if no bbox + if len(mmdet_box[0])<1: + # save rendered image + frame_name = img_path.split('/')[-1] + save_path_img = os.path.join(args.output_folder, 'img') + os.makedirs(save_path_img, exist_ok= True) + cv2.imwrite(os.path.join(save_path_img, f'{frame_name}'), vis_img[:, :, ::-1]) + continue + + if not multi_person: + # only select the largest bbox + num_bbox = 1 + mmdet_box = mmdet_box[0] + else: + # keep bbox by NMS with iou_thr + mmdet_box = non_max_suppression(mmdet_box[0], args.iou_thr) + num_bbox = len(mmdet_box) + + ## loop all detected bboxes + for bbox_id in range(num_bbox): + mmdet_box_xywh = np.zeros((4)) + mmdet_box_xywh[0] = mmdet_box[bbox_id][0] + mmdet_box_xywh[1] = mmdet_box[bbox_id][1] + mmdet_box_xywh[2] = abs(mmdet_box[bbox_id][2]-mmdet_box[bbox_id][0]) + mmdet_box_xywh[3] = abs(mmdet_box[bbox_id][3]-mmdet_box[bbox_id][1]) + + # skip small bboxes by bbox_thr in pixel + if mmdet_box_xywh[2] < args.bbox_thr or mmdet_box_xywh[3] < args.bbox_thr * 3: + continue + + # for bbox visualization + start_point = (int(mmdet_box[bbox_id][0]), int(mmdet_box[bbox_id][1])) + end_point = (int(mmdet_box[bbox_id][2]), int(mmdet_box[bbox_id][3])) + + bbox = process_bbox(mmdet_box_xywh, original_img_width, original_img_height) + img, img2bb_trans, bb2img_trans = generate_patch_image(original_img, bbox, 1.0, 0.0, False, cfg.input_img_shape) + img = transform(img.astype(np.float32))/255 + img = img.cuda()[None,:,:,:] + inputs = {'img': img} + targets = {} + meta_info = {} + + # mesh recovery + with torch.no_grad(): + out = demoer.model(inputs, targets, meta_info, 'test') + mesh = out['smplx_mesh_cam'].detach().cpu().numpy()[0] + + ## save mesh + if args.save_mesh: + save_path_mesh = os.path.join(args.output_folder, 'mesh') + os.makedirs(save_path_mesh, exist_ok= True) + save_obj(mesh, smpl_x.face, os.path.join(save_path_mesh, f'{frame:05}_{bbox_id}.obj')) + + ## save single person param + smplx_pred = {} + smplx_pred['global_orient'] = out['smplx_root_pose'].reshape(-1,3).cpu().numpy() + smplx_pred['body_pose'] = out['smplx_body_pose'].reshape(-1,3).cpu().numpy() + smplx_pred['left_hand_pose'] = out['smplx_lhand_pose'].reshape(-1,3).cpu().numpy() + smplx_pred['right_hand_pose'] = out['smplx_rhand_pose'].reshape(-1,3).cpu().numpy() + smplx_pred['jaw_pose'] = out['smplx_jaw_pose'].reshape(-1,3).cpu().numpy() + smplx_pred['leye_pose'] = np.zeros((1, 3)) + smplx_pred['reye_pose'] = np.zeros((1, 3)) + smplx_pred['betas'] = out['smplx_shape'].reshape(-1,10).cpu().numpy() + smplx_pred['expression'] = out['smplx_expr'].reshape(-1,10).cpu().numpy() + smplx_pred['transl'] = out['cam_trans'].reshape(-1,3).cpu().numpy() + save_path_smplx = os.path.join(args.output_folder, 'smplx') + os.makedirs(save_path_smplx, exist_ok= True) + + npz_path = os.path.join(save_path_smplx, f'{frame:05}_{bbox_id}.npz') + np.savez(npz_path, **smplx_pred) + + ## render single person mesh + focal = [cfg.focal[0] / cfg.input_body_shape[1] * bbox[2], cfg.focal[1] / cfg.input_body_shape[0] * bbox[3]] + princpt = [cfg.princpt[0] / cfg.input_body_shape[1] * bbox[2] + bbox[0], cfg.princpt[1] / cfg.input_body_shape[0] * bbox[3] + bbox[1]] + vis_img = render_mesh(vis_img, mesh, smpl_x.face, {'focal': focal, 'princpt': princpt}, + mesh_as_vertices=args.show_verts) + if args.show_bbox: + vis_img = cv2.rectangle(vis_img, start_point, end_point, (255, 0, 0), 2) + + ## save single person meta + meta = {'focal': focal, + 'princpt': princpt, + 'bbox': bbox.tolist(), + 'bbox_mmdet': mmdet_box_xywh.tolist(), + 'bbox_id': bbox_id, + 'img_path': img_path} + json_object = json.dumps(meta, indent=4) + + save_path_meta = os.path.join(args.output_folder, 'meta') + os.makedirs(save_path_meta, exist_ok= True) + with open(os.path.join(save_path_meta, f'{frame:05}_{bbox_id}.json'), "w") as outfile: + outfile.write(json_object) + + ## save rendered image with all person + frame_name = img_path.split('/')[-1] + save_path_img = os.path.join(args.output_folder, 'img') + os.makedirs(save_path_img, exist_ok= True) + cv2.imwrite(os.path.join(save_path_img, f'{frame_name}'), vis_img[:, :, ::-1]) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/main/slurm_inference.sh b/main/slurm_inference.sh new file mode 100644 index 0000000..a010ba6 --- /dev/null +++ b/main/slurm_inference.sh @@ -0,0 +1,58 @@ +#!/usr/bin/env bash +set -x + +PARTITION=Zoetrope + +INPUT_VIDEO=$1 +APPENDIX=$2 +FPS=$3 +RES_PATH=$4 +CKPT=$5 +GPUS=1 +JOB_NAME=inference_${INPUT_VIDEO} + +GPUS_PER_NODE=$((${GPUS}<8?${GPUS}:8)) +CPUS_PER_TASK=4 # ${CPUS_PER_TASK:-2} +SRUN_ARGS=${SRUN_ARGS:-""} + +IMG_PATH=../demo/images/${INPUT_VIDEO} +SAVE_DIR=../demo/results/${INPUT_VIDEO} + +# video to images +mkdir $IMG_PATH +mkdir $SAVE_DIR +ffmpeg -i ../demo/videos/${INPUT_VIDEO}.${APPENDIX} -f image2 -vf fps=${FPS}/1 -qscale 0 ../demo/images/${INPUT_VIDEO}/%06d.jpg + +end_count=$(find "$IMG_PATH" -type f | wc -l) +echo $end_count + +# inference +PYTHONPATH="$(dirname $0)/..":$PYTHONPATH \ +srun -p ${PARTITION} \ + --job-name=${JOB_NAME} \ + --gres=gpu:${GPUS_PER_NODE} \ + --ntasks=${GPUS} \ + --ntasks-per-node=${GPUS_PER_NODE} \ + --cpus-per-task=${CPUS_PER_TASK} \ + --kill-on-bad-exit=1 \ + ${SRUN_ARGS} \ + python inference.py \ + --num_gpus ${GPUS_PER_NODE} \ + --exp_name output/demo_${JOB_NAME} \ + --result_path ${RES_PATH} \ + --ckpt_idx ${CKPT} \ + --agora_benchmark agora_model \ + --img_path ${IMG_PATH} \ + --start 1 \ + --end $end_count \ + --output_folder ${SAVE_DIR} \ + --show_verts \ + --show_bbox \ + --save_mesh \ + # --multi_person \ + # --iou_thr 0.2 \ + # --bbox_thr 20 + + +# images to video +ffmpeg -y -f image2 -r ${FPS} -i ${SAVE_DIR}/img/%06d.jpg -vcodec mjpeg -qscale 0 -pix_fmt yuv420p ../demo/results/${INPUT_VIDEO}.mp4