def __init__()

in data/STB/dataset.py [0:0]

• 53 lines of code
• 7 McCabe index (conditional complexity)

    def __init__(self, transform, mode):
        self.mode = mode
        self.root_path = '../data/STB/data'
        self.rootnet_output_path = '../data/STB/rootnet_output/rootnet_stb_output.json'
        self.original_img_shape = (480, 640) # height, width
        self.transform = transform
        self.joint_num = 21 # single hand
        self.joint_type = {'right': np.arange(self.joint_num,self.joint_num*2), 'left': np.arange(0,self.joint_num)}
        self.root_joint_idx = {'right': self.joint_num, 'left': 0}
        self.skeleton = load_skeleton(osp.join(self.root_path, 'skeleton.txt'), self.joint_num*2)
        
        self.datalist = [];
        self.annot_path = osp.join(self.root_path, 'STB_' + self.mode + '.json')
        db = COCO(self.annot_path)

        if self.mode == 'test' and cfg.trans_test == 'rootnet':
            print("Get bbox and root depth from " + self.rootnet_output_path)
            rootnet_result = {}
            with open(self.rootnet_output_path) as f:
                annot = json.load(f)
            for i in range(len(annot)):
                rootnet_result[str(annot[i]['annot_id'])] = annot[i]
        else:
            print("Get bbox and root depth from groundtruth annotation")

        for aid in db.anns.keys():
            ann = db.anns[aid]
            image_id = ann['image_id']
            img = db.loadImgs(image_id)[0]
            
            seq_name = img['seq_name']
            img_path = osp.join(self.root_path, 'images', seq_name, img['file_name'])
            img_width, img_height = img['width'], img['height']
            cam_param = img['cam_param']
            focal, princpt = np.array(cam_param['focal'],dtype=np.float32), np.array(cam_param['princpt'],dtype=np.float32)
            
            joint_img = np.array(ann['joint_img'],dtype=np.float32)
            joint_cam = np.array(ann['joint_cam'],dtype=np.float32)
            joint_valid = np.array(ann['joint_valid'],dtype=np.float32)
            
            # transform single hand data to double hand data structure
            hand_type = ann['hand_type']
            joint_img_dh = np.zeros((self.joint_num*2,2),dtype=np.float32)
            joint_cam_dh = np.zeros((self.joint_num*2,3),dtype=np.float32)
            joint_valid_dh = np.zeros((self.joint_num*2),dtype=np.float32)
            joint_img_dh[self.joint_type[hand_type]] = joint_img
            joint_cam_dh[self.joint_type[hand_type]] = joint_cam
            joint_valid_dh[self.joint_type[hand_type]] = joint_valid
            joint_img = joint_img_dh; joint_cam = joint_cam_dh; joint_valid = joint_valid_dh;

            if self.mode == 'test' and cfg.trans_test == 'rootnet':
                bbox = np.array(rootnet_result[str(aid)]['bbox'],dtype=np.float32)
                abs_depth = rootnet_result[str(aid)]['abs_depth']
            else:
                bbox = np.array(ann['bbox'],dtype=np.float32) # x,y,w,h
                bbox = process_bbox(bbox, (img_height, img_width))
                abs_depth = joint_cam[self.root_joint_idx[hand_type],2] # single hand abs depth
            
            cam_param = {'focal': focal, 'princpt': princpt}
            joint = {'cam_coord': joint_cam, 'img_coord': joint_img, 'valid': joint_valid}
            data = {'img_path': img_path, 'bbox': bbox, 'cam_param': cam_param, 'joint': joint, 'hand_type': hand_type, 'abs_depth': abs_depth}
            self.datalist.append(data)