(self, img_dir=None, out_path=None)
| 22 | |
| 23 | class INFERENCE_demo(torch.utils.data.Dataset): |
| 24 | def __init__(self, img_dir=None, out_path=None): |
| 25 | |
| 26 | self.output_path = out_path |
| 27 | self.img_dir = img_dir |
| 28 | self.is_vid = False |
| 29 | body_model_cfg = dict( |
| 30 | type='smplx', |
| 31 | keypoint_src='smplx', |
| 32 | num_expression_coeffs=10, |
| 33 | num_betas=10, |
| 34 | gender='neutral', |
| 35 | keypoint_dst='smplx_137', |
| 36 | model_path='data/body_models/smplx', |
| 37 | use_pca=False, |
| 38 | use_face_contour=True, |
| 39 | batch_size = cfg.batch_size) |
| 40 | self.body_model = build_body_model(body_model_cfg).to('cuda') |
| 41 | |
| 42 | rank, _ = get_dist_info() |
| 43 | if self.img_dir.endswith('.mp4'): |
| 44 | self.is_vid = True |
| 45 | self.img_name = self.img_dir.split('/')[-1][:-4] |
| 46 | # self.img_dir = self.img_dir[:-4] |
| 47 | else: |
| 48 | self.img_name = self.img_dir.split('/')[-1] |
| 49 | |
| 50 | self.output_path = os.path.join(self.output_path, self.img_name) |
| 51 | os.makedirs(self.output_path, exist_ok=True) |
| 52 | self.mesh_path = os.path.join(self.output_path, 'mesh') |
| 53 | os.makedirs(self.mesh_path, exist_ok=True) |
| 54 | self.tmp_dir = os.path.join(self.output_path, 'temp_img') |
| 55 | os.makedirs(self.tmp_dir, exist_ok=True) |
| 56 | self.result_img_dir = os.path.join(self.output_path, 'res_img') |
| 57 | |
| 58 | if not self.is_vid: |
| 59 | if rank == 0: |
| 60 | image_files = sorted(glob(self.img_dir + '/*.jpg') + glob(self.img_dir + '/*.png')) |
| 61 | for i, image_file in enumerate(image_files): |
| 62 | new_name = os.path.join(self.tmp_dir, '%06d.png'%i) |
| 63 | shutil.copy(image_file, new_name) |
| 64 | dist.barrier() |
| 65 | else: |
| 66 | if rank == 0: |
| 67 | video_to_images(self.img_dir, self.tmp_dir) |
| 68 | dist.barrier() |
| 69 | |
| 70 | self.img_paths = sorted(glob(self.tmp_dir +'/*',recursive=True)) |
| 71 | |
| 72 | self.num_person = cfg.num_person if 'num_person' in cfg else 0.1 |
| 73 | self.score_threshold = cfg.threshold if 'threshold' in cfg else 0.1 |
| 74 | self.format = DefaultFormatBundle() |
| 75 | self.normalize = Normalize(mean=[123.675, 116.28, 103.53], std=[58.395, 57.12, 57.375]) |
| 76 | |
| 77 | def __len__(self): |
| 78 | return len(self.img_paths) |
nothing calls this directly
no test coverage detected