def work_op(loader, queue_op):
while 1:
if queue_op.qsize() < 5:
batch = opcaffe.Batch()
myClass.load(batch)
data = torch.tensor(batch.data)
label = torch.tensor(batch.label)
queue_op.put([data, label])
time.sleep(0.1)
def work(loader, queue, control):
while 1:
if control.value == 0:
break
if queue.qsize() < 5:
batch = opcaffe.Batch()
myClass.load(batch)
data = torch.tensor(batch.data)
label = torch.tensor(batch.label)
queue.put([data, label])
time.sleep(0.1)
def get_index(self, i):
img = cv2.imread(self.db_data["img_dirs"][i])
K = self.db_data["K"][i]
R = self.db_data["R"][i]
t = self.db_data["t"][i]
distCoef = self.db_data["distCoef"][i]
calib = {"K": K, "R": R, "t": t, "distCoef": distCoef}
joints = self.db_data["body"][i]
# Convert to COCO Format
joints = convert(joints)
# Project
pt = project2D(joints,
calib,
imgwh=(img.shape[1], img.shape[0]),
applyDistort=True)
# Create elems
points_3d = np.zeros((17, 3))
points_2d = np.zeros((17, 3))
for j in range(0, joints.shape[0]):
point_2d = pt[0][:, j]
valid = pt[1][j]
point_3d = pt[2][:, j]
if valid:
points_3d[j, :] = point_3d
points_2d[j, :] = np.array([point_2d[0], point_2d[1], 1])
else:
points_3d[j, :] = np.array([0, 0, 0])
points_2d[j, :] = np.array([0, 0, 2])
# Stuff
batch = opcaffe.Batch()
metaData = create_meta(points_3d, points_2d, img)
self.opTransformer.load(img, metaData, batch)
#viz_pof(batch)
# Return
image = batch.data.copy()
paf_mask = batch.label[:, 0:72, :, :].copy()
hm_mask = batch.label[:, 72:97, :, :].copy()
paf = batch.label[:, 97:169, :, :].copy()
hm = batch.label[:, 169:194, :, :].copy()
pof_mask = np.zeros((1, 24 * 3, 46, 46), dtype=np.float32)
counter = 0
for j in range(0, 24 * 2, 2):
pof_mask[:, counter * 3 + 0, :, :] = paf_mask[:, j, :, :]
pof_mask[:, counter * 3 + 1, :, :] = paf_mask[:, j, :, :]
pof_mask[:, counter * 3 + 2, :, :] = paf_mask[:, j, :, :]
counter += 1
pof = batch.other.copy()
return image, paf_mask, paf, pof_mask, pof, hm_mask, hm
params = {
"batch_size": 5,
"stride": 8,
"max_degree_rotations": "45.0",
"crop_size_x": 368,
"crop_size_y": 368,
"center_perterb_max": 40.0,
"center_swap_prob": 0.0,
"scale_prob": 1.0,
"scale_mins": "0.333333333333",
"scale_maxs": "1.5",
"target_dist": 0.600000023842,
"number_max_occlusions": "2",
"sigmas": "7.0",
"models": "COCO_25B_23;COCO_25B_17;MPII_25B_16;PT_25B_15",
"sources":
"/media/raaj/Storage/openpose_train/dataset/lmdb_coco2017_foot;/media/raaj/Storage/openpose_train/dataset/lmdb_coco;/media/raaj/Storage/openpose_train/dataset/lmdb_mpii;/media/raaj/Storage/openpose_train/dataset/lmdb_pt2_train",
"probabilities": "0.05;0.85;0.05;0.05",
"source_background":
"/media/raaj/Storage/openpose_train/dataset/lmdb_background",
"normalization": 0,
"add_distance": 0
}
myClass = opcaffe.OPCaffe(params)
while 1:
batch = opcaffe.Batch()
myClass.load(batch)
#print batch.label.shape
#print batch.data.shape
def __getitem__(self, index):
pid = os.getpid() % self.WORKER_SIZE
batch = opcaffe.Batch()
self.workers[pid].load(batch)
return torch.tensor(batch.data), torch.tensor(batch.label)