From 3c537e72b9c5b22af78a81d4b3cfb1b0b7e96980 Mon Sep 17 00:00:00 2001 From: Sourcery AI Date: Tue, 7 Feb 2023 01:46:15 +0000 Subject: [PATCH] 'Refactored by Sourcery' --- demo/picture_demo.py | 8 +- demo/web_demo.py | 10 +- evaluate/coco_eval.py | 28 +++-- lib/datasets/datasets.py | 50 ++++----- lib/datasets/heatmap.py | 6 +- lib/datasets/paf.py | 8 +- lib/datasets/preprocessing.py | 29 ++--- lib/datasets/test_dataloader.py | 48 ++++----- lib/datasets/utils.py | 15 +-- lib/network/atrous_model.py | 44 ++++---- lib/network/atrous_model_share_stages.py | 36 +++---- lib/network/atrouspose.py | 9 +- lib/network/im_transform.py | 14 +-- lib/network/openpose.py | 22 ++-- lib/network/post.py | 41 ++++--- lib/network/rtpose_hourglass.py | 34 +++--- lib/network/rtpose_mobilenetV2.py | 5 +- lib/network/rtpose_shufflenetV2.py | 91 +++++++++++----- lib/network/rtpose_vgg.py | 92 +++++++--------- lib/utils/common.py | 21 ++-- lib/utils/paf_to_pose.py | 28 ++--- train/train_SH.py | 130 +++++++++++------------ train/train_ShuffleNetV2.py | 130 +++++++++++------------ train/train_VGG19.py | 61 +++++------ video_demo.py | 28 ++--- 25 files changed, 493 insertions(+), 495 deletions(-) diff --git a/demo/picture_demo.py b/demo/picture_demo.py index e374501..0d0a43b 100644 --- a/demo/picture_demo.py +++ b/demo/picture_demo.py @@ -42,7 +42,7 @@ -model = get_model('vgg19') +model = get_model('vgg19') model.load_state_dict(torch.load(args.weight)) model = torch.nn.DataParallel(model).cuda() model.float() @@ -50,16 +50,16 @@ test_image = './readme/ski.jpg' oriImg = cv2.imread(test_image) # B,G,R order -shape_dst = np.min(oriImg.shape[0:2]) +shape_dst = np.min(oriImg.shape[:2]) # Get results of original image with torch.no_grad(): paf, heatmap, im_scale = get_outputs(oriImg, model, 'rtpose') - + print(im_scale) humans = paf_to_pose_cpp(heatmap, paf, cfg) - + out = draw_humans(oriImg, humans) cv2.imwrite('result.png',out) diff --git a/demo/web_demo.py b/demo/web_demo.py index 960e0b4..3aa4273 100644 --- a/demo/web_demo.py +++ b/demo/web_demo.py @@ -40,7 +40,7 @@ # update config file update_config(cfg, args) -model = get_model('vgg19') +model = get_model('vgg19') model.load_state_dict(torch.load(args.weight)) model.cuda() model.float() @@ -53,15 +53,15 @@ while True: # Capture frame-by-frame ret, oriImg = video_capture.read() - - shape_dst = np.min(oriImg.shape[0:2]) + + shape_dst = np.min(oriImg.shape[:2]) with torch.no_grad(): paf, heatmap, imscale = get_outputs( oriImg, model, 'rtpose') - + humans = paf_to_pose_cpp(heatmap, paf, cfg) - + out = draw_humans(oriImg, humans) # Display the resulting frame diff --git a/evaluate/coco_eval.py b/evaluate/coco_eval.py index d552ba6..22d4ce1 100644 --- a/evaluate/coco_eval.py +++ b/evaluate/coco_eval.py @@ -124,14 +124,13 @@ def append_result(image_id, humans, upsample_keypoints, outputs): """ for human in humans: one_result = { - "image_id": 0, "category_id": 1, "keypoints": [], - "score": 0 + "score": 0, + "image_id": image_id, } - one_result["image_id"] = image_id keypoints = np.zeros((18, 3)) - + all_scores = [] for i in range(cfg.MODEL.NUM_KEYPOINTS): if i not in human.body_parts.keys(): @@ -146,7 +145,7 @@ def append_result(image_id, humans, upsample_keypoints, outputs): keypoints[i, 2] = 1 score = human.body_parts[i].score all_scores.append(score) - + keypoints = keypoints[ORDER_COCO,:] one_result["score"] = 1. one_result["keypoints"] = list(keypoints.reshape(51)) @@ -165,20 +164,19 @@ def append_result_legacy(image_id, person_to_joint_assoc, joint_list, outputs): for ridxPred in range(len(person_to_joint_assoc)): one_result = { - "image_id": 0, "category_id": 1, "keypoints": [], - "score": 0 + "score": 0, + "image_id": image_id, } - one_result["image_id"] = image_id keypoints = np.zeros((17, 3)) for part in range(17): ind = ORDER_COCO[part] index = int(person_to_joint_assoc[ridxPred, ind]) - if -1 == index: + if index == -1: keypoints[part, 0] = 0 keypoints[part, 1] = 0 keypoints[part, 2] = 0 @@ -248,31 +246,31 @@ def run_eval(image_dir, anno_file, vis_dir, model, preprocess): :returns: float, the reported mAP score """ coco = COCO(anno_file) - cat_ids = coco.getCatIds(catNms=['person']) + cat_ids = coco.getCatIds(catNms=['person']) img_ids = coco.getImgIds(catIds=cat_ids) - print("Total number of validation images {}".format(len(img_ids))) + print(f"Total number of validation images {len(img_ids)}") # iterate all val images outputs = [] print("Processing Images in validation set") for i in range(len(img_ids)): if i % 10 == 0 and i != 0: - print("Processed {} images".format(i)) + print(f"Processed {i} images") img = coco.loadImgs(img_ids[i])[0] file_name = img['file_name'] file_path = os.path.join(image_dir, file_name) oriImg = cv2.imread(file_path) # Get the shortest side of the image (either height or width) - shape_dst = np.min(oriImg.shape[0:2]) + shape_dst = np.min(oriImg.shape[:2]) # Get results of original image paf, heatmap, scale_img = get_outputs(oriImg, model, preprocess) humans = paf_to_pose_cpp(heatmap, paf, cfg) - + out = draw_humans(oriImg, humans) - + vis_path = os.path.join(vis_dir, file_name) cv2.imwrite(vis_path, out) # subset indicated how many peoples foun in this image. diff --git a/lib/datasets/datasets.py b/lib/datasets/datasets.py index ac47141..4c0ccd1 100644 --- a/lib/datasets/datasets.py +++ b/lib/datasets/datasets.py @@ -11,39 +11,36 @@ from . import transforms, utils def kp_connections(keypoints): - kp_lines = [ - [keypoints.index('neck'), keypoints.index('right_hip')], + return [ + [keypoints.index('neck'), keypoints.index('right_hip')], [keypoints.index('right_hip'), keypoints.index('right_knee')], [keypoints.index('right_knee'), keypoints.index('right_ankle')], - [keypoints.index('neck'), keypoints.index('left_hip')], + [keypoints.index('neck'), keypoints.index('left_hip')], [keypoints.index('left_hip'), keypoints.index('left_knee')], [keypoints.index('left_knee'), keypoints.index('left_ankle')], - [keypoints.index('neck'), keypoints.index('right_shoulder')], + [keypoints.index('neck'), keypoints.index('right_shoulder')], [keypoints.index('right_shoulder'), keypoints.index('right_elbow')], - [keypoints.index('right_elbow'), keypoints.index('right_wrist')], - [keypoints.index('right_shoulder'), keypoints.index('right_eye')], - [keypoints.index('neck'), keypoints.index('left_shoulder')], + [keypoints.index('right_elbow'), keypoints.index('right_wrist')], + [keypoints.index('right_shoulder'), keypoints.index('right_eye')], + [keypoints.index('neck'), keypoints.index('left_shoulder')], [keypoints.index('left_shoulder'), keypoints.index('left_elbow')], [keypoints.index('left_elbow'), keypoints.index('left_wrist')], - [keypoints.index('left_shoulder'), keypoints.index('left_eye')], - [keypoints.index('neck'), keypoints.index('nose')], + [keypoints.index('left_shoulder'), keypoints.index('left_eye')], + [keypoints.index('neck'), keypoints.index('nose')], [keypoints.index('nose'), keypoints.index('right_eye')], - [keypoints.index('nose'), keypoints.index('left_eye')], + [keypoints.index('nose'), keypoints.index('left_eye')], [keypoints.index('right_eye'), keypoints.index('right_ear')], - [keypoints.index('left_eye'), keypoints.index('left_ear')] + [keypoints.index('left_eye'), keypoints.index('left_ear')], ] - return kp_lines def get_keypoints(): """Get the COCO keypoints and their left/right flip coorespondence map.""" - # Keypoints are not available in the COCO json for the test split, so we - # provide them here. - keypoints = [ + return [ 'nose', 'neck', 'right_shoulder', 'right_elbow', - 'right_wrist', + 'right_wrist', 'left_shoulder', 'left_elbow', 'left_wrist', @@ -53,12 +50,11 @@ def get_keypoints(): 'left_hip', 'left_knee', 'left_ankle', - 'right_eye', + 'right_eye', 'left_eye', 'right_ear', - 'left_ear'] - - return keypoints + 'left_ear', + ] def collate_images_anns_meta(batch): images = torch.utils.data.dataloader.default_collate([b[0] for b in batch]) @@ -124,16 +120,16 @@ def __init__(self, root, annFile, image_transform=None, target_transforms=None, self.filter_for_keypoint_annotations() if n_images: self.ids = self.ids[:n_images] - print('Images: {}'.format(len(self.ids))) + print(f'Images: {len(self.ids)}') self.preprocess = preprocess or transforms.Normalize() self.image_transform = image_transform or transforms.image_transform self.target_transforms = target_transforms - + self.HEATMAP_COUNT = len(get_keypoints()) self.LIMB_IDS = kp_connections(get_keypoints()) self.input_y = input_y - self.input_x = input_x + self.input_x = input_x self.stride = stride self.log = logging.getLogger(self.__class__.__name__) @@ -257,13 +253,13 @@ def add_neck(self, keypoint): return keypoint def get_ground_truth(self, anns): - + grid_y = int(self.input_y / self.stride) grid_x = int(self.input_x / self.stride) channels_heat = (self.HEATMAP_COUNT + 1) channels_paf = 2 * len(self.LIMB_IDS) - heatmaps = np.zeros((int(grid_y), int(grid_x), channels_heat)) - pafs = np.zeros((int(grid_y), int(grid_x), channels_paf)) + heatmaps = np.zeros((grid_y, grid_x, channels_heat)) + pafs = np.zeros((grid_y, grid_x, channels_paf)) keypoints = [] for ann in anns: @@ -286,7 +282,7 @@ def get_ground_truth(self, anns): # pafs for i, (k1, k2) in enumerate(self.LIMB_IDS): # limb - count = np.zeros((int(grid_y), int(grid_x)), dtype=np.uint32) + count = np.zeros((grid_y, grid_x), dtype=np.uint32) for joint in keypoints: if joint[k1, 2] > 0.5 and joint[k2, 2] > 0.5: centerA = joint[k1, :2] diff --git a/lib/datasets/heatmap.py b/lib/datasets/heatmap.py index 1bc9d1c..9bf7588 100644 --- a/lib/datasets/heatmap.py +++ b/lib/datasets/heatmap.py @@ -20,8 +20,8 @@ def putGaussianMaps(center, accumulate_confid_map, sigma, grid_y, grid_x, stride): start = stride / 2.0 - 0.5 - y_range = [i for i in range(int(grid_y))] - x_range = [i for i in range(int(grid_x))] + y_range = list(range(int(grid_y))) + x_range = list(range(int(grid_x))) xx, yy = np.meshgrid(x_range, y_range) xx = xx * stride + start yy = yy * stride + start @@ -32,5 +32,5 @@ def putGaussianMaps(center, accumulate_confid_map, sigma, grid_y, grid_x, stride cofid_map = np.multiply(mask, cofid_map) accumulate_confid_map += cofid_map accumulate_confid_map[accumulate_confid_map > 1.0] = 1.0 - + return accumulate_confid_map diff --git a/lib/datasets/paf.py b/lib/datasets/paf.py index 4cf8cd0..2f9b587 100644 --- a/lib/datasets/paf.py +++ b/lib/datasets/paf.py @@ -37,8 +37,8 @@ def putVecMaps(centerA, centerB, accumulate_vec_map, count, grid_y, grid_x, stri min_y = max(int(round(min(centerA[1], centerB[1]) - thre)), 0) max_y = min(int(round(max(centerA[1], centerB[1]) + thre)), grid_y) - range_x = list(range(int(min_x), int(max_x), 1)) - range_y = list(range(int(min_y), int(max_y), 1)) + range_x = list(range(int(min_x), int(max_x))) + range_y = list(range(int(min_y), int(max_y))) xx, yy = np.meshgrid(range_x, range_y) ba_x = xx - centerA[0] # the vector from (x,y) to centerA ba_y = yy - centerA[1] @@ -60,9 +60,9 @@ def putVecMaps(centerA, centerB, accumulate_vec_map, count, grid_y, grid_x, stri mask = count == 0 - count[mask == True] = 1 + count[mask] = 1 accumulate_vec_map = np.divide(accumulate_vec_map, count[:, :, np.newaxis]) - count[mask == True] = 0 + count[mask] = 0 return accumulate_vec_map, count diff --git a/lib/datasets/preprocessing.py b/lib/datasets/preprocessing.py index 1ac9ccf..7efe01d 100644 --- a/lib/datasets/preprocessing.py +++ b/lib/datasets/preprocessing.py @@ -16,18 +16,13 @@ def rtpose_preprocess(image): image = image.astype(np.float32) image = image / 256. - 0.5 - image = image.transpose((2, 0, 1)).astype(np.float32) - - return image + return image.transpose((2, 0, 1)).astype(np.float32) def inverse_rtpose_preprocess(image): image = image.astype(np.float32) - image = image.transpose((1, 2, 0)).astype(np.float32) - image = (image + 0.5) * 256. - image = image.astype(np.uint8) - - - return image + image = image.transpose((1, 2, 0)).astype(np.float32) + image = (image + 0.5) * 256. + return image.astype(np.uint8) def vgg_preprocess(image): image = image.astype(np.float32) / 255. @@ -47,9 +42,7 @@ def inception_preprocess(image): image = image.copy()[:, :, ::-1] image = image.astype(np.float32) image = image / 128. - 1. - image = image.transpose((2, 0, 1)).astype(np.float32) - - return image + return image.transpose((2, 0, 1)).astype(np.float32) def inverse_vgg_preprocess(image): means = [0.485, 0.456, 0.406] @@ -70,9 +63,7 @@ def inverse_inception_preprocess(image): image = image.transpose((1, 2, 0)).astype(np.float32) image = image[:, :, ::-1] image = (image + 1.)*128. - image = image.astype(np.uint8) - - return image + return image.astype(np.uint8) def ssd_preprocess(image): image = image.astype(np.float32) @@ -93,9 +84,7 @@ def preprocess(image, mode): 'inception': inception_preprocess, 'ssd': ssd_preprocess } - if mode not in preprocessors: - return image - return preprocessors[mode](image) + return image if mode not in preprocessors else preprocessors[mode](image) def put_vec_maps(centerA, centerB, accumulate_vec_map, count, params_transform): @@ -132,8 +121,8 @@ def put_vec_maps(centerA, centerB, accumulate_vec_map, count, params_transform): min_y = max(int(round(min(centerA[1], centerB[1]) - thre)), 0) max_y = min(int(round(max(centerA[1], centerB[1]) + thre)), grid_y) - range_x = list(range(int(min_x), int(max_x), 1)) - range_y = list(range(int(min_y), int(max_y), 1)) + range_x = list(range(int(min_x), int(max_x))) + range_y = list(range(int(min_y), int(max_y))) xx, yy = np.meshgrid(range_x, range_y) ba_x = xx - centerA[0] # the vector from (x,y) to centerA ba_y = yy - centerA[1] diff --git a/lib/datasets/test_dataloader.py b/lib/datasets/test_dataloader.py index ebee503..4dd1a7e 100755 --- a/lib/datasets/test_dataloader.py +++ b/lib/datasets/test_dataloader.py @@ -84,7 +84,7 @@ def get_result(person_to_joint_assoc, joint_list): for part in range(18): index = int(person_to_joint_assoc[ridxPred, part]) - if -1 == index: + if index == -1: keypoints[part, 0] = 0 keypoints[part, 1] = 0 keypoints[part, 2] = 0 @@ -112,39 +112,36 @@ def get_result(person_to_joint_assoc, joint_list): ]) def kp_connections(keypoints): - kp_lines = [ - [keypoints.index('neck'), keypoints.index('right_hip')], + return [ + [keypoints.index('neck'), keypoints.index('right_hip')], [keypoints.index('right_hip'), keypoints.index('right_knee')], [keypoints.index('right_knee'), keypoints.index('right_ankle')], - [keypoints.index('neck'), keypoints.index('left_hip')], + [keypoints.index('neck'), keypoints.index('left_hip')], [keypoints.index('left_hip'), keypoints.index('left_knee')], [keypoints.index('left_knee'), keypoints.index('left_ankle')], - [keypoints.index('neck'), keypoints.index('right_shoulder')], + [keypoints.index('neck'), keypoints.index('right_shoulder')], [keypoints.index('right_shoulder'), keypoints.index('right_elbow')], - [keypoints.index('right_elbow'), keypoints.index('right_wrist')], - [keypoints.index('right_shoulder'), keypoints.index('right_eye')], - [keypoints.index('neck'), keypoints.index('left_shoulder')], + [keypoints.index('right_elbow'), keypoints.index('right_wrist')], + [keypoints.index('right_shoulder'), keypoints.index('right_eye')], + [keypoints.index('neck'), keypoints.index('left_shoulder')], [keypoints.index('left_shoulder'), keypoints.index('left_elbow')], [keypoints.index('left_elbow'), keypoints.index('left_wrist')], - [keypoints.index('left_shoulder'), keypoints.index('left_eye')], - [keypoints.index('neck'), keypoints.index('nose')], + [keypoints.index('left_shoulder'), keypoints.index('left_eye')], + [keypoints.index('neck'), keypoints.index('nose')], [keypoints.index('nose'), keypoints.index('right_eye')], - [keypoints.index('nose'), keypoints.index('left_eye')], + [keypoints.index('nose'), keypoints.index('left_eye')], [keypoints.index('right_eye'), keypoints.index('right_ear')], - [keypoints.index('left_eye'), keypoints.index('left_ear')] + [keypoints.index('left_eye'), keypoints.index('left_ear')], ] - return kp_lines def get_keypoints(): """Get the COCO keypoints and their left/right flip coorespondence map.""" - # Keypoints are not available in the COCO json for the test split, so we - # provide them here. - keypoints = [ + return [ 'nose', 'neck', 'right_shoulder', 'right_elbow', - 'right_wrist', + 'right_wrist', 'left_shoulder', 'left_elbow', 'left_wrist', @@ -154,12 +151,11 @@ def get_keypoints(): 'left_hip', 'left_knee', 'left_ankle', - 'right_eye', + 'right_eye', 'left_eye', 'right_ear', - 'left_ear'] - - return keypoints + 'left_ear', + ] limb_thickness = 4 @@ -216,7 +212,7 @@ def inverse_vgg_preprocess(image): train_loader = torch.utils.data.DataLoader( train_data, batch_size=cfg.TRAIN.BATCH_SIZE_PER_GPU*len(cfg.GPUS), shuffle=True, pin_memory=cfg.PIN_MEMORY, num_workers=cfg.WORKERS, drop_last=True) - + val_data = datasets.CocoKeypoints( root=cfg.DATASET.VAL_IMAGE_DIR, @@ -244,7 +240,7 @@ def inverse_vgg_preprocess(image): heatmap = heatmaps[index].cpu().numpy().transpose((1,2,0)) paf = pafs[index].cpu().numpy().transpose((1,2,0)) param = {'thre1': 0.1, 'thre2': 0.05, 'thre3': 0.5} - + joint_list, person_to_joint_assoc = paf_to_pose(heatmap, paf) keypoints = get_result(person_to_joint_assoc, joint_list) @@ -257,7 +253,7 @@ def inverse_vgg_preprocess(image): joint = keypoint[j] print(joint) if joint[2]==1: - cv2.circle(img, tuple(joint[0:2].astype(int)), 2, colors[j], thickness=-1) + cv2.circle(img, tuple(joint[:2].astype(int)), 2, colors[j], thickness=-1) for k in range(NUM_LIMBS): src_joint = keypoint[joint_to_limb_heatmap_relationship[k][0]] dst_joint = keypoint[joint_to_limb_heatmap_relationship[k][1]] @@ -276,7 +272,7 @@ def inverse_vgg_preprocess(image): # copying it polygon = cv2.ellipse2Poly( coords_center, (int(limb_length / 2), limb_thickness), int(angle), 0, 360, 1) - cv2.fillConvexPoly(img, polygon, colors[k]) + cv2.fillConvexPoly(img, polygon, colors[k]) # print(one_dict.keys()) #if output_picture_name: - cv2.imwrite(os.path.join(vis_dir, str(batch_idx)+'.png'), img) + cv2.imwrite(os.path.join(vis_dir, f'{str(batch_idx)}.png'), img) diff --git a/lib/datasets/utils.py b/lib/datasets/utils.py index 1082fa3..712319b 100644 --- a/lib/datasets/utils.py +++ b/lib/datasets/utils.py @@ -10,8 +10,7 @@ def horizontal_swap_coco(keypoints): for source_i, xyv in enumerate(keypoints): source_name = COCO_KEYPOINTS[source_i] - target_name = HFLIP.get(source_name) - if target_name: + if target_name := HFLIP.get(source_name): target_i = COCO_KEYPOINTS.index(target_name) else: target_i = source_i @@ -26,11 +25,13 @@ def create_sink(side): return np.zeros((2, 1, 1)) sink1d = np.linspace((side - 1.0) / 2.0, -(side - 1.0) / 2.0, num=side, dtype=np.float32) - sink = np.stack(( - sink1d.reshape(1, -1).repeat(side, axis=0), - sink1d.reshape(-1, 1).repeat(side, axis=1), - ), axis=0) - return sink + return np.stack( + ( + sink1d.reshape(1, -1).repeat(side, axis=0), + sink1d.reshape(-1, 1).repeat(side, axis=1), + ), + axis=0, + ) def mask_valid_area(intensities, valid_area): diff --git a/lib/network/atrous_model.py b/lib/network/atrous_model.py index 68470c0..68c1d0a 100755 --- a/lib/network/atrous_model.py +++ b/lib/network/atrous_model.py @@ -101,8 +101,7 @@ def forward(self, x): x_cat = torch.cat([x_1, x_2, x_3, x_4, x_5], 1) x_out = self.Mconv_6(x_cat) x_out = self.Mrelu_6(x_out) - x_pred = self.Mconv_7(x_out) - return x_pred + return self.Mconv_7(x_out) class feature_extractor(nn.Module): @@ -151,10 +150,20 @@ def __init__(self, stages=5, have_bn = True, have_bias = False): in_channels=256, out_channels=128, kernel_size=3, padding=1), nn.ReLU(inplace=True)) for i in range(stages): - setattr(self, 'PAF_stage{}'.format(i + 2), stage_block(in_channels=128, out_channels=38) if i == 0 else - stage_block(in_channels=185, out_channels=38)) - setattr(self, 'heatmap_stage{}'.format(i + 2), stage_block(in_channels=128, out_channels=19) if i == 0 else - stage_block(in_channels=185, out_channels=19)) + setattr( + self, + f'PAF_stage{i + 2}', + stage_block(in_channels=128, out_channels=38) + if i == 0 + else stage_block(in_channels=185, out_channels=38), + ) + setattr( + self, + f'heatmap_stage{i + 2}', + stage_block(in_channels=128, out_channels=19) + if i == 0 + else stage_block(in_channels=185, out_channels=19), + ) self.init_weight() def forward(self, x): @@ -163,11 +172,9 @@ def forward(self, x): x_in_0 = self.stage_0(x_in) x_in = x_in_0 for i in range(self.stages): - x_PAF_pred = getattr(self, 'PAF_stage{}'.format(i + 2))(x_in) - x_heatmap_pred = getattr( - self, 'heatmap_stage{}'.format(i + 2))(x_in) - saved_for_loss.append(x_PAF_pred) - saved_for_loss.append(x_heatmap_pred) + x_PAF_pred = getattr(self, f'PAF_stage{i + 2}')(x_in) + x_heatmap_pred = getattr(self, f'heatmap_stage{i + 2}')(x_in) + saved_for_loss.extend((x_PAF_pred, x_heatmap_pred)) if i != self.stages - 1: x_in = torch.cat([x_PAF_pred, x_heatmap_pred, x_in_0], 1) return [x_PAF_pred, x_heatmap_pred], saved_for_loss @@ -226,8 +233,7 @@ def build_names(): names = [] for j in range(1, 6): - for k in range(1, 3): - names.append('loss_stage%d_L%d' % (j, k)) + names.extend('loss_stage%d_L%d' % (j, k) for k in range(1, 3)) return names """Load pretrained model on Imagenet @@ -239,13 +245,13 @@ def use_inception(model): url = 'https://download.pytorch.org/models/inception_v3_google-1a9a5a14.pth' incep_state_dict = model_zoo.load_url(url) incep_keys = incep_state_dict.keys() - - # load weights of vgg - weights_load = {} - # weight+bias,weight+bias.....(repeat 10 times) - for i in range(60): - weights_load[list(model.state_dict().keys())[i]] = incep_state_dict[list(incep_keys)[i]] + weights_load = { + list(model.state_dict().keys())[i]: incep_state_dict[ + list(incep_keys)[i] + ] + for i in range(60) + } state = model.state_dict() state.update(weights_load) model.load_state_dict(state) diff --git a/lib/network/atrous_model_share_stages.py b/lib/network/atrous_model_share_stages.py index 86a8add..45c015f 100755 --- a/lib/network/atrous_model_share_stages.py +++ b/lib/network/atrous_model_share_stages.py @@ -175,26 +175,28 @@ def __init__(self, stages=5, have_bn=True, have_bias=False): in_channels=256, out_channels=128, kernel_size=3, padding=1), nn.ReLU(inplace=True)) for i in range(stages): - setattr(self, 'stage{}'.format(i + 2), stage_block(in_channels=128, out_channels=14) if i == 0 else - stage_block(in_channels=151, out_channels=14)) + setattr( + self, + f'stage{i + 2}', + stage_block(in_channels=128, out_channels=14) + if i == 0 + else stage_block(in_channels=151, out_channels=14), + ) self._initialize_weights_norm() def forward(self, x): - saved_for_loss = [] - paf_ret, heat_ret = [], [] + paf_ret, heat_ret = [], [] x_in = self.feature_extractor(x) x_in_0 = self.stage_0(x_in) x_in = x_in_0 for i in range(self.stages): - x_PAF_pred, x_heatmap_pred = getattr( - self, 'stage{}'.format(i + 2))(x_in) + x_PAF_pred, x_heatmap_pred = getattr(self, f'stage{i + 2}')(x_in) paf_ret.append(x_PAF_pred) heat_ret.append(x_heatmap_pred) if i != self.stages - 1: x_in = torch.cat([x_PAF_pred, x_heatmap_pred, x_in_0], 1) - saved_for_loss.append(paf_ret) - saved_for_loss.append(heat_ret) + saved_for_loss = [paf_ret, heat_ret] return [(paf_ret[-2], heat_ret[-2]), (paf_ret[-1], heat_ret[-1])], saved_for_loss def _initialize_weights_norm(self): @@ -251,23 +253,21 @@ def use_inception(model, model_path): incep_state_dict = model_zoo.load_url(url, model_dir=model_path) incep_keys = incep_state_dict.keys() - # load weights of vgg - weights_load = {} - # weight+bias,weight+bias.....(repeat 10 times) - for i in range(60): - weights_load[list(model.state_dict().keys())[i] - ] = incep_state_dict[list(incep_keys)[i]] - + weights_load = { + list(model.state_dict().keys())[i]: incep_state_dict[ + list(incep_keys)[i] + ] + for i in range(60) + } state = model.state_dict() state.update(weights_load) model.load_state_dict(state) - print('load imagenet pretrained model: {}'.format(model_path)) + print(f'load imagenet pretrained model: {model_path}') def build_names(): names = [] for j in range(1, 6): - for k in range(1, 3): - names.append('loss_stage%d_L%d' % (j, k)) + names.extend('loss_stage%d_L%d' % (j, k) for k in range(1, 3)) return names diff --git a/lib/network/atrouspose.py b/lib/network/atrouspose.py index 70c455e..2e5b6ef 100644 --- a/lib/network/atrouspose.py +++ b/lib/network/atrouspose.py @@ -44,9 +44,7 @@ def forward(self, feature_map): add2 = add1+out_3x3_1 add3 = add2+out_3x3_2 add4 = add3+out_3x3_3 - out = F.relu(self.bn_out(torch.cat([add1, add2, add3, add4], 1))) # (shape: (batch_size, 1280, h/16, w/16)) - - return out + return F.relu(self.bn_out(torch.cat([add1, add2, add3, add4], 1))) def conv_dw(in_channels, out_channels, kernel_size=3, padding=1, stride=1, dilation=1): @@ -98,8 +96,7 @@ def __init__(self, paf_out_channels=38, heat_out_channels=19): ) def _lateral(self, input_size, factor): - layers = [] - layers.append(nn.Conv2d(input_size, 256, kernel_size=1, stride=1, bias=False)) + layers = [nn.Conv2d(input_size, 256, kernel_size=1, stride=1, bias=False)] layers.append(nn.BatchNorm2d(256)) layers.append(nn.ReLU(inplace=True)) layers.append(Upsample(scale_factor=factor, mode='bilinear')) @@ -141,7 +138,7 @@ def model_info(model): # Plots a line-by-line description of a PyTorch model image = torch.from_numpy(image).type(torch.FloatTensor).permute(2, 0, 1).reshape(1, 3, 256, 256).cuda() print(net) model_info(net) - for i in range(30): + for _ in range(30): vec1, heat1 = net(image) print(vec1.shape, heat1.shape) diff --git a/lib/network/im_transform.py b/lib/network/im_transform.py index ea74b88..c22afe6 100644 --- a/lib/network/im_transform.py +++ b/lib/network/im_transform.py @@ -5,7 +5,7 @@ def resize(frame, desired_size): old_size = frame.shape[:2] ratio = float(desired_size) / max(old_size) - new_size = tuple([int(x*ratio) for x in old_size]) + new_size = tuple(int(x*ratio) for x in old_size) frame = cv2.resize(frame, (new_size[1], new_size[0])) delta_w = desired_size - new_size[1] @@ -89,21 +89,21 @@ def offset_boxes(boxes, scale, offs, flip, im_shape): boxes[:, 0::2] -= offs[0] boxes[:, 1::2] -= offs[1] - is_box = boxes.shape[-1] % 4 == 0 # if is_box: # boxes = clip_boxes(boxes, im_shape) if flip: boxes[:, 0::2] = im_shape[1] - boxes[:, 0::2] + is_box = boxes.shape[-1] % 4 == 0 if is_box: for i in range(boxes.shape[-1] // 4): tmp = boxes[:, i].copy() boxes[:, i] = boxes[:, i + 2] boxes[:, i + 2] = tmp - # boxes_x = np.copy(boxes[:, 0]) - # boxes[:, 0] = im_shape[1] - boxes[:, 2] - # boxes[:, 2] = im_shape[1] - boxes_x + # boxes_x = np.copy(boxes[:, 0]) + # boxes[:, 0] = im_shape[1] - boxes[:, 2] + # boxes[:, 2] = im_shape[1] - boxes_x if expand: boxes = boxes[0] @@ -118,8 +118,8 @@ def _factor_closest(num, factor, is_ceil=True): def crop_with_factor(im, dest_size=None, factor=32, is_ceil=True): im_shape = im.shape - im_size_min = np.min(im_shape[0:2]) - im_size_max = np.max(im_shape[0:2]) + im_size_min = np.min(im_shape[:2]) + im_size_max = np.max(im_shape[:2]) # im_scale = 1. # if max_size is not None and im_size_min > max_size: im_scale = float(dest_size) / im_size_min diff --git a/lib/network/openpose.py b/lib/network/openpose.py index 73587df..8a14660 100644 --- a/lib/network/openpose.py +++ b/lib/network/openpose.py @@ -31,8 +31,8 @@ def make_vgg19_block(): {'conv4_3_CPM': [512, 256, 3, 1, 1]}, {'conv4_4_CPM': [256, 128, 3, 1, 1]}] layers = [] - for i in range(len(block)): - one_ = block[i] + for item in block: + one_ = item for k, v in one_.items(): if 'pool' in k: layers += [nn.MaxPool2d(kernel_size=v[0], stride=v[1], @@ -105,8 +105,7 @@ def forward(self, x): out3_5 = self.Mconv5_2(out2_5) x_cat_5 = torch.cat([out1_5, out2_5, out3_5], 1) out_6 = self.Mconv6(x_cat_5) - stage_output = self.Mconv7(out_6) - return stage_output + return self.Mconv7(out_6) class OpenPose_Model(nn.Module): def __init__(self, l2_stages=4, l1_stages=2, @@ -158,7 +157,6 @@ def __init__(self, l2_stages=4, l1_stages=2, self._initialize_weights_norm() def forward(self, x): - saved_for_loss = [] features = self.feature_extractor(x) paf_ret, heat_ret = [], [] x_in = features @@ -172,8 +170,7 @@ def forward(self, x): heat_pred = l1_stage(x_in) x_in = torch.cat([features, heat_pred, paf_pred], 1) heat_ret.append(heat_pred) - saved_for_loss.append(paf_ret) - saved_for_loss.append(heat_ret) + saved_for_loss = [paf_ret, heat_ret] return [(paf_ret[-2], heat_ret[-2]), (paf_ret[-1], heat_ret[-1])], saved_for_loss def _initialize_weights_norm(self): @@ -215,13 +212,10 @@ def use_vgg(model): vgg_state_dict = model_zoo.load_url(url) vgg_keys = vgg_state_dict.keys() - # load weights of vgg - weights_load = {} - # weight+bias,weight+bias.....(repeat 10 times) - for i in range(20): - weights_load[list(model.state_dict().keys())[i] - ] = vgg_state_dict[list(vgg_keys)[i]] - + weights_load = { + list(model.state_dict().keys())[i]: vgg_state_dict[list(vgg_keys)[i]] + for i in range(20) + } state = model.state_dict() state.update(weights_load) model.load_state_dict(state) diff --git a/lib/network/post.py b/lib/network/post.py index 23c91bd..54a9664 100644 --- a/lib/network/post.py +++ b/lib/network/post.py @@ -153,8 +153,11 @@ def NMS(param, heatmaps, upsampFactor=1., bool_refine_center=True, bool_gaussian refined_center = [0, 0] # Flip peak coordinates since they are [x,y] instead of [y,x] peak_score = map_orig[tuple(peak[::-1])] - peaks[i, :] = tuple([int(round(x)) for x in compute_resized_coords( - peak_coords[i], upsampFactor) + refined_center[::-1]]) + (peak_score, cnt_total_joints) + peaks[i, :] = tuple( + int(round(x)) + for x in compute_resized_coords(peak_coords[i], upsampFactor) + + refined_center[::-1] + ) + (peak_score, cnt_total_joints) cnt_total_joints += 1 joint_list_per_joint_type.append(peaks) @@ -250,7 +253,7 @@ def find_connected_joints(param, paf_upsamp, joint_list_per_joint_type, num_inte max_connections = min(len(joints_src), len(joints_dst)) # Traverse all potential joint connections (sorted by their score) for potential_connection in connection_candidates: - i, j, s = potential_connection[0:3] + i, j, s = potential_connection[:3] # Make sure joints_src[i] or joints_dst[j] haven't already been # connected to other joints_dst or joints_src if i not in connections[:, 3] and j not in connections[:, 4]: @@ -284,11 +287,12 @@ def group_limbs_of_same_person(connected_limbs, joint_list): joint_src_type, joint_dst_type = joint_to_limb_heatmap_relationship[limb_type] for limb_info in connected_limbs[limb_type]: - person_assoc_idx = [] - for person, person_limbs in enumerate(person_to_joint_assoc): - if person_limbs[joint_src_type] == limb_info[0] or person_limbs[joint_dst_type] == limb_info[1]: - person_assoc_idx.append(person) - + person_assoc_idx = [ + person + for person, person_limbs in enumerate(person_to_joint_assoc) + if person_limbs[joint_src_type] == limb_info[0] + or person_limbs[joint_dst_type] == limb_info[1] + ] # If one of the joints has been associated to a person, and either # the other joint is also associated with the same person or not # associated to anyone yet: @@ -337,11 +341,11 @@ def group_limbs_of_same_person(connected_limbs, joint_list): ) + limb_info[2] person_to_joint_assoc.append(row) - # Delete people who have very few parts connected - people_to_delete = [] - for person_id, person_info in enumerate(person_to_joint_assoc): - if person_info[-1] < 3 or person_info[-2] / person_info[-1] < 0.2: - people_to_delete.append(person_id) + people_to_delete = [ + person_id + for person_id, person_info in enumerate(person_to_joint_assoc) + if person_info[-1] < 3 or person_info[-2] / person_info[-1] < 0.2 + ] # Traverse the list in reverse order so we delete indices starting from the # last one (otherwise, removing item for example 0 would modify the indices of # the remaining people to be deleted!) @@ -377,10 +381,15 @@ def plot_pose(img_orig, joint_list, person_to_joint_assoc, bool_fast_plot=True, # joint_coords[:,0] represents Y coords of both joints; # joint_coords[:,1], X coords joint_coords = joint_list[joint_indices, 0:2] - + for joint in joint_coords: # Draw circles at every joint - cv2.circle(canvas, tuple(joint[0:2].astype( - int)), 4, (255,255,255), thickness=-1) + cv2.circle( + canvas, + tuple(joint[:2].astype(int)), + 4, + (255, 255, 255), + thickness=-1, + ) # mean along the axis=0 computes meanYcoord and meanXcoord -> Round # and make int to avoid errors coords_center = tuple( diff --git a/lib/network/rtpose_hourglass.py b/lib/network/rtpose_hourglass.py index d22e949..42ec74c 100644 --- a/lib/network/rtpose_hourglass.py +++ b/lib/network/rtpose_hourglass.py @@ -55,17 +55,13 @@ def __init__(self, block, num_blocks, planes, depth): self.hg = self._make_hour_glass(block, num_blocks, planes, depth) def _make_residual(self, block, num_blocks, planes): - layers = [] - for i in range(0, num_blocks): - layers.append(block(planes * block.expansion, planes)) + layers = [block(planes * block.expansion, planes) for _ in range(num_blocks)] return nn.Sequential(*layers) def _make_hour_glass(self, block, num_blocks, planes, depth): hg = [] for i in range(depth): - res = [] - for j in range(3): - res.append(self._make_residual(block, num_blocks, planes)) + res = [self._make_residual(block, num_blocks, planes) for _ in range(3)] if i == 0: res.append(self._make_residual(block, num_blocks, planes)) hg.append(nn.ModuleList(res)) @@ -82,8 +78,7 @@ def _hour_glass_forward(self, n, x): low2 = self.hg[n - 1][3](low1) low3 = self.hg[n - 1][2](low2) up2 = self.upsample(low3) - out = up1 + up2 - return out + return up1 + up2 def forward(self, x): return self._hour_glass_forward(self.depth, x) @@ -142,12 +137,9 @@ def _make_residual(self, block, planes, blocks, stride=1): kernel_size=1, stride=stride, bias=True), ) - layers = [] - layers.append(block(self.inplanes, planes, stride, downsample)) + layers = [block(self.inplanes, planes, stride, downsample)] self.inplanes = planes * block.expansion - for i in range(1, blocks): - layers.append(block(self.inplanes, planes)) - + layers.extend(block(self.inplanes, planes) for _ in range(1, blocks)) return nn.Sequential(*layers) def _make_fc(self, inplanes, outplanes): @@ -160,7 +152,6 @@ def _make_fc(self, inplanes, outplanes): ) def forward(self, x): - saved_for_loss = [] out = [] x = self.conv1(x) x = self.bn1(x) @@ -182,10 +173,8 @@ def forward(self, x): paf_score_ = self.paf_score_[i](score_paf) ht_score_ = self.ht_score_[i](score_ht) x = x + fc_ + paf_score_ + ht_score_ - - saved_for_loss.append(score_paf) - saved_for_loss.append(score_ht) + saved_for_loss = [score_paf, score_ht] return (score_paf, score_ht), saved_for_loss def _initialize_weights_norm(self): @@ -199,7 +188,10 @@ def _initialize_weights_norm(self): m.bias.data.zero_() def hg(**kwargs): - model = HourglassNet(Bottleneck, num_stacks=kwargs['num_stacks'], - num_blocks=kwargs['num_blocks'], paf_classes=kwargs['paf_classes'], - ht_classes=kwargs['ht_classes']) - return model + return HourglassNet( + Bottleneck, + num_stacks=kwargs['num_stacks'], + num_blocks=kwargs['num_blocks'], + paf_classes=kwargs['paf_classes'], + ht_classes=kwargs['ht_classes'], + ) diff --git a/lib/network/rtpose_mobilenetV2.py b/lib/network/rtpose_mobilenetV2.py index a9b3900..5f9f1c4 100644 --- a/lib/network/rtpose_mobilenetV2.py +++ b/lib/network/rtpose_mobilenetV2.py @@ -53,10 +53,7 @@ def __init__(self, inp, oup, stride, expand_ratio): ) def forward(self, x): - if self.use_res_connect: - return x + self.conv(x) - else: - return self.conv(x) + return x + self.conv(x) if self.use_res_connect else self.conv(x) class MobileNetV2(nn.Module): diff --git a/lib/network/rtpose_shufflenetV2.py b/lib/network/rtpose_shufflenetV2.py index 50c3fa2..40b6090 100644 --- a/lib/network/rtpose_shufflenetV2.py +++ b/lib/network/rtpose_shufflenetV2.py @@ -31,27 +31,48 @@ def __init__(self, name, in_channels, out_channels, stride, downsample, dilation if not self.downsample and self.stride==1: assert in_channels == out_channels self.conv = nn.Sequential( - slim.conv_bn_relu(name + '/conv1', channels, channels, 1), - slim.conv_bn(name + '/conv2', - channels, channels, 3, stride=stride, - dilation=dilation, padding=dilation, groups=channels), - slim.conv_bn_relu(name + '/conv3', channels, channels, 1), + slim.conv_bn_relu(f'{name}/conv1', channels, channels, 1), + slim.conv_bn( + f'{name}/conv2', + channels, + channels, + 3, + stride=stride, + dilation=dilation, + padding=dilation, + groups=channels, + ), + slim.conv_bn_relu(f'{name}/conv3', channels, channels, 1), ) else: self.conv = nn.Sequential( - slim.conv_bn_relu(name + '/conv1', in_channels, channels, 1), - slim.conv_bn(name + '/conv2', - channels, channels, 3, stride=stride, - dilation=dilation, padding=dilation, groups=channels), - slim.conv_bn_relu(name + '/conv3', channels, channels, 1), + slim.conv_bn_relu(f'{name}/conv1', in_channels, channels, 1), + slim.conv_bn( + f'{name}/conv2', + channels, + channels, + 3, + stride=stride, + dilation=dilation, + padding=dilation, + groups=channels, + ), + slim.conv_bn_relu(f'{name}/conv3', channels, channels, 1), ) self.conv0 = nn.Sequential( - slim.conv_bn(name + '/conv4', - in_channels, in_channels, 3, stride=stride, - dilation=dilation, padding=dilation, groups=in_channels), - slim.conv_bn_relu(name + '/conv5', in_channels, channels, 1), + slim.conv_bn( + f'{name}/conv4', + in_channels, + in_channels, + 3, + stride=stride, + dilation=dilation, + padding=dilation, + groups=in_channels, + ), + slim.conv_bn_relu(f'{name}/conv5', in_channels, channels, 1), ) - self.shuffle = slim.channel_shuffle(name + '/shuffle', 2) + self.shuffle = slim.channel_shuffle(f'{name}/shuffle', 2) def forward(self, x): if not self.downsample: @@ -65,14 +86,14 @@ def forward(self, x): def generate_caffe_prototxt(self, caffe_net, layer): if self.stride == 1: layer_x1, layer_x2 = L.Slice(layer, ntop=2, axis=1, slice_point=[self.in_channels//2]) - caffe_net[self.g_name + '/slice1'] = layer_x1 - caffe_net[self.g_name + '/slice2'] = layer_x2 + caffe_net[f'{self.g_name}/slice1'] = layer_x1 + caffe_net[f'{self.g_name}/slice2'] = layer_x2 layer_x2 = slim.generate_caffe_prototxt(self.conv, caffe_net, layer_x2) else: layer_x1 = slim.generate_caffe_prototxt(self.conv0, caffe_net, layer) layer_x2 = slim.generate_caffe_prototxt(self.conv, caffe_net, layer) layer = L.Concat(layer_x1, layer_x2, axis=1) - caffe_net[self.g_name + '/concat'] = layer + caffe_net[f'{self.g_name}/concat'] = layer layer = slim.generate_caffe_prototxt(self.shuffle, caffe_net, layer) return layer @@ -112,12 +133,28 @@ def __init__(self, width_multiplier): out_channels, stride, dilation, num_blocks, stage_type = config if stride==2: downsample=True - stage_prefix = 'stage_{}'.format(i - 1) - blocks = [BasicBlock(stage_prefix + '_1', in_channels, - out_channels, stride, downsample, dilation)] - for i in range(1, num_blocks): - blocks.append(BasicBlock(stage_prefix + '_{}'.format(i + 1), - out_channels, out_channels, 1, False, dilation)) + stage_prefix = f'stage_{i - 1}' + blocks = [ + BasicBlock( + f'{stage_prefix}_1', + in_channels, + out_channels, + stride, + downsample, + dilation, + ) + ] + blocks.extend( + BasicBlock( + f'{stage_prefix}_{i + 1}', + out_channels, + out_channels, + 1, + False, + dilation, + ) + for i in range(1, num_blocks) + ) self.network += [nn.Sequential(*blocks)] in_channels = out_channels @@ -158,11 +195,11 @@ def convert_to_caffe(self, name): caffe_net.tops['data'] = layer slim.generate_caffe_prototxt(self, caffe_net, layer) print(caffe_net.to_proto()) - with open(name + '.prototxt', 'wb') as f: + with open(f'{name}.prototxt', 'wb') as f: f.write(str(caffe_net.to_proto()).encode()) - caffe_net = caffe.Net(name + '.prototxt', caffe.TEST) + caffe_net = caffe.Net(f'{name}.prototxt', caffe.TEST) slim.convert_pytorch_to_caffe(self, caffe_net) - caffe_net.save(name + '.caffemodel') + caffe_net.save(f'{name}.caffemodel') if __name__ == '__main__': diff --git a/lib/network/rtpose_vgg.py b/lib/network/rtpose_vgg.py index fc34214..47a33df 100644 --- a/lib/network/rtpose_vgg.py +++ b/lib/network/rtpose_vgg.py @@ -63,9 +63,17 @@ def get_model(trunk='vgg19'): trunk: string, 'vgg19' or 'mobilenet' Returns: Module, the defined model """ - blocks = {} # block0 is the preprocessing stage - if trunk == 'vgg19': + if trunk == 'mobilenet': + block0 = [{'conv_bn': [3, 32, 2]}, # out: 3, 32, 184, 184 + {'conv_dw1': [32, 64, 1]}, # out: 32, 64, 184, 184 + {'conv_dw2': [64, 128, 2]}, # out: 64, 128, 92, 92 + {'conv_dw3': [128, 128, 1]}, # out: 128, 256, 92, 92 + {'conv_dw4': [128, 256, 2]}, # out: 256, 256, 46, 46 + {'conv4_3_CPM': [256, 256, 1, 3, 1]}, + {'conv4_4_CPM': [256, 128, 1, 3, 1]}] + + elif trunk == 'vgg19': block0 = [{'conv1_1': [3, 64, 3, 1, 1]}, {'conv1_2': [64, 64, 3, 1, 1]}, {'pool1_stage1': [2, 2, 0]}, @@ -82,28 +90,22 @@ def get_model(trunk='vgg19'): {'conv4_3_CPM': [512, 256, 3, 1, 1]}, {'conv4_4_CPM': [256, 128, 3, 1, 1]}] - elif trunk == 'mobilenet': - block0 = [{'conv_bn': [3, 32, 2]}, # out: 3, 32, 184, 184 - {'conv_dw1': [32, 64, 1]}, # out: 32, 64, 184, 184 - {'conv_dw2': [64, 128, 2]}, # out: 64, 128, 92, 92 - {'conv_dw3': [128, 128, 1]}, # out: 128, 256, 92, 92 - {'conv_dw4': [128, 256, 2]}, # out: 256, 256, 46, 46 - {'conv4_3_CPM': [256, 256, 1, 3, 1]}, - {'conv4_4_CPM': [256, 128, 1, 3, 1]}] - - # Stage 1 - blocks['block1_1'] = [{'conv5_1_CPM_L1': [128, 128, 3, 1, 1]}, - {'conv5_2_CPM_L1': [128, 128, 3, 1, 1]}, - {'conv5_3_CPM_L1': [128, 128, 3, 1, 1]}, - {'conv5_4_CPM_L1': [128, 512, 1, 1, 0]}, - {'conv5_5_CPM_L1': [512, 38, 1, 1, 0]}] - - blocks['block1_2'] = [{'conv5_1_CPM_L2': [128, 128, 3, 1, 1]}, - {'conv5_2_CPM_L2': [128, 128, 3, 1, 1]}, - {'conv5_3_CPM_L2': [128, 128, 3, 1, 1]}, - {'conv5_4_CPM_L2': [128, 512, 1, 1, 0]}, - {'conv5_5_CPM_L2': [512, 19, 1, 1, 0]}] - + blocks = { + 'block1_1': [ + {'conv5_1_CPM_L1': [128, 128, 3, 1, 1]}, + {'conv5_2_CPM_L1': [128, 128, 3, 1, 1]}, + {'conv5_3_CPM_L1': [128, 128, 3, 1, 1]}, + {'conv5_4_CPM_L1': [128, 512, 1, 1, 0]}, + {'conv5_5_CPM_L1': [512, 38, 1, 1, 0]}, + ], + 'block1_2': [ + {'conv5_1_CPM_L2': [128, 128, 3, 1, 1]}, + {'conv5_2_CPM_L2': [128, 128, 3, 1, 1]}, + {'conv5_3_CPM_L2': [128, 128, 3, 1, 1]}, + {'conv5_4_CPM_L2': [128, 512, 1, 1, 0]}, + {'conv5_5_CPM_L2': [512, 19, 1, 1, 0]}, + ], + } # Stages 2 - 6 for i in range(2, 7): blocks['block%d_1' % i] = [ @@ -135,6 +137,8 @@ def get_model(trunk='vgg19'): for k, v in blocks.items(): models[k] = make_stages(list(v)) + + class rtpose_model(nn.Module): def __init__(self, model_dict): super(rtpose_model, self).__init__() @@ -157,44 +161,31 @@ def __init__(self, model_dict): def forward(self, x): - saved_for_loss = [] out1 = self.model0(x) out1_1 = self.model1_1(out1) out1_2 = self.model1_2(out1) out2 = torch.cat([out1_1, out1_2, out1], 1) - saved_for_loss.append(out1_1) - saved_for_loss.append(out1_2) - + saved_for_loss = [out1_1, out1_2] out2_1 = self.model2_1(out2) out2_2 = self.model2_2(out2) out3 = torch.cat([out2_1, out2_2, out1], 1) - saved_for_loss.append(out2_1) - saved_for_loss.append(out2_2) - + saved_for_loss.extend((out2_1, out2_2)) out3_1 = self.model3_1(out3) out3_2 = self.model3_2(out3) out4 = torch.cat([out3_1, out3_2, out1], 1) - saved_for_loss.append(out3_1) - saved_for_loss.append(out3_2) - + saved_for_loss.extend((out3_1, out3_2)) out4_1 = self.model4_1(out4) out4_2 = self.model4_2(out4) out5 = torch.cat([out4_1, out4_2, out1], 1) - saved_for_loss.append(out4_1) - saved_for_loss.append(out4_2) - + saved_for_loss.extend((out4_1, out4_2)) out5_1 = self.model5_1(out5) out5_2 = self.model5_2(out5) out6 = torch.cat([out5_1, out5_2, out1], 1) - saved_for_loss.append(out5_1) - saved_for_loss.append(out5_2) - + saved_for_loss.extend((out5_1, out5_2)) out6_1 = self.model6_1(out6) out6_2 = self.model6_2(out6) - saved_for_loss.append(out6_1) - saved_for_loss.append(out6_2) - + saved_for_loss.extend((out6_1, out6_2)) return (out6_1, out6_2), saved_for_loss def _initialize_weights_norm(self): @@ -221,8 +212,8 @@ def _initialize_weights_norm(self): init.normal_(self.model5_2[12].weight, std=0.01) init.normal_(self.model6_2[12].weight, std=0.01) - model = rtpose_model(models) - return model + + return rtpose_model(models) """Load pretrained model on Imagenet @@ -238,13 +229,10 @@ def use_vgg(model): vgg_state_dict = model_zoo.load_url(url) vgg_keys = vgg_state_dict.keys() - # load weights of vgg - weights_load = {} - # weight+bias,weight+bias.....(repeat 10 times) - for i in range(20): - weights_load[list(model.state_dict().keys())[i] - ] = vgg_state_dict[list(vgg_keys)[i]] - + weights_load = { + list(model.state_dict().keys())[i]: vgg_state_dict[list(vgg_keys)[i]] + for i in range(20) + } state = model.state_dict() state.update(weights_load) model.load_state_dict(state) diff --git a/lib/utils/common.py b/lib/utils/common.py index 41241d3..fa30d55 100644 --- a/lib/utils/common.py +++ b/lib/utils/common.py @@ -64,7 +64,7 @@ def part_count(self): return len(self.body_parts.keys()) def get_max_score(self): - return max([x.score for _, x in self.body_parts.items()]) + return max(x.score for _, x in self.body_parts.items()) def get_face_box(self, img_w, img_h, mode=0): """ @@ -101,9 +101,8 @@ def get_face_box(self, img_w, img_h, mode=0): size = max(size, img_w * math.sqrt((part_reye.x - part_leye.x) ** 2 + (part_reye.y - part_leye.y) ** 2) * 2.0) - if mode == 1: - if not is_reye and not is_leye: - return None + if mode == 1 and not is_reye and not is_leye: + return None is_rear, part_rear = _include_part(parts, _REar) is_lear, part_lear = _include_part(parts, _LEar) @@ -154,7 +153,7 @@ def get_upper_body_box(self, img_w, img_h): :return: """ - if not (img_w > 0 and img_h > 0): + if img_w <= 0 or img_h <= 0: raise Exception("img size should be positive") _NOSE = CocoPart.Nose.value @@ -170,10 +169,10 @@ def get_upper_body_box(self, img_w, img_h): return None # Initial Bounding Box - x = min([part[0] for part in part_coords]) - y = min([part[1] for part in part_coords]) - x2 = max([part[0] for part in part_coords]) - y2 = max([part[1] for part in part_coords]) + x = min(part[0] for part in part_coords) + y = min(part[1] for part in part_coords) + x2 = max(part[0] for part in part_coords) + y2 = max(part[1] for part in part_coords) # # ------ Adjust heuristically + # if face points are detcted, adjust y value @@ -194,11 +193,11 @@ def get_upper_body_box(self, img_w, img_h): x -= dx x2 += dx elif is_neck: - if is_lshoulder and not is_rshoulder: + if is_lshoulder: half_w = abs(part_lshoulder.x - part_neck.x) * img_w * 1.15 x = min(part_neck.x * img_w - half_w, x) x2 = max(part_neck.x * img_w + half_w, x2) - elif not is_lshoulder and is_rshoulder: + elif is_rshoulder: half_w = abs(part_rshoulder.x - part_neck.x) * img_w * 1.15 x = min(part_neck.x * img_w - half_w, x) x2 = max(part_neck.x * img_w + half_w, x2) diff --git a/lib/utils/paf_to_pose.py b/lib/utils/paf_to_pose.py index 1c43e10..b6ff106 100644 --- a/lib/utils/paf_to_pose.py +++ b/lib/utils/paf_to_pose.py @@ -137,8 +137,9 @@ def NMS(heatmaps, upsampFactor=1., bool_refine_center=True, bool_gaussian_filt=F # Flip peak coordinates since they are [x,y] instead of [y,x] peak_score = map_orig[tuple(peak[::-1])] peaks[i, :] = tuple( - x for x in compute_resized_coords(peak_coords[i], upsampFactor) + refined_center[::-1]) + ( - peak_score, cnt_total_joints) + compute_resized_coords(peak_coords[i], upsampFactor) + + refined_center[::-1] + ) + (peak_score, cnt_total_joints) cnt_total_joints += 1 joint_list_per_joint_type.append(peaks) @@ -239,7 +240,7 @@ def find_connected_joints(paf_upsamp, joint_list_per_joint_type, num_intermed_pt max_connections = min(len(joints_src), len(joints_dst)) # Traverse all potential joint connections (sorted by their score) for potential_connection in connection_candidates: - i, j, s = potential_connection[0:3] + i, j, s = potential_connection[:3] # Make sure joints_src[i] or joints_dst[j] haven't already been # connected to other joints_dst or joints_src if i not in connections[:, 3] and j not in connections[:, 4]: @@ -273,11 +274,12 @@ def group_limbs_of_same_person(connected_limbs, joint_list, config): joint_src_type, joint_dst_type = joint_to_limb_heatmap_relationship[limb_type] for limb_info in connected_limbs[limb_type]: - person_assoc_idx = [] - for person, person_limbs in enumerate(person_to_joint_assoc): - if person_limbs[joint_src_type] == limb_info[0] or person_limbs[joint_dst_type] == limb_info[1]: - person_assoc_idx.append(person) - + person_assoc_idx = [ + person + for person, person_limbs in enumerate(person_to_joint_assoc) + if person_limbs[joint_src_type] == limb_info[0] + or person_limbs[joint_dst_type] == limb_info[1] + ] # If one of the joints has been associated to a person, and either # the other joint is also associated with the same person or not # associated to anyone yet: @@ -326,11 +328,11 @@ def group_limbs_of_same_person(connected_limbs, joint_list, config): ) + limb_info[2] person_to_joint_assoc.append(row) - # Delete people who have very few parts connected - people_to_delete = [] - for person_id, person_info in enumerate(person_to_joint_assoc): - if person_info[-1] < 3 or person_info[-2] / person_info[-1] < 0.2: - people_to_delete.append(person_id) + people_to_delete = [ + person_id + for person_id, person_info in enumerate(person_to_joint_assoc) + if person_info[-1] < 3 or person_info[-2] / person_info[-1] < 0.2 + ] # Traverse the list in reverse order so we delete indices starting from the # last one (otherwise, removing item for example 0 would modify the indices of # the remaining people to be deleted!) diff --git a/train/train_SH.py b/train/train_SH.py index ef64641..398076f 100644 --- a/train/train_SH.py +++ b/train/train_SH.py @@ -17,64 +17,58 @@ # Hyper-params parser = argparse.ArgumentParser(description='PyTorch rtpose Training') parser.add_argument('--data_dir', default='/data/coco/images', type=str, metavar='DIR', - help='path to where coco images stored') + help='path to where coco images stored') parser.add_argument('--mask_dir', default='/data/coco/', type=str, metavar='DIR', - help='path to where coco images stored') + help='path to where coco images stored') parser.add_argument('--logdir', default='/extra/tensorboy', type=str, metavar='DIR', - help='path to where tensorboard log restore') + help='path to where tensorboard log restore') parser.add_argument('--json_path', default='/data/coco/COCO.json', type=str, metavar='PATH', help='path to where coco images stored') parser.add_argument('--model_path', default='./network/weight/', type=str, metavar='DIR', help='path to where the model saved') - + parser.add_argument('--lr', '--learning-rate', default=1., type=float, metavar='LR', help='initial learning rate') parser.add_argument('--momentum', default=0.9, type=float, metavar='M', help='momentum') - + parser.add_argument('--epochs', default=200, type=int, metavar='N', help='number of total epochs to run') - + parser.add_argument('--weight-decay', '--wd', default=0.000, type=float, - metavar='W', help='weight decay (default: 1e-4)') + metavar='W', help='weight decay (default: 1e-4)') parser.add_argument('--nesterov', dest='nesterov', action='store_true') - + parser.add_argument('-o', '--optim', default='sgd', type=str) #Device options parser.add_argument('--gpu_ids', dest='gpu_ids', help='which gpu to use', nargs="+", default=[0,1,2,3], type=int) - + parser.add_argument('-b', '--batch_size', default=80, type=int, metavar='N', help='mini-batch size (default: 256)') parser.add_argument('--print_freq', default=20, type=int, metavar='N', help='number of iterations to print the training statistics') -from tensorboardX import SummaryWriter +from tensorboardX import SummaryWriter args = parser.parse_args() - -os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(str(e) for e in args.gpu_ids) -params_transform = dict() -params_transform['mode'] = 5 -# === aug_scale === -params_transform['scale_min'] = 0.5 -params_transform['scale_max'] = 1.1 -params_transform['scale_prob'] = 1 -params_transform['target_dist'] = 0.6 -# === aug_rotate === -params_transform['max_rotate_degree'] = 40 - -# === -params_transform['center_perterb_max'] = 40 - -# === aug_flip === -params_transform['flip_prob'] = 0.5 +os.environ['CUDA_VISIBLE_DEVICES'] = ','.join(str(e) for e in args.gpu_ids) -params_transform['np'] = 56 -params_transform['sigma'] = 4.416 -params_transform['limb_width'] = 1.289 +params_transform = { + 'mode': 5, + 'scale_min': 0.5, + 'scale_max': 1.1, + 'scale_prob': 1, + 'target_dist': 0.6, + 'max_rotate_degree': 40, + 'center_perterb_max': 40, + 'flip_prob': 0.5, + 'np': 56, + 'sigma': 4.416, + 'limb_width': 1.289, +} def get_loss(saved_for_loss, heat_temp, heat_weight, @@ -130,17 +124,16 @@ def train(train_loader, model, optimizer, epoch): batch_time = AverageMeter() data_time = AverageMeter() losses = AverageMeter() - - meter_dict = {} - meter_dict['paf'] = AverageMeter() - meter_dict['heatmap'] = AverageMeter() + meter_dict = {'paf': AverageMeter()} + + meter_dict['heatmap'] = AverageMeter() meter_dict['max_ht'] = AverageMeter() - meter_dict['min_ht'] = AverageMeter() - meter_dict['max_paf'] = AverageMeter() + meter_dict['min_ht'] = AverageMeter() + meter_dict['max_paf'] = AverageMeter() meter_dict['min_paf'] = AverageMeter() - - + + # switch to train mode model.train() @@ -148,7 +141,7 @@ def train(train_loader, model, optimizer, epoch): for i, (img, heatmap_target, heat_mask, paf_target, paf_mask) in enumerate(train_loader): # measure data loading time #writer.add_text('Text', 'text logged at step:' + str(i), i) - + #for name, param in model.named_parameters(): # writer.add_histogram(name, param.clone().cpu().data.numpy(),i) data_time.update(time.time() - end) @@ -158,14 +151,14 @@ def train(train_loader, model, optimizer, epoch): heat_mask = heat_mask.cuda() paf_target = paf_target.cuda() paf_mask = paf_mask.cuda() - + # compute output _,saved_for_loss = model(img) - + total_loss, saved_for_log = get_loss(saved_for_loss, heatmap_target, heat_mask, paf_target, paf_mask) - - for name,_ in meter_dict.items(): + + for name in meter_dict: meter_dict[name].update(saved_for_log[name], img.size(0)) losses.update(total_loss, img.size(0)) @@ -178,8 +171,11 @@ def train(train_loader, model, optimizer, epoch): batch_time.update(time.time() - end) end = time.time() if i % args.print_freq == 0: - print_string = 'Epoch: [{0}][{1}/{2}]\t'.format(epoch, i, len(train_loader)) - print_string +='Data time {data_time.val:.3f} ({data_time.avg:.3f})\t'.format( data_time=data_time) + print_string = 'Epoch: [{0}][{1}/{2}]\t'.format( + epoch, i, len(train_loader) + ) + 'Data time {data_time.val:.3f} ({data_time.avg:.3f})\t'.format( + data_time=data_time + ) print_string += 'Loss {loss.val:.4f} ({loss.avg:.4f})'.format(loss=losses) for name, value in meter_dict.items(): @@ -191,13 +187,12 @@ def validate(val_loader, model, epoch): batch_time = AverageMeter() data_time = AverageMeter() losses = AverageMeter() - - meter_dict = {} - meter_dict['paf'] = AverageMeter() - meter_dict['heatmap'] = AverageMeter() + + meter_dict = {'paf': AverageMeter()} + meter_dict['heatmap'] = AverageMeter() meter_dict['max_ht'] = AverageMeter() - meter_dict['min_ht'] = AverageMeter() - meter_dict['max_paf'] = AverageMeter() + meter_dict['min_ht'] = AverageMeter() + meter_dict['max_paf'] = AverageMeter() meter_dict['min_paf'] = AverageMeter() # switch to train mode model.train() @@ -212,16 +207,16 @@ def validate(val_loader, model, epoch): heat_mask = heat_mask.cuda() paf_target = paf_target.cuda() paf_mask = paf_mask.cuda() - + # compute output _,saved_for_loss = model(img) - + total_loss, saved_for_log = get_loss(saved_for_loss, heatmap_target, heat_mask, paf_target, paf_mask) - + #for name,_ in meter_dict.items(): # meter_dict[name].update(saved_for_log[name], img.size(0)) - + losses.update(total_loss.item(), img.size(0)) # compute gradient and do SGD step @@ -231,16 +226,19 @@ def validate(val_loader, model, epoch): # measure elapsed time batch_time.update(time.time() - end) - end = time.time() + end = time.time() if i % args.print_freq == 0: - print_string = 'Epoch: [{0}][{1}/{2}]\t'.format(epoch, i, len(val_loader)) - print_string +='Data time {data_time.val:.3f} ({data_time.avg:.3f})\t'.format( data_time=data_time) + print_string = 'Epoch: [{0}][{1}/{2}]\t'.format( + epoch, i, len(val_loader) + ) + 'Data time {data_time.val:.3f} ({data_time.avg:.3f})\t'.format( + data_time=data_time + ) print_string += 'Loss {loss.val:.4f} ({loss.avg:.4f})'.format(loss=losses) for name, value in meter_dict.items(): print_string+='{name}: {loss.val:.4f} ({loss.avg:.4f})\t'.format(name=name, loss=value) print(print_string) - + return losses.avg class AverageMeter(object): @@ -267,19 +265,19 @@ def update(self, val, n=1): args.mask_dir, 256, 4, 'rtpose', args.batch_size, shuffle=True, params_transform=params_transform, training=True, num_workers=16) -print('train dataset len: {}'.format(len(train_data.dataset))) +print(f'train dataset len: {len(train_data.dataset)}') # validation data valid_data = get_loader(args.json_path, args.data_dir, args.mask_dir, 256, 4, preprocess='rtpose', params_transform=params_transform, training=False, batch_size=args.batch_size, shuffle=False, num_workers=4) -print('val dataset len: {}'.format(len(valid_data.dataset))) +print(f'val dataset len: {len(valid_data.dataset)}') # model model = hourglass.hg(num_stacks=8, num_blocks=1, paf_classes=38, ht_classes=19) #model = encoding.nn.DataParallelModel(model, device_ids=args.gpu_ids) model = torch.nn.DataParallel(model).cuda() - + writer = SummaryWriter(log_dir=args.logdir) trainable_vars = [param for param in model.parameters() if param.requires_grad] @@ -287,7 +285,7 @@ def update(self, val, n=1): momentum=args.momentum, weight_decay=args.weight_decay, nesterov=args.nesterov) - + lr_scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.8, patience=5, verbose=True, threshold=0.0001, threshold_mode='rel', cooldown=3, min_lr=0, eps=1e-08) best_val_loss = np.inf @@ -301,15 +299,15 @@ def update(self, val, n=1): # evaluate on validation set val_loss = validate(valid_data, model, epoch) - + writer.add_scalars('data/scalar_group', {'train loss': train_loss, 'val loss': val_loss}, epoch) lr_scheduler.step(val_loss) - + is_best = val_loss