def __init__(self,
img_folder,
meta_dir,
random_bg_dir,
cam_name,
inp_res=224,
train=True,
training_set_percentage=0.9):
self.actor_name = 'RobotArmActor_3'
self.img_folder = img_folder # root image folders
self.meta_dir = meta_dir
self.is_train = train # training set or test set
self.inp_res = inp_res
self.cam_name = cam_name
self.dataset = vdb.Dataset(img_folder)
ids = self.dataset.get_ids
self.color = self.dataset.get_annotation_color()[self.actor_name]
split = round(
len(ids) *
training_set_percentage) # 90% for training, 10% for validation
# create train/val split
self.train = ids[:split]
self.valid = ids[split:]
self.mean, self.std = self._compute_mean()
def __init__(self,
img_folder,
meta_dir,
random_bg_dir,
cam_name,
anno_type,
inp_res=256,
out_res=64,
train=True,
sigma=1,
training_set_percentage=0.9,
label_type='Gaussian',
scales=[0.7, 0.8, 0.9, 1, 1.2, 1.4, 1.6],
multi_scale=False,
ignore_invis_pts=False,
replace_bg=False):
self.scales = scales
self.actor_name = "RobotArmActor_3"
self.img_folder = img_folder # root image folders
self.meta_dir = meta_dir
self.is_train = train # training set or test set
self.inp_res = inp_res
self.out_res = out_res
self.sigma = sigma
self.label_type = label_type
self.cam_name = cam_name
self.replace_bg = replace_bg
anno_type = anno_type.lower()
assert anno_type == '3d' or anno_type == '2d' or anno_type == 'none'
self.anno_type = anno_type
if self.anno_type == '3d':
self.dataset = vdb.Dataset(img_folder)
ids = self.dataset.get_ids()
self.color = self.dataset.get_annotation_color()[self.actor_name]
split = round(
len(ids) *
training_set_percentage) #90% for training, 10% for validation
self.train = ids[:split]
self.valid = ids[split:]
with open(os.path.join(meta_dir, 'lr_pairs.json'), 'r') as f:
self.lr_pairs = json.load(f)
if self.anno_type == '2d':
if os.path.isfile(os.path.join(img_folder, 'valid_img_list.json')):
with open(os.path.join(img_folder,
'valid_img_list.json')) as f:
self.dataset = json.load(f)
else:
_, self.dataset = read_jpg(os.path.join(img_folder, 'imgs'))
ids = list(range(len(self.dataset)))
with open(os.path.join(img_folder, 'pts.json')) as f:
self.bbox_anno = json.load(f)
split = round(
len(ids) *
training_set_percentage) #90% for training, 10% for validation
self.train = ids[:split]
self.valid = ids[split:]
with open(os.path.join(meta_dir, 'lr_pairs.json'), 'r') as f:
self.lr_pairs = json.load(f)
if self.anno_type == 'none':
L, F = read_jpg(img_folder)
ids = range(1, len(L) + 1)
split = round(len(ids) * training_set_percentage)
self.train = ids[:split]
self.valid = ids[split:]
# self.train = []
# self.valid = ids
self.dataset = L
self.F = F
self.anno = None
if isfile(os.path.join(img_folder, 'pts.json')):
self.anno = json.load(
open(os.path.join(img_folder, 'pts.json')))
self.multi_scale = multi_scale
self.ignore_invis_pts = ignore_invis_pts
if self.replace_bg:
self.background_replace = background_replace(random_bg_dir)
# create train/val split
self.mean, self.std = self._compute_mean()