def __init__(self, root, imageset_path, syn_data_path=None, use_real_img = True, syn_range=70000, num_syn_images=70000 ,target_h=76, target_w=76

, bg_path = None, kp_path="data/YCB-Video/YCB_bbox.npy", data_cfg="data/data-YCB.cfg",

use_bg_img = True, one_syn_per_batch = False, batch_size = None):

self.root = root

data_options = read_data_cfg(data_cfg)

self.input_width = int(data_options['width']) # 608, width of CNN input

self.input_height = int(data_options['height'])

self.original_width = 640 # width of original img

self.original_height = 480

self.target_h = target_h # after network

self.target_w = target_w

self.num_classes = int(data_options['classes'])

self.train_paths = []

self.gen_train_list(imageset_path)

self.use_real_img = use_real_img

self.syn_data_path = syn_data_path

self.syn_range = syn_range # YCB has 80000 syn images in total, continuously indexed. We keep a few images for testing.

self.syn_bg_image_paths = get_img_list_from(bg_path) if bg_path is not None else []

self.use_bg_img = use_bg_img

self.num_syn_images = num_syn_images # syn images for training

self.weight_cross_entropy = None

self.set_balancing_weight()

self.object_names_ycbvideo = ['002_master_chef_can', '003_cracker_box', '004_sugar_box', '005_tomato_soup_can',

'006_mustard_bottle', '007_tuna_fish_can', '008_pudding_box', '009_gelatin_box',

'010_potted_meat_can', '011_banana', '019_pitcher_base', '021_bleach_cleanser',

'024_bowl', '025_mug', '035_power_drill',

'036_wood_block', '037_scissors', '040_large_marker', '051_large_clamp',

'052_extra_large_clamp',

'061_foam_brick']

self.ycb_class_to_idx = {}

for i, item in enumerate(self.object_names_ycbvideo):

self.ycb_class_to_idx[item] = i

self.kp3d = np.load(kp_path)

self.n_kp = 8

self.counter = -1

self.one_syn_per_batch = one_syn_per_batch

self.batch_size = batch_size

self.num_real_images = len(self.train_paths)

def gen_train_list(self, imageset_path, out_pkl="data/real_train_path.pkl"):

with open(opj(imageset_path, "trainval.txt"), 'r') as file:

trainlines = file.readlines()

real_train_path = [opj(self.root,x.rstrip('\n')) for x in trainlines]

with open(out_pkl, 'wb') as f:

pickle.dump(real_train_path, f)

self.train_paths = real_train_path

def gen_kp_gt_for_item(self, item):

# item is a path prefix. e.g. /media/data_2/YCB/data_syn/000104

out_pkl = item + '-bb8_2d.pkl'

meta = loadmat(item + '-meta.mat')

# bbox_file = item + '-box.txt'

# with open(bbox_file, 'rb') as f:

# bboxes = f.readlines()

intrinsic = meta['intrinsic_matrix'] # Note that this may vary between frames

poses = meta['poses'].transpose(2, 0, 1)

cls_idxs = meta['cls_indexes'] - 1

cls_idxs = cls_idxs.squeeze()

kp_2d = np.zeros((len(cls_idxs), self.n_kp, 2))

for idx, pose in enumerate(poses):

vertex = self.kp3d[int(cls_idxs[idx])].squeeze()

kp_2d[idx] = vertices_reprojection(vertex, pose, intrinsic)

kp_2d[:, :, 0] /= self.original_width

kp_2d[:, :, 1] /= self.original_height

print(kp_2d)

with open(out_pkl, 'wb') as f:

pickle.dump(kp_2d, f)

def gen_kp_gt(self, for_syn = True, for_real = True):

if for_real:

print("generate and save kp gt for real images.")

for item in tqdm(self.train_paths):

self.gen_kp_gt_for_item(item)

if for_syn:

print("generate and save kp gt for synthetic images.")

syn_prefix = self.syn_data_path

for id in tqdm(range(self.syn_range)):

item = opj(syn_prefix, "%06d" % id)

self.gen_kp_gt_for_item(item)

def gen_synthetic(self):

