cakirin1212 · sourcery-ai · Oct 5, 2023
diff --git a/lib/core/config.py b/lib/core/config.py
@@ -125,9 +125,5 @@ def parse_args():
     print(args, end='\n\n')
 
     cfg_file = args.cfg
-    if args.cfg is not None:
-        cfg = update_cfg(args.cfg)
-    else:
-        cfg = get_cfg_defaults()
-
+    cfg = update_cfg(args.cfg) if args.cfg is not None else get_cfg_defaults()
     return cfg, cfg_file
diff --git a/lib/core/loss.py b/lib/core/loss.py
@@ -155,8 +155,12 @@ def keypoint_loss(self, pred_keypoints_2d, gt_keypoints_2d, openpose_weight, gt_
         conf = gt_keypoints_2d[:, :, -1].unsqueeze(-1).clone()
         conf[:, :25] *= openpose_weight
         conf[:, 25:] *= gt_weight
-        loss = (conf * self.criterion_keypoints(pred_keypoints_2d, gt_keypoints_2d[:, :, :-1])).mean()
-        return loss
+        return (
+            conf
+            * self.criterion_keypoints(
+                pred_keypoints_2d, gt_keypoints_2d[:, :, :-1]
+            )
+        ).mean()
 
     def keypoint_3d_loss(self, pred_keypoints_3d, gt_keypoints_3d):
         """

diff --git a/lib/core/trainer.py b/lib/core/trainer.py
@@ -217,11 +217,13 @@ def train(self):
             start = time.time()
 
             summary_string = f'({i + 1}/{self.num_iters_per_epoch}) | Total: {bar.elapsed_td} | ' \
-                             f'ETA: {bar.eta_td:} | loss: {losses.avg:.4f}'
+                                 f'ETA: {bar.eta_td:} | loss: {losses.avg:.4f}'
 
             for k, v in loss_dict.items():
                 summary_string += f' | {k}: {v:.2f}'
-                self.writer.add_scalar('train_loss/'+k, v, global_step=self.train_global_step)
+                self.writer.add_scalar(
+                    f'train_loss/{k}', v, global_step=self.train_global_step
+                )
 
             for k,v in timer.items():
                 summary_string += f' | {k}: {v:.2f}'
@@ -243,7 +245,7 @@ def train(self):
 
             if torch.isnan(total_loss):
                 exit('Nan value in loss, exiting!...')
-            # =======>
+                # =======>
 
         bar.finish()
 

diff --git a/lib/data_utils/feature_extractor.py b/lib/data_utils/feature_extractor.py
@@ -35,7 +35,7 @@ def extract_features(model, video, bbox, debug=False, batch_size=200, kp_2d=None
     '''
     device = 'cuda'
 
-    if isinstance(video, torch.Tensor) or isinstance(video, np.ndarray):
+    if isinstance(video, (torch.Tensor, np.ndarray)):
         video = video
     elif isinstance(video, str):
         if os.path.isfile(video):

diff --git a/lib/data_utils/img_utils.py b/lib/data_utils/img_utils.py
@@ -40,7 +40,7 @@ def do_augmentation(scale_factor=0.3, color_factor=0.2):
 def trans_point2d(pt_2d, trans):
     src_pt = np.array([pt_2d[0], pt_2d[1], 1.]).T
     dst_pt = np.dot(trans, src_pt)
-    return dst_pt[0:2]
+    return dst_pt[:2]
 
 def rotate_2d(pt_2d, rot_rad):
     x = pt_2d[0]
@@ -78,12 +78,11 @@ def gen_trans_from_patch_cv(c_x, c_y, src_width, src_height, dst_width, dst_heig
     dst[1, :] = dst_center + dst_downdir
     dst[2, :] = dst_center + dst_rightdir
 
-    if inv:
-        trans = cv2.getAffineTransform(np.float32(dst), np.float32(src))
-    else:
-        trans = cv2.getAffineTransform(np.float32(src), np.float32(dst))
-
-    return trans
+    return (
+        cv2.getAffineTransform(np.float32(dst), np.float32(src))
+        if inv
+        else cv2.getAffineTransform(np.float32(src), np.float32(dst))
+    )
 
 def generate_patch_image_cv(cvimg, c_x, c_y, bb_width, bb_height, patch_width, patch_height, do_flip, scale, rot):
     img = cvimg.copy()
@@ -174,8 +173,7 @@ def get_image_crops(image_file, bboxes):
         crop_image = convert_cvimg_to_tensor(crop_image)
         crop_images.append(crop_image)
 
-    batch_image = torch.cat([x.unsqueeze(0) for x in crop_images])
-    return batch_image
+    return torch.cat([x.unsqueeze(0) for x in crop_images])
 
 def get_single_image_crop(image, bbox, scale=1.3):
     if isinstance(image, str):
@@ -295,29 +293,24 @@ def get_bbox_from_kp2d(kp_2d):
     w = h = np.where(w / h > 1, w, h)
     w = h = h * 1.1
 
-    bbox = np.array([c_x, c_y, w, h])  # shape = (4,N)
-    return bbox
+    return np.array([c_x, c_y, w, h])
 
 def normalize_2d_kp(kp_2d, crop_size=224, inv=False):
+    ratio = 1.0 / crop_size
     # Normalize keypoints between -1, 1
-    if not inv:
-        ratio = 1.0 / crop_size
-        kp_2d = 2.0 * kp_2d * ratio - 1.0
-    else:
-        ratio = 1.0 / crop_size
-        kp_2d = (kp_2d + 1.0)/(2*ratio)
-
+    kp_2d = 2.0 * kp_2d * ratio - 1.0 if not inv else (kp_2d + 1.0)/(2*ratio)
     return kp_2d
 
 def get_default_transform():
     normalize = transforms.Normalize(
         mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]
     )
-    transform = transforms.Compose([
-        transforms.ToTensor(),
-        normalize,
-    ])
-    return transform
+    return transforms.Compose(
+        [
+            transforms.ToTensor(),
+            normalize,
+        ]
+    )
 
 def split_into_chunks(vid_names, seqlen, stride):
     video_start_end_indices = []

diff --git a/lib/data_utils/insta_utils.py b/lib/data_utils/insta_utils.py
@@ -76,18 +76,14 @@ def decode_jpeg(self, image_data):
         return image
 
     def encode_jpeg(self, image):
-        image_data = self._sess.run(
-            self._encode_jpeg, feed_dict={
-                self._encode_jpeg_data: image
-            })
-        return image_data
+        return self._sess.run(
+            self._encode_jpeg, feed_dict={self._encode_jpeg_data: image}
+        )
 
     def encode_png(self, image):
-        image_data = self._sess.run(
-            self._encode_png, feed_dict={
-                self._encode_png_data: image
-            })
-        return image_data
+        return self._sess.run(
+            self._encode_png, feed_dict={self._encode_png_data: image}
+        )
 
     def decode_png(self, image_data):
         image = self._sess.run(
@@ -326,7 +322,7 @@ def read_single_record(fname):
         print(features.shape)
         assert features.shape[0] == N
 
-    for k in dataset.keys():
+    for k in dataset:
         dataset[k] = np.concatenate(dataset[k])
 
     for k,v in dataset.items():
@@ -355,10 +351,10 @@ def concatenate_annotations():
         filename = osp.join(VIBE_DB_DIR, 'insta_parts', f'insta_train_part_{i}.h5')
         print(filename)
         with h5py.File(filename, 'r') as f:
-            for k in ds.keys():
+            for k in ds:
                 ds[k].append(f[k].value)
 
-    for k in ds.keys():
+    for k in ds:
         ds[k] = np.concatenate(ds[k])
 
     print('Saving Insta Variety dataset!..')

diff --git a/lib/data_utils/kp_utils.py b/lib/data_utils/kp_utils.py
@@ -40,8 +40,7 @@ def convert_kps(joints2d, src, dst):
 def get_perm_idxs(src, dst):
     src_names = eval(f'get_{src}_joint_names')()
     dst_names = eval(f'get_{dst}_joint_names')()
-    idxs = [src_names.index(h) for h in dst_names if h in src_names]
-    return idxs
+    return [src_names.index(h) for h in dst_names if h in src_names]
 
 def get_mpii3d_test_joint_names():
     return [

diff --git a/lib/data_utils/mpii3d_utils.py b/lib/data_utils/mpii3d_utils.py
@@ -19,43 +19,42 @@
 
 
 def read_openpose(json_file, gt_part, dataset):
-    # get only the arms/legs joints
-    op_to_12 = [11, 10, 9, 12, 13, 14, 4, 3, 2, 5, 6, 7]
     # read the openpose detection
     json_data = json.load(open(json_file, 'r'))
     people = json_data['people']
     if len(people) == 0:
         # no openpose detection
-        keyp25 = np.zeros([25,3])
-    else:
-        # size of person in pixels
-        scale = max(max(gt_part[:,0])-min(gt_part[:,0]),max(gt_part[:,1])-min(gt_part[:,1]))
-        # go through all people and find a match
-        dist_conf = np.inf*np.ones(len(people))
-        for i, person in enumerate(people):
-            # openpose keypoints
-            op_keyp25 = np.reshape(person['pose_keypoints_2d'], [25,3])
-            op_keyp12 = op_keyp25[op_to_12, :2]
-            op_conf12 = op_keyp25[op_to_12, 2:3] > 0
+        return np.zeros([25,3])
+    # size of person in pixels
+    scale = max(max(gt_part[:,0])-min(gt_part[:,0]),max(gt_part[:,1])-min(gt_part[:,1]))
+    # go through all people and find a match
+    dist_conf = np.inf*np.ones(len(people))
+    # get only the arms/legs joints
+    op_to_12 = [11, 10, 9, 12, 13, 14, 4, 3, 2, 5, 6, 7]
+    for i, person in enumerate(people):
+        # openpose keypoints
+        op_keyp25 = np.reshape(person['pose_keypoints_2d'], [25,3])
+        op_conf12 = op_keyp25[op_to_12, 2:3] > 0
             # all the relevant joints should be detected
-            if min(op_conf12) > 0:
-                # weighted distance of keypoints
-                dist_conf[i] = np.mean(np.sqrt(np.sum(op_conf12*(op_keyp12 - gt_part[:12, :2])**2, axis=1)))
-        # closest match
-        p_sel = np.argmin(dist_conf)
+        if min(op_conf12) > 0:
+            op_keyp12 = op_keyp25[op_to_12, :2]
+            # weighted distance of keypoints
+            dist_conf[i] = np.mean(np.sqrt(np.sum(op_conf12*(op_keyp12 - gt_part[:12, :2])**2, axis=1)))
+    # closest match
+    p_sel = np.argmin(dist_conf)
         # the exact threshold is not super important but these are the values we used
-        if dataset == 'mpii':
-            thresh = 30
-        elif dataset == 'coco':
-            thresh = 10
-        else:
-            thresh = 0
+    if dataset == 'coco':
+        thresh = 10
+    elif dataset == 'mpii':
+        thresh = 30
+    else:
+        thresh = 0
         # dataset-specific thresholding based on pixel size of person
-        if min(dist_conf)/scale > 0.1 and min(dist_conf) < thresh:
-            keyp25 = np.zeros([25,3])
-        else:
-            keyp25 = np.reshape(people[p_sel]['pose_keypoints_2d'], [25,3])
-    return keyp25
+    return (
+        np.zeros([25, 3])
+        if min(dist_conf) / scale > 0.1 and min(dist_conf) < thresh
+        else np.reshape(people[p_sel]['pose_keypoints_2d'], [25, 3])
+    )
 
 
 def read_calibration(calib_file, vid_list):
@@ -99,26 +98,23 @@ def read_data_train(dataset_path, debug=False):
 
     for user_i in user_list:
         for seq_i in seq_list:
-            seq_path = os.path.join(dataset_path,
-                                    'S' + str(user_i),
-                                    'Seq' + str(seq_i))
+            seq_path = os.path.join(dataset_path, f'S{str(user_i)}', f'Seq{str(seq_i)}')
             # mat file with annotations
             annot_file = os.path.join(seq_path, 'annot.mat')
             annot2 = sio.loadmat(annot_file)['annot2']
             annot3 = sio.loadmat(annot_file)['annot3']
             # calibration file and camera parameters
             for j, vid_i in enumerate(vid_list):
                 # image folder
-                imgs_path = os.path.join(seq_path,
-                                         'video_' + str(vid_i))
+                imgs_path = os.path.join(seq_path, f'video_{str(vid_i)}')
                 # per frame
                 pattern = os.path.join(imgs_path, '*.jpg')
                 img_list = sorted(glob.glob(pattern))
                 vid_used_frames = []
                 vid_used_joints = []
                 vid_used_bbox = []
                 vid_segments = []
-                vid_uniq_id = "subj" + str(user_i) + '_seq' + str(seq_i) + "_vid" + str(vid_i) + "_seg0"
+                vid_uniq_id = f"subj{str(user_i)}_seq{str(seq_i)}_vid{str(vid_i)}_seg0"
                 for i, img_i in tqdm_enumerate(img_list):
 
                     # for each image we store the relevant annotations
@@ -191,7 +187,7 @@ def read_data_train(dataset_path, debug=False):
                                                 dataset='spin', debug=False)
                     dataset['features'].append(features)
 
-    for k in dataset.keys():
+    for k in dataset:
         dataset[k] = np.array(dataset[k])
     dataset['features'] = np.concatenate(dataset['features'])
 
@@ -217,9 +213,9 @@ def read_test_data(dataset_path):
 
     for user_i in user_list:
         print('Subject', user_i)
-        seq_path = os.path.join(dataset_path,
-                                'mpi_inf_3dhp_test_set',
-                                'TS' + str(user_i))
+        seq_path = os.path.join(
+            dataset_path, 'mpi_inf_3dhp_test_set', f'TS{str(user_i)}'
+        )
         # mat file with annotations
         annot_file = os.path.join(seq_path, 'annot_data.mat')
         mat_as_h5 = h5py.File(annot_file, 'r')
@@ -231,15 +227,17 @@ def read_test_data(dataset_path):
         vid_used_joints = []
         vid_used_bbox = []
         vid_segments = []
-        vid_uniq_id = "subj" + str(user_i) + "_seg0"
+        vid_uniq_id = f"subj{str(user_i)}_seg0"
 
 
         for frame_i, valid_i in tqdm(enumerate(valid)):
 
-            img_i = os.path.join('mpi_inf_3dhp_test_set',
-                                    'TS' + str(user_i),
-                                    'imageSequence',
-                                    'img_' + str(frame_i + 1).zfill(6) + '.jpg')
+            img_i = os.path.join(
+                'mpi_inf_3dhp_test_set',
+                f'TS{str(user_i)}',
+                'imageSequence',
+                f'img_{str(frame_i + 1).zfill(6)}.jpg',
+            )
 
             joints_2d_raw = np.expand_dims(annot2[frame_i, 0, :, :], axis = 0)
             joints_2d_raw = np.append(joints_2d_raw, np.ones((1, 17, 1)), axis=2)
@@ -317,7 +315,7 @@ def read_test_data(dataset_path):
                                         dataset='spin', debug=False)
             dataset['features'].append(features)
 
-    for k in dataset.keys():
+    for k in dataset:
         dataset[k] = np.array(dataset[k])
     dataset['features'] = np.concatenate(dataset['features'])
 

diff --git a/lib/data_utils/penn_action_utils.py b/lib/data_utils/penn_action_utils.py
@@ -71,11 +71,17 @@ def read_data(folder):
 
     model = spin.get_pretrained_hmr()
 
-    file_names = sorted(glob.glob(folder + '/labels/'+'*.mat'))
+    file_names = sorted(glob.glob(f'{folder}/labels/*.mat'))
 
     for fname in tqdm(file_names):
         vid_dict=load_mat(fname)
-        imgs = sorted(glob.glob(folder + '/frames/'+ fname.strip().split('/')[-1].split('.')[0]+'/*.jpg'))
+        imgs = sorted(
+            glob.glob(
+                f'{folder}/frames/'
+                + fname.strip().split('/')[-1].split('.')[0]
+                + '/*.jpg'
+            )
+        )
         kp_2d = np.zeros((vid_dict['nframes'], 13, 3))
         perm_idxs = get_perm_idxs('pennaction', 'common')
 
@@ -109,9 +115,9 @@ def read_data(folder):
         features = extract_features(model, np.array(imgs) , bbox, dataset='pennaction', debug=False)
         dataset['features'].append(features)
 
-    for k in dataset.keys():
+    for k in dataset:
         dataset[k] = np.array(dataset[k])
-    for k in dataset.keys():
+    for k in dataset:
         dataset[k] = np.concatenate(dataset[k])
 
     return dataset