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)