AiOS / datasets /SynBody.py
ttxskk
update
d7e58f0
raw
history blame
2.11 kB
import os
import os.path as osp
import numpy as np
import torch
import cv2
import json
import copy
from pycocotools.coco import COCO
from config.config import cfg
from util.human_models import smpl_x
from util.preprocessing import (
load_img, process_bbox, augmentation_instance_sample
,process_human_model_output_batch_simplify,process_db_coord_batch_no_valid)
from util.transforms import world2cam, cam2pixel, rigid_align
from humandata import HumanDataset
class SynBody(HumanDataset):
def __init__(self, transform, data_split):
super(SynBody, self).__init__(transform, data_split)
self.img_dir = 'data/datasets/synbody'
self.annot_path = 'data/preprocessed_npz/multihuman_data/synbody_v1.1_multi_new.npz'
self.annot_path_cache = 'data/preprocessed_npz/cache/synbody_v1.1_cache_new_10.npz'
self.use_cache = getattr(cfg, 'use_cache', False)
self.img_shape = (720, 1280) # (h, w)
self.cam_param = {
'focal': (540, 540), # (fx, fy)
'princpt': (640, 360) # (cx, cy)
}
# check image shape
img_path = osp.join(self.img_dir,
np.load(self.annot_path)['image_path'][0])
img_shape = cv2.imread(img_path).shape[:2]
assert self.img_shape == img_shape, 'image shape is incorrect: {} vs {}'.format(
self.img_shape, img_shape)
# load data or cache
if self.use_cache and osp.isfile(self.annot_path_cache):
print(
f'[{self.__class__.__name__}] loading cache from {self.annot_path_cache}'
)
self.datalist = self.load_cache(self.annot_path_cache)
else:
if self.use_cache:
print(
f'[{self.__class__.__name__}] Cache not found, generating cache...'
)
self.datalist = self.load_data(train_sample_interval=getattr(
cfg, f'{self.__class__.__name__}_train_sample_interval', 15))
if self.use_cache:
self.save_cache(self.annot_path_cache, self.datalist)