if len(self.syn_bg_image_paths)<1 :

print("you need to give bg images folder!")

# generate a synthetic image on the fly

prefix = self.syn_data_path

id = random.randint(0, self.syn_range-1)

item = opj(prefix, "%06d"%id)

raw = cv2.imread(item + "-color.png")

img = cv2.resize(raw, (self.input_height, self.input_width))

img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

# get segmentation gt

seg_img = cv2.imread(item + "-label.png")

seg_img = cv2.resize(seg_img, (self.input_height, self.input_width), interpolation=cv2.INTER_NEAREST)

mask_front = ma.getmaskarray(ma.masked_not_equal(seg_img, 0)).astype(int)

mask_back = ma.getmaskarray(ma.masked_equal(seg_img, 0)).astype(int)

# random erase some parts to make the network robust to occlusions

random_erasing = RandomErasing(sl=0.01,sh=0.1)

img = random_erasing(img)

# get bg image and combine them together

back_img_path = random.choice(self.syn_bg_image_paths)

bg_raw = cv2.imread(back_img_path)

bg_img = cv2.resize(bg_raw, (self.input_height, self.input_width))

bg_img = cv2.cvtColor(bg_img, cv2.COLOR_BGR2RGB)

if self.use_bg_img:

combined_img = bg_img * mask_back + img * mask_front

else:

combined_img = img * mask_front

color_augmentation = transforms.ColorJitter(0.02, 0.02, 0.02, 0.05)

combined_img = Image.fromarray(combined_img.astype('uint8')).convert('RGB')

combined_img = color_augmentation(combined_img)

combined_img = np.array(combined_img)

# get segmentation label

seg_label = seg_img[:, :, 0] # RGB channels are the same

seg_label = cv2.resize(seg_label, (self.target_h, self.target_w), interpolation=cv2.INTER_NEAREST)

# generate kp gt map of (nH, nW, nV)

kp_gt_map_x = np.zeros((self.target_h, self.target_w, self.n_kp))

kp_gt_map_y = np.zeros((self.target_h, self.target_w, self.n_kp))

in_pkl = item + '-bb8_2d.pkl'

# load class info

meta = loadmat(item + '-meta.mat')

#class_ids = meta['cls_indexes'] CORRECTION BELOW

class_ids = meta['cls_indexes'][0]

with open(in_pkl, 'rb') as f:

bb8_2d = pickle.load(f)

for i, cid in enumerate(class_ids):

class_mask = np.where(seg_label == cid)

# class_mask = np.where(seg_label == cid[0]) CORRECTION ABOVE

# print("CID: ", cid, "bb82d0: ", bb8_2d[:, :, 0][i], "bb82d1", bb8_2d[:, :, 1][i])

kp_gt_map_x[class_mask] = bb8_2d[:,:,0][i]

kp_gt_map_y[class_mask] = bb8_2d[:,:,1][i]

# DEBUG: save 100 syn image

# cv2.imwrite("./syn_images/syn-img-" + str(random.randint(0, 100)) + ".jpg",cv2.cvtColor(combined_img, cv2.COLOR_RGB2BGR))

# this mask front is used to compute loss

mask_front = cv2.resize(mask_front, (self.target_h, self.target_w), interpolation=cv2.INTER_NEAREST)

return (torch.from_numpy(combined_img.transpose(2, 0, 1)).float().div(255.0), torch.from_numpy(seg_label).long(),

torch.from_numpy(kp_gt_map_x).float(), torch.from_numpy(kp_gt_map_y).float(),

torch.from_numpy(mask_front[:,:,0]).float(),

torch.zeros(self.target_w, self.target_h).long())

def gen_balancing_weight(self, save_pkl="data/balancing_weight.pkl"):

# get pixel-wise balancing weight for cross entropy loss

pixels_per_img = (self.target_h * self.target_w)

real_frequency = [0 for x in range(self.num_classes)]

print("collect weight for real images")

for prefix in tqdm(self.train_paths):

label_img = cv2.imread(prefix + "-label.png")[: , : , 0]

label_img = cv2.resize(label_img, (self.target_h, self.target_w), interpolation=cv2.INTER_NEAREST)

