add inference scripts and tools
This commit is contained in:
yinwanqi
+17
-1
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
+1
-1
@@ -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
|
||||
|
||||
@@ -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()
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user