diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 00000000..d976407b --- /dev/null +++ b/Dockerfile @@ -0,0 +1,31 @@ +FROM pytorch/pytorch:2.5.1-cuda12.4-cudnn9-devel + +ENV TZ=Asia/Seoul +ENV COCOAPI=/workspace/cocoapi + +WORKDIR /workspace + +RUN apt-get update -y && \ + apt-get upgrade -y && \ + apt-get install git ffmpeg libsm6 libxext6 -y + +RUN mkdir /workspace/deep-high-resolution-net.pytorch +COPY . /workspace/deep-high-resolution-net.pytorch + +RUN pip install --no-cache-dir -r requirements.txt + +WORKDIR ./deep-high-resolution-net.pytorch/lib +RUN make + +RUN git clone https://github.com/cocodataset/cocoapi.git $COCOAPI +WORKDIR $COCOAPI/PythonAPI +RUN make install +RUN python setup.py install --user + +ENV NVIDIA_VISIBLE_DEVICES all +ENV NVIDIA_DRIVER_CAPABILITIES compute,utility + +WORKDIR /workspace/deep-high-resolution-net.pytorch +RUN mkdir output log models data data/coco data/mpii + +CMD ["bash", "start.sh"] \ No newline at end of file diff --git a/lib/dataset/coco.py b/lib/dataset/coco.py index b8bad408..52fd698e 100755 --- a/lib/dataset/coco.py +++ b/lib/dataset/coco.py @@ -27,7 +27,7 @@ class COCODataset(JointsDataset): - ''' + """ "keypoints": { 0: "nose", 1: "left_eye", @@ -47,10 +47,11 @@ class COCODataset(JointsDataset): 15: "left_ankle", 16: "right_ankle" }, - "skeleton": [ + "skeleton": [ [16,14],[14,12],[17,15],[15,13],[12,13],[6,12],[7,13], [6,7],[6,8], [7,9],[8,10],[9,11],[2,3],[1,2],[1,3],[2,4],[3,5],[4,6],[5,7]] - ''' + """ + def __init__(self, cfg, root, image_set, is_train, transform=None): super().__init__(cfg, root, image_set, is_train, transform) self.nms_thre = cfg.TEST.NMS_THRE @@ -68,38 +69,30 @@ def __init__(self, cfg, root, image_set, is_train, transform=None): self.coco = COCO(self._get_ann_file_keypoint()) # deal with class names - cats = [cat['name'] - for cat in self.coco.loadCats(self.coco.getCatIds())] - self.classes = ['__background__'] + cats - logger.info('=> classes: {}'.format(self.classes)) + cats = [cat["name"] for cat in self.coco.loadCats(self.coco.getCatIds())] + self.classes = ["__background__"] + cats + logger.info("=> classes: {}".format(self.classes)) self.num_classes = len(self.classes) self._class_to_ind = dict(zip(self.classes, range(self.num_classes))) self._class_to_coco_ind = dict(zip(cats, self.coco.getCatIds())) self._coco_ind_to_class_ind = dict( - [ - (self._class_to_coco_ind[cls], self._class_to_ind[cls]) - for cls in self.classes[1:] - ] + [(self._class_to_coco_ind[cls], self._class_to_ind[cls]) for cls in self.classes[1:]] ) # load image file names self.image_set_index = self._load_image_set_index() self.num_images = len(self.image_set_index) - logger.info('=> num_images: {}'.format(self.num_images)) + logger.info("=> num_images: {}".format(self.num_images)) self.num_joints = 17 - self.flip_pairs = [[1, 2], [3, 4], [5, 6], [7, 8], - [9, 10], [11, 12], [13, 14], [15, 16]] + self.flip_pairs = [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12], [13, 14], [15, 16]] self.parent_ids = None self.upper_body_ids = (0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10) self.lower_body_ids = (11, 12, 13, 14, 15, 16) self.joints_weight = np.array( - [ - 1., 1., 1., 1., 1., 1., 1., 1.2, 1.2, - 1.5, 1.5, 1., 1., 1.2, 1.2, 1.5, 1.5 - ], - dtype=np.float32 + [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.2, 1.2, 1.5, 1.5, 1.0, 1.0, 1.2, 1.2, 1.5, 1.5], + dtype=np.float32, ).reshape((self.num_joints, 1)) self.db = self._get_db() @@ -107,20 +100,15 @@ def __init__(self, cfg, root, image_set, is_train, transform=None): if is_train and cfg.DATASET.SELECT_DATA: self.db = self.select_data(self.db) - logger.info('=> load {} samples'.format(len(self.db))) + logger.info("=> load {} samples".format(len(self.db))) def _get_ann_file_keypoint(self): - """ self.root / annotations / person_keypoints_train2017.json """ - prefix = 'person_keypoints' \ - if 'test' not in self.image_set else 'image_info' - return os.path.join( - self.root, - 'annotations', - prefix + '_' + self.image_set + '.json' - ) + """self.root / annotations / person_keypoints_train2017.json""" + prefix = "person_keypoints" if "test" not in self.image_set else "image_info" + return os.path.join(self.root, "annotations", prefix + "_" + self.image_set + ".json") def _load_image_set_index(self): - """ image id: int """ + """image id: int""" image_ids = self.coco.getImgIds() return image_ids @@ -134,7 +122,7 @@ def _get_db(self): return gt_db def _load_coco_keypoint_annotations(self): - """ ground truth bbox and keypoints """ + """ground truth bbox and keypoints""" gt_db = [] for index in self.image_set_index: gt_db.extend(self._load_coco_keypoint_annotation_kernal(index)) @@ -152,8 +140,8 @@ def _load_coco_keypoint_annotation_kernal(self, index): :return: db entry """ im_ann = self.coco.loadImgs(index)[0] - width = im_ann['width'] - height = im_ann['height'] + width = im_ann["width"] + height = im_ann["height"] annIds = self.coco.getAnnIds(imgIds=index, iscrowd=False) objs = self.coco.loadAnns(annIds) @@ -161,49 +149,51 @@ def _load_coco_keypoint_annotation_kernal(self, index): # sanitize bboxes valid_objs = [] for obj in objs: - x, y, w, h = obj['bbox'] + x, y, w, h = obj["bbox"] x1 = np.max((0, x)) y1 = np.max((0, y)) x2 = np.min((width - 1, x1 + np.max((0, w - 1)))) y2 = np.min((height - 1, y1 + np.max((0, h - 1)))) - if obj['area'] > 0 and x2 >= x1 and y2 >= y1: - obj['clean_bbox'] = [x1, y1, x2-x1, y2-y1] + if obj["area"] > 0 and x2 >= x1 and y2 >= y1: + obj["clean_bbox"] = [x1, y1, x2 - x1, y2 - y1] valid_objs.append(obj) objs = valid_objs rec = [] for obj in objs: - cls = self._coco_ind_to_class_ind[obj['category_id']] + cls = self._coco_ind_to_class_ind[obj["category_id"]] if cls != 1: continue # ignore objs without keypoints annotation - if max(obj['keypoints']) == 0: + if max(obj["keypoints"]) == 0: continue - joints_3d = np.zeros((self.num_joints, 3), dtype=np.float) - joints_3d_vis = np.zeros((self.num_joints, 3), dtype=np.float) + joints_3d = np.zeros((self.num_joints, 3), dtype=np.float32) + joints_3d_vis = np.zeros((self.num_joints, 3), dtype=np.float32) for ipt in range(self.num_joints): - joints_3d[ipt, 0] = obj['keypoints'][ipt * 3 + 0] - joints_3d[ipt, 1] = obj['keypoints'][ipt * 3 + 1] + joints_3d[ipt, 0] = obj["keypoints"][ipt * 3 + 0] + joints_3d[ipt, 1] = obj["keypoints"][ipt * 3 + 1] joints_3d[ipt, 2] = 0 - t_vis = obj['keypoints'][ipt * 3 + 2] + t_vis = obj["keypoints"][ipt * 3 + 2] if t_vis > 1: t_vis = 1 joints_3d_vis[ipt, 0] = t_vis joints_3d_vis[ipt, 1] = t_vis joints_3d_vis[ipt, 2] = 0 - center, scale = self._box2cs(obj['clean_bbox'][:4]) - rec.append({ - 'image': self.image_path_from_index(index), - 'center': center, - 'scale': scale, - 'joints_3d': joints_3d, - 'joints_3d_vis': joints_3d_vis, - 'filename': '', - 'imgnum': 0, - }) + center, scale = self._box2cs(obj["clean_bbox"][:4]) + rec.append( + { + "image": self.image_path_from_index(index), + "center": center, + "scale": scale, + "joints_3d": joints_3d, + "joints_3d_vis": joints_3d_vis, + "filename": "", + "imgnum": 0, + } + ) return rec @@ -220,49 +210,46 @@ def _xywh2cs(self, x, y, w, h): h = w * 1.0 / self.aspect_ratio elif w < self.aspect_ratio * h: w = h * self.aspect_ratio - scale = np.array( - [w * 1.0 / self.pixel_std, h * 1.0 / self.pixel_std], - dtype=np.float32) + scale = np.array([w * 1.0 / self.pixel_std, h * 1.0 / self.pixel_std], dtype=np.float32) if center[0] != -1: scale = scale * 1.25 return center, scale def image_path_from_index(self, index): - """ example: images / train2017 / 000000119993.jpg """ - file_name = '%012d.jpg' % index - if '2014' in self.image_set: - file_name = 'COCO_%s_' % self.image_set + file_name + """example: images / train2017 / 000000119993.jpg""" + file_name = "%012d.jpg" % index + if "2014" in self.image_set: + file_name = "COCO_%s_" % self.image_set + file_name - prefix = 'test2017' if 'test' in self.image_set else self.image_set + prefix = "test2017" if "test" in self.image_set else self.image_set - data_name = prefix + '.zip@' if self.data_format == 'zip' else prefix + data_name = prefix + ".zip@" if self.data_format == "zip" else prefix - image_path = os.path.join( - self.root, 'images', data_name, file_name) + image_path = os.path.join(self.root, "images", data_name, file_name) return image_path def _load_coco_person_detection_results(self): all_boxes = None - with open(self.bbox_file, 'r') as f: + with open(self.bbox_file, "r") as f: all_boxes = json.load(f) if not all_boxes: - logger.error('=> Load %s fail!' % self.bbox_file) + logger.error("=> Load %s fail!" % self.bbox_file) return None - logger.info('=> Total boxes: {}'.format(len(all_boxes))) + logger.info("=> Total boxes: {}".format(len(all_boxes))) kpt_db = [] num_boxes = 0 for n_img in range(0, len(all_boxes)): det_res = all_boxes[n_img] - if det_res['category_id'] != 1: + if det_res["category_id"] != 1: continue - img_name = self.image_path_from_index(det_res['image_id']) - box = det_res['bbox'] - score = det_res['score'] + img_name = self.image_path_from_index(det_res["image_id"]) + box = det_res["bbox"] + score = det_res["score"] if score < self.image_thre: continue @@ -270,53 +257,51 @@ def _load_coco_person_detection_results(self): num_boxes = num_boxes + 1 center, scale = self._box2cs(box) - joints_3d = np.zeros((self.num_joints, 3), dtype=np.float) - joints_3d_vis = np.ones( - (self.num_joints, 3), dtype=np.float) - kpt_db.append({ - 'image': img_name, - 'center': center, - 'scale': scale, - 'score': score, - 'joints_3d': joints_3d, - 'joints_3d_vis': joints_3d_vis, - }) - - logger.info('=> Total boxes after fliter low score@{}: {}'.format( - self.image_thre, num_boxes)) + joints_3d = np.zeros((self.num_joints, 3), dtype=np.float32) + joints_3d_vis = np.ones((self.num_joints, 3), dtype=np.float32) + kpt_db.append( + { + "image": img_name, + "center": center, + "scale": scale, + "score": score, + "joints_3d": joints_3d, + "joints_3d_vis": joints_3d_vis, + } + ) + + logger.info("=> Total boxes after fliter low score@{}: {}".format(self.image_thre, num_boxes)) return kpt_db - def evaluate(self, cfg, preds, output_dir, all_boxes, img_path, - *args, **kwargs): + def evaluate(self, cfg, preds, output_dir, all_boxes, img_path, *args, **kwargs): rank = cfg.RANK - res_folder = os.path.join(output_dir, 'results') + res_folder = os.path.join(output_dir, "results") if not os.path.exists(res_folder): try: os.makedirs(res_folder) except Exception: - logger.error('Fail to make {}'.format(res_folder)) + logger.error("Fail to make {}".format(res_folder)) - res_file = os.path.join( - res_folder, 'keypoints_{}_results_{}.json'.format( - self.image_set, rank) - ) + res_file = os.path.join(res_folder, "keypoints_{}_results_{}.json".format(self.image_set, rank)) # person x (keypoints) _kpts = [] for idx, kpt in enumerate(preds): - _kpts.append({ - 'keypoints': kpt, - 'center': all_boxes[idx][0:2], - 'scale': all_boxes[idx][2:4], - 'area': all_boxes[idx][4], - 'score': all_boxes[idx][5], - 'image': int(img_path[idx][-16:-4]) - }) + _kpts.append( + { + "keypoints": kpt, + "center": all_boxes[idx][0:2], + "scale": all_boxes[idx][2:4], + "area": all_boxes[idx][4], + "score": all_boxes[idx][5], + "image": int(img_path[idx][-16:-4]), + } + ) # image x person x (keypoints) kpts = defaultdict(list) for kpt in _kpts: - kpts[kpt['image']].append(kpt) + kpts[kpt["image"]].append(kpt) # rescoring and oks nms num_joints = self.num_joints @@ -326,87 +311,77 @@ def evaluate(self, cfg, preds, output_dir, all_boxes, img_path, for img in kpts.keys(): img_kpts = kpts[img] for n_p in img_kpts: - box_score = n_p['score'] + box_score = n_p["score"] kpt_score = 0 valid_num = 0 for n_jt in range(0, num_joints): - t_s = n_p['keypoints'][n_jt][2] + t_s = n_p["keypoints"][n_jt][2] if t_s > in_vis_thre: kpt_score = kpt_score + t_s valid_num = valid_num + 1 if valid_num != 0: kpt_score = kpt_score / valid_num # rescoring - n_p['score'] = kpt_score * box_score + n_p["score"] = kpt_score * box_score if self.soft_nms: - keep = soft_oks_nms( - [img_kpts[i] for i in range(len(img_kpts))], - oks_thre - ) + keep = soft_oks_nms([img_kpts[i] for i in range(len(img_kpts))], oks_thre) else: - keep = oks_nms( - [img_kpts[i] for i in range(len(img_kpts))], - oks_thre - ) + keep = oks_nms([img_kpts[i] for i in range(len(img_kpts))], oks_thre) if len(keep) == 0: oks_nmsed_kpts.append(img_kpts) else: oks_nmsed_kpts.append([img_kpts[_keep] for _keep in keep]) - self._write_coco_keypoint_results( - oks_nmsed_kpts, res_file) - if 'test' not in self.image_set: - info_str = self._do_python_keypoint_eval( - res_file, res_folder) + self._write_coco_keypoint_results(oks_nmsed_kpts, res_file) + if "test" not in self.image_set: + info_str = self._do_python_keypoint_eval(res_file, res_folder) name_value = OrderedDict(info_str) - return name_value, name_value['AP'] + return name_value, name_value["AP"] else: - return {'Null': 0}, 0 + return {"Null": 0}, 0 def _write_coco_keypoint_results(self, keypoints, res_file): data_pack = [ { - 'cat_id': self._class_to_coco_ind[cls], - 'cls_ind': cls_ind, - 'cls': cls, - 'ann_type': 'keypoints', - 'keypoints': keypoints + "cat_id": self._class_to_coco_ind[cls], + "cls_ind": cls_ind, + "cls": cls, + "ann_type": "keypoints", + "keypoints": keypoints, } - for cls_ind, cls in enumerate(self.classes) if not cls == '__background__' + for cls_ind, cls in enumerate(self.classes) + if not cls == "__background__" ] results = self._coco_keypoint_results_one_category_kernel(data_pack[0]) - logger.info('=> writing results json to %s' % res_file) - with open(res_file, 'w') as f: + logger.info("=> writing results json to %s" % res_file) + with open(res_file, "w") as f: json.dump(results, f, sort_keys=True, indent=4) try: json.load(open(res_file)) except Exception: content = [] - with open(res_file, 'r') as f: + with open(res_file, "r") as f: for line in f: content.append(line) - content[-1] = ']' - with open(res_file, 'w') as f: + content[-1] = "]" + with open(res_file, "w") as f: for c in content: f.write(c) def _coco_keypoint_results_one_category_kernel(self, data_pack): - cat_id = data_pack['cat_id'] - keypoints = data_pack['keypoints'] + cat_id = data_pack["cat_id"] + keypoints = data_pack["keypoints"] cat_results = [] for img_kpts in keypoints: if len(img_kpts) == 0: continue - _key_points = np.array([img_kpts[k]['keypoints'] - for k in range(len(img_kpts))]) - key_points = np.zeros( - (_key_points.shape[0], self.num_joints * 3), dtype=np.float - ) + _key_points = np.array([img_kpts[k]["keypoints"] for k in range(len(img_kpts))]) + key_points = np.zeros((_key_points.shape[0], self.num_joints * 3), dtype=np.float32) for ipt in range(self.num_joints): key_points[:, ipt * 3 + 0] = _key_points[:, ipt, 0] @@ -415,12 +390,12 @@ def _coco_keypoint_results_one_category_kernel(self, data_pack): result = [ { - 'image_id': img_kpts[k]['image'], - 'category_id': cat_id, - 'keypoints': list(key_points[k]), - 'score': img_kpts[k]['score'], - 'center': list(img_kpts[k]['center']), - 'scale': list(img_kpts[k]['scale']) + "image_id": img_kpts[k]["image"], + "category_id": cat_id, + "keypoints": list(key_points[k]), + "score": img_kpts[k]["score"], + "center": list(img_kpts[k]["center"]), + "scale": list(img_kpts[k]["scale"]), } for k in range(len(img_kpts)) ] @@ -430,13 +405,24 @@ def _coco_keypoint_results_one_category_kernel(self, data_pack): def _do_python_keypoint_eval(self, res_file, res_folder): coco_dt = self.coco.loadRes(res_file) - coco_eval = COCOeval(self.coco, coco_dt, 'keypoints') + coco_eval = COCOeval(self.coco, coco_dt, "keypoints") coco_eval.params.useSegm = None coco_eval.evaluate() coco_eval.accumulate() coco_eval.summarize() - stats_names = ['AP', 'Ap .5', 'AP .75', 'AP (M)', 'AP (L)', 'AR', 'AR .5', 'AR .75', 'AR (M)', 'AR (L)'] + stats_names = [ + "AP", + "Ap .5", + "AP .75", + "AP (M)", + "AP (L)", + "AR", + "AR .5", + "AR .75", + "AR (M)", + "AR (L)", + ] info_str = [] for ind, name in enumerate(stats_names): diff --git a/lib/dataset/mpii.py b/lib/dataset/mpii.py index c935001a..1d916dec 100644 --- a/lib/dataset/mpii.py +++ b/lib/dataset/mpii.py @@ -38,22 +38,20 @@ def __init__(self, cfg, root, image_set, is_train, transform=None): if is_train and cfg.DATASET.SELECT_DATA: self.db = self.select_data(self.db) - logger.info('=> load {} samples'.format(len(self.db))) + logger.info("=> load {} samples".format(len(self.db))) def _get_db(self): # create train/val split - file_name = os.path.join( - self.root, 'annot', self.image_set+'.json' - ) + file_name = os.path.join(self.root, "annot", self.image_set + ".json") with open(file_name) as anno_file: anno = json.load(anno_file) gt_db = [] for a in anno: - image_name = a['image'] + image_name = a["image"] - c = np.array(a['center'], dtype=np.float) - s = np.array([a['scale'], a['scale']], dtype=np.float) + c = np.array(a["center"], dtype=float) + s = np.array([a["scale"], a["scale"]], dtype=float) # Adjust center/scale slightly to avoid cropping limbs if c[0] != -1: @@ -64,30 +62,30 @@ def _get_db(self): # we should first convert to 0-based index c = c - 1 - joints_3d = np.zeros((self.num_joints, 3), dtype=np.float) - joints_3d_vis = np.zeros((self.num_joints, 3), dtype=np.float) - if self.image_set != 'test': - joints = np.array(a['joints']) + joints_3d = np.zeros((self.num_joints, 3), dtype=float) + joints_3d_vis = np.zeros((self.num_joints, 3), dtype=float) + if self.image_set != "test": + joints = np.array(a["joints"]) joints[:, 0:2] = joints[:, 0:2] - 1 - joints_vis = np.array(a['joints_vis']) - assert len(joints) == self.num_joints, \ - 'joint num diff: {} vs {}'.format(len(joints), - self.num_joints) + joints_vis = np.array(a["joints_vis"]) + assert len(joints) == self.num_joints, "joint num diff: {} vs {}".format( + len(joints), self.num_joints + ) joints_3d[:, 0:2] = joints[:, 0:2] joints_3d_vis[:, 0] = joints_vis[:] joints_3d_vis[:, 1] = joints_vis[:] - image_dir = 'images.zip@' if self.data_format == 'zip' else 'images' + image_dir = "images.zip@" if self.data_format == "zip" else "images" gt_db.append( { - 'image': os.path.join(self.root, image_dir, image_name), - 'center': c, - 'scale': s, - 'joints_3d': joints_3d, - 'joints_3d_vis': joints_3d_vis, - 'filename': '', - 'imgnum': 0, + "image": os.path.join(self.root, image_dir, image_name), + "center": c, + "scale": s, + "joints_3d": joints_3d, + "joints_3d_vis": joints_3d_vis, + "filename": "", + "imgnum": 0, } ) @@ -98,40 +96,38 @@ def evaluate(self, cfg, preds, output_dir, *args, **kwargs): preds = preds[:, :, 0:2] + 1.0 if output_dir: - pred_file = os.path.join(output_dir, 'pred.mat') - savemat(pred_file, mdict={'preds': preds}) + pred_file = os.path.join(output_dir, "pred.mat") + savemat(pred_file, mdict={"preds": preds}) - if 'test' in cfg.DATASET.TEST_SET: - return {'Null': 0.0}, 0.0 + if "test" in cfg.DATASET.TEST_SET: + return {"Null": 0.0}, 0.0 SC_BIAS = 0.6 threshold = 0.5 - gt_file = os.path.join(cfg.DATASET.ROOT, - 'annot', - 'gt_{}.mat'.format(cfg.DATASET.TEST_SET)) + gt_file = os.path.join(cfg.DATASET.ROOT, "annot", "gt_{}.mat".format(cfg.DATASET.TEST_SET)) gt_dict = loadmat(gt_file) - dataset_joints = gt_dict['dataset_joints'] - jnt_missing = gt_dict['jnt_missing'] - pos_gt_src = gt_dict['pos_gt_src'] - headboxes_src = gt_dict['headboxes_src'] + dataset_joints = gt_dict["dataset_joints"] + jnt_missing = gt_dict["jnt_missing"] + pos_gt_src = gt_dict["pos_gt_src"] + headboxes_src = gt_dict["headboxes_src"] pos_pred_src = np.transpose(preds, [1, 2, 0]) - head = np.where(dataset_joints == 'head')[1][0] - lsho = np.where(dataset_joints == 'lsho')[1][0] - lelb = np.where(dataset_joints == 'lelb')[1][0] - lwri = np.where(dataset_joints == 'lwri')[1][0] - lhip = np.where(dataset_joints == 'lhip')[1][0] - lkne = np.where(dataset_joints == 'lkne')[1][0] - lank = np.where(dataset_joints == 'lank')[1][0] - - rsho = np.where(dataset_joints == 'rsho')[1][0] - relb = np.where(dataset_joints == 'relb')[1][0] - rwri = np.where(dataset_joints == 'rwri')[1][0] - rkne = np.where(dataset_joints == 'rkne')[1][0] - rank = np.where(dataset_joints == 'rank')[1][0] - rhip = np.where(dataset_joints == 'rhip')[1][0] + head = np.where(dataset_joints == "head")[1][0] + lsho = np.where(dataset_joints == "lsho")[1][0] + lelb = np.where(dataset_joints == "lelb")[1][0] + lwri = np.where(dataset_joints == "lwri")[1][0] + lhip = np.where(dataset_joints == "lhip")[1][0] + lkne = np.where(dataset_joints == "lkne")[1][0] + lank = np.where(dataset_joints == "lank")[1][0] + + rsho = np.where(dataset_joints == "rsho")[1][0] + relb = np.where(dataset_joints == "relb")[1][0] + rwri = np.where(dataset_joints == "rwri")[1][0] + rkne = np.where(dataset_joints == "rkne")[1][0] + rank = np.where(dataset_joints == "rank")[1][0] + rhip = np.where(dataset_joints == "rhip")[1][0] jnt_visible = 1 - jnt_missing uv_error = pos_pred_src - pos_gt_src @@ -143,39 +139,36 @@ def evaluate(self, cfg, preds, output_dir, *args, **kwargs): scaled_uv_err = np.divide(uv_err, scale) scaled_uv_err = np.multiply(scaled_uv_err, jnt_visible) jnt_count = np.sum(jnt_visible, axis=1) - less_than_threshold = np.multiply((scaled_uv_err <= threshold), - jnt_visible) - PCKh = np.divide(100.*np.sum(less_than_threshold, axis=1), jnt_count) + less_than_threshold = np.multiply((scaled_uv_err <= threshold), jnt_visible) + PCKh = np.divide(100.0 * np.sum(less_than_threshold, axis=1), jnt_count) # save - rng = np.arange(0, 0.5+0.01, 0.01) + rng = np.arange(0, 0.5 + 0.01, 0.01) pckAll = np.zeros((len(rng), 16)) for r in range(len(rng)): threshold = rng[r] - less_than_threshold = np.multiply(scaled_uv_err <= threshold, - jnt_visible) - pckAll[r, :] = np.divide(100.*np.sum(less_than_threshold, axis=1), - jnt_count) + less_than_threshold = np.multiply(scaled_uv_err <= threshold, jnt_visible) + pckAll[r, :] = np.divide(100.0 * np.sum(less_than_threshold, axis=1), jnt_count) PCKh = np.ma.array(PCKh, mask=False) PCKh.mask[6:8] = True jnt_count = np.ma.array(jnt_count, mask=False) jnt_count.mask[6:8] = True - jnt_ratio = jnt_count / np.sum(jnt_count).astype(np.float64) + jnt_ratio = jnt_count / np.sum(jnt_count).astype(float64) name_value = [ - ('Head', PCKh[head]), - ('Shoulder', 0.5 * (PCKh[lsho] + PCKh[rsho])), - ('Elbow', 0.5 * (PCKh[lelb] + PCKh[relb])), - ('Wrist', 0.5 * (PCKh[lwri] + PCKh[rwri])), - ('Hip', 0.5 * (PCKh[lhip] + PCKh[rhip])), - ('Knee', 0.5 * (PCKh[lkne] + PCKh[rkne])), - ('Ankle', 0.5 * (PCKh[lank] + PCKh[rank])), - ('Mean', np.sum(PCKh * jnt_ratio)), - ('Mean@0.1', np.sum(pckAll[11, :] * jnt_ratio)) + ("Head", PCKh[head]), + ("Shoulder", 0.5 * (PCKh[lsho] + PCKh[rsho])), + ("Elbow", 0.5 * (PCKh[lelb] + PCKh[relb])), + ("Wrist", 0.5 * (PCKh[lwri] + PCKh[rwri])), + ("Hip", 0.5 * (PCKh[lhip] + PCKh[rhip])), + ("Knee", 0.5 * (PCKh[lkne] + PCKh[rkne])), + ("Ankle", 0.5 * (PCKh[lank] + PCKh[rank])), + ("Mean", np.sum(PCKh * jnt_ratio)), + ("Mean@0.1", np.sum(pckAll[11, :] * jnt_ratio)), ] name_value = OrderedDict(name_value) - return name_value, name_value['Mean'] + return name_value, name_value["Mean"] diff --git a/lib/nms/cpu_nms.pyx b/lib/nms/cpu_nms.pyx index 3cd0d740..fc8e2676 100644 --- a/lib/nms/cpu_nms.pyx +++ b/lib/nms/cpu_nms.pyx @@ -25,10 +25,10 @@ def cpu_nms(np.ndarray[np.float32_t, ndim=2] dets, np.float thresh): cdef np.ndarray[np.float32_t, ndim=1] scores = dets[:, 4] cdef np.ndarray[np.float32_t, ndim=1] areas = (x2 - x1 + 1) * (y2 - y1 + 1) - cdef np.ndarray[np.int_t, ndim=1] order = scores.argsort()[::-1].astype('i') + cdef np.ndarray[np.int32_t, ndim=1] order = scores.argsort()[::-1].astype('i') cdef int ndets = dets.shape[0] - cdef np.ndarray[np.int_t, ndim=1] suppressed = \ + cdef np.ndarray[np.int32_t, ndim=1] suppressed = \ np.zeros((ndets), dtype=np.int) # nominal indices diff --git a/lib/nms/setup_linux.py b/lib/nms/setup_linux.py index 9120a93b..0b27e6d8 100644 --- a/lib/nms/setup_linux.py +++ b/lib/nms/setup_linux.py @@ -33,26 +33,33 @@ def locate_cuda(): """ # first check if the CUDAHOME env variable is in use - if 'CUDAHOME' in os.environ: - home = os.environ['CUDAHOME'] - nvcc = pjoin(home, 'bin', 'nvcc') + if "CUDAHOME" in os.environ: + home = os.environ["CUDAHOME"] + nvcc = pjoin(home, "bin", "nvcc") else: # otherwise, search the PATH for NVCC - default_path = pjoin(os.sep, 'usr', 'local', 'cuda', 'bin') - nvcc = find_in_path('nvcc', os.environ['PATH'] + os.pathsep + default_path) + default_path = pjoin(os.sep, "usr", "local", "cuda", "bin") + nvcc = find_in_path("nvcc", os.environ["PATH"] + os.pathsep + default_path) if nvcc is None: - raise EnvironmentError('The nvcc binary could not be ' - 'located in your $PATH. Either add it to your path, or set $CUDAHOME') + raise EnvironmentError( + "The nvcc binary could not be " + "located in your $PATH. Either add it to your path, or set $CUDAHOME" + ) home = os.path.dirname(os.path.dirname(nvcc)) - cudaconfig = {'home':home, 'nvcc':nvcc, - 'include': pjoin(home, 'include'), - 'lib64': pjoin(home, 'lib64')} + cudaconfig = { + "home": home, + "nvcc": nvcc, + "include": pjoin(home, "include"), + "lib64": pjoin(home, "lib64"), + } for k, v in cudaconfig.items(): if not os.path.exists(v): - raise EnvironmentError('The CUDA %s path could not be located in %s' % (k, v)) + raise EnvironmentError("The CUDA %s path could not be located in %s" % (k, v)) return cudaconfig + + CUDA = locate_cuda() @@ -73,7 +80,7 @@ def customize_compiler_for_nvcc(self): subclassing going on.""" # tell the compiler it can processes .cu - self.src_extensions.append('.cu') + self.src_extensions.append(".cu") # save references to the default compiler_so and _comple methods default_compiler_so = self.compiler_so @@ -83,14 +90,14 @@ def customize_compiler_for_nvcc(self): # object but distutils doesn't have the ability to change compilers # based on source extension: we add it. def _compile(obj, src, ext, cc_args, extra_postargs, pp_opts): - if os.path.splitext(src)[1] == '.cu': + if os.path.splitext(src)[1] == ".cu": # use the cuda for .cu files - self.set_executable('compiler_so', CUDA['nvcc']) + self.set_executable("compiler_so", CUDA["nvcc"]) # use only a subset of the extra_postargs, which are 1-1 translated # from the extra_compile_args in the Extension class - postargs = extra_postargs['nvcc'] + postargs = extra_postargs["nvcc"] else: - postargs = extra_postargs['gcc'] + postargs = extra_postargs["gcc"] super(obj, src, ext, cc_args, postargs, pp_opts) # reset the default compiler_so, which we might have changed for cuda @@ -111,31 +118,30 @@ def build_extensions(self): Extension( "cpu_nms", ["cpu_nms.pyx"], - extra_compile_args={'gcc': ["-Wno-cpp", "-Wno-unused-function"]}, - include_dirs = [numpy_include] + extra_compile_args={"gcc": ["-Wno-cpp", "-Wno-unused-function"]}, + include_dirs=[numpy_include], ), - Extension('gpu_nms', - ['nms_kernel.cu', 'gpu_nms.pyx'], - library_dirs=[CUDA['lib64']], - libraries=['cudart'], - language='c++', - runtime_library_dirs=[CUDA['lib64']], + Extension( + "gpu_nms", + ["nms_kernel.cu", "gpu_nms.pyx"], + library_dirs=[CUDA["lib64"]], + libraries=["cudart"], + language="c++", + runtime_library_dirs=[CUDA["lib64"]], # this syntax is specific to this build system # we're only going to use certain compiler args with nvcc and not with # gcc the implementation of this trick is in customize_compiler() below - extra_compile_args={'gcc': ["-Wno-unused-function"], - 'nvcc': ['-arch=sm_35', - '--ptxas-options=-v', - '-c', - '--compiler-options', - "'-fPIC'"]}, - include_dirs = [numpy_include, CUDA['include']] + extra_compile_args={ + "gcc": ["-Wno-unused-function"], + "nvcc": ["-arch=sm_86", "--ptxas-options=-v", "-c", "--compiler-options", "'-fPIC'"], + }, + include_dirs=[numpy_include, CUDA["include"]], ), ] setup( - name='nms', + name="nms", ext_modules=ext_modules, # inject our custom trigger - cmdclass={'build_ext': custom_build_ext}, + cmdclass={"build_ext": custom_build_ext}, ) diff --git a/requirements.txt b/requirements.txt index 14f225c7..a89823ba 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,11 +1,10 @@ EasyDict==1.7 -opencv-python==3.4.1.15 -shapely==1.6.4 -Cython -scipy -pandas +opencv-python==4.11.0.86 +shapely==2.0.6 +Cython==3.0.11 +scipy==1.15.1 +pandas==2.2.3 pyyaml -json_tricks -scikit-image -yacs>=0.1.5 -tensorboardX==1.6 +json_tricks==3.17.3 +scikit-image==0.25.0 +yacs>=0.1.5 \ No newline at end of file