labels_per_img = np.unique(label_img)

for img_id in labels_per_img:

if len(np.where(label_img==img_id)) <1:

real_frequency[img_id] += 0

else:

real_frequency[img_id] += len(np.where(label_img==img_id)[0]) / pixels_per_img

real_frequency = np.array(real_frequency)

real_frequency/=len(self.train_paths)

print("collect weights for syn images")

syn_frequency = [0 for x in range(self.num_classes)]

prefix = self.syn_data_path

for id in tqdm(range(self.syn_range - 1)):

item = opj(prefix, "%06d"%id)

seg_img = cv2.imread(item + "-label.png")

seg_img = cv2.resize(seg_img, (self.target_h, self.target_w), interpolation=cv2.INTER_NEAREST)

labels_per_img = np.unique(seg_img)

for img_id in labels_per_img:

if len(np.where(seg_img==img_id)) <1:

syn_frequency[img_id] += 0

else:

syn_frequency[img_id] += len(np.where(seg_img==img_id)[0]) / pixels_per_img

syn_frequency = np.array(syn_frequency)

syn_frequency/=self.syn_range

frequencies = {'real':real_frequency, 'syn':syn_frequency}

with open(save_pkl, 'wb') as f:

pickle.dump(frequencies, f)

def set_balancing_weight(self, save_pkl="data/balancing_weight.pkl"):

print("Loading weight from file ", save_pkl)

with open(save_pkl, 'rb') as f:

frequencies = pickle.load(f)

real_frequency = frequencies['real']

syn_frequency = frequencies['syn']

combined_frequency = self.num_syn_images * syn_frequency + len(self.train_paths) * real_frequency

median_frequency = np.median(combined_frequency)

weight = [median_frequency/x for x in combined_frequency]

self.weight_cross_entropy = torch.from_numpy(np.array(weight)).float()

def __getitem__(self, index):

self.counter += 1

if not self.use_real_img:

return self.gen_synthetic()

if index > self.num_real_images - 1:

return self.gen_synthetic()

# for DA, enforces that each batch has at least one source sample

if self.one_syn_per_batch and self.counter % self.batch_size == 0:

return self.gen_synthetic()

else:

prefix = self.train_paths[index]

# get raw image

raw = cv2.imread(prefix + "-color.png")

img = cv2.resize(raw, (self.input_height, self.input_width))

img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

# load class info

meta = loadmat(prefix + '-meta.mat')

class_ids = meta['cls_indexes']

# get segmentation gt, note 0 is for background

label_img = cv2.imread(prefix + "-label.png")[: , : , 0]

label_img = cv2.resize(label_img, (self.target_h, self.target_w), interpolation=cv2.INTER_NEAREST)

# generate kp gt map of (nH, nW, nV)

kp_gt_map_x = np.zeros((self.target_h, self.target_w, self.n_kp))

kp_gt_map_y = np.zeros((self.target_h, self.target_w, self.n_kp))

in_pkl = prefix + '-bb8_2d.pkl'

with open(in_pkl, 'rb') as f:

bb8_2d = pickle.load(f)

for i, cid in enumerate(class_ids):

class_mask = np.where(label_img == cid[0])

kp_gt_map_x[class_mask] = bb8_2d[:,:,0][i]

kp_gt_map_y[class_mask] = bb8_2d[:,:,1][i]

mask_front = ma.getmaskarray(ma.masked_not_equal(label_img, 0)).astype(int)

# debug: save 100 real imgages

#cv2.imwrite("./real_images/img-" + str(random.randint(0, 100)) + ".jpg",cv2.cvtColor(img, cv2.COLOR_RGB2BGR))

#TODO: get mask weighted by class

return (torch.from_numpy(img.transpose(2, 0, 1)).float().div(255.0),

torch.from_numpy(label_img).long(),

torch.from_numpy(kp_gt_map_x).float(), torch.from_numpy(kp_gt_map_y).float(),

torch.from_numpy(mask_front).float(),

torch.ones(self.target_w, self.target_h).long() )

def __len__(self):

if self.use_real_img:

return self.num_real_images + self.num_syn_images

else: