diff --git a/inf_demo.py b/inf_demo.py new file mode 100644 index 0000000..c4f81a9 --- /dev/null +++ b/inf_demo.py @@ -0,0 +1,67 @@ +import os +import argparse +import cv2 +import readline +import tqdm +import time + +import torch +from torch.utils.data import DataLoader +from PIL import Image + +# from smplerx.main.inference import Inferer +from main.inference import SmplerxData, Inferer + + +import pdb + + +def inference(args): + + # load model + num_gpus = 1 if torch.cuda.is_available() else -1 + inferer = Inferer(args.pretrained_model, num_gpus) + + # test annotations + annotations = [ + {'image_path': '/home/weichen/wc_workspace/laoyouji/frames/Season_1/S01E01/frame025000.jpg', 'bbox': [0, 0, 1000, 1000]}, + {'image_path': '/home/weichen/wc_workspace/laoyouji/frames/Season_1/S01E01/frame026000.jpg', 'bbox': [0, 0, 1000, 1000]} + ] + + annotations = annotations*500 + anno_len = len(annotations) + + start_time = time.time() + batch_size = 1 + + dataset = SmplerxData(annotations=annotations) + dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=4) + + preproces_time = time.time() + + for batch in tqdm.tqdm(dataloader): + + smplx_pred, meta, mesh = inferer.batch_infer_given_bbox(batch['image'], batch['bbox']) + + end_time = time.time() + + # print report, time in seconds + print(f'Instance number: {anno_len}, Batch size: {batch_size}') + print(f'Preprocess time: {preproces_time-start_time:02f}, FPS: {anno_len/(preproces_time-start_time):02f}') + print(f'Inference time: {end_time-preproces_time:02f}, FPS: {anno_len/(end_time-preproces_time):02f}') + + +if __name__ == "__main__": + + parser = argparse.ArgumentParser() + + parser.add_argument('--in_threshold', type=float, default=0.5) + parser.add_argument('--pretrained_model', type=str, default='smpler_x_h32') + + args = parser.parse_args() + + inference(args) + + + + diff --git a/main/inference.py b/main/inference.py index 730cbcb..28276e9 100644 --- a/main/inference.py +++ b/main/inference.py @@ -18,9 +18,9 @@ CUR_DIR = osp.dirname(os.path.abspath(__file__)) sys.path.insert(0, osp.join(CUR_DIR, '..', 'main')) sys.path.insert(0, osp.join(CUR_DIR , '..', 'common')) from config import cfg, model_path_dict -from ..common.base import Demoer -from ..common.utils.preprocessing import process_bbox, generate_patch_image -from ..common.utils.human_models import smpl_x +from base import Demoer +from utils.preprocessing import process_bbox, generate_patch_image +from utils.human_models import smpl_x class SmplerxData(Dataset): def __init__(self, annotations): @@ -39,7 +39,6 @@ class SmplerxData(Dataset): # prepare input image transform = transforms.ToTensor() - vis_img = image.copy() original_img_height, original_img_width = image.shape[:2] bbox = process_bbox(bbox, original_img_width, original_img_height) img, img2bb_trans, bb2img_trans = generate_patch_image(image, bbox, 1.0, 0.0, False, (512, 384))