完善部署并训练5个新发型 + 换发型集成文档
部署修复： - torch.load 增加 weights_only=False patch，兼容 PyTorch 2.6+ 加载旧权重 - OSS 改为懒加载，本地用 output_format=base64 无需配凭证即可启动 - 补全被 gitignore 误排除的必需代码：core/models/layers/data、models/layers/data、keypoints/lib - webui 训练命令 --xformers 改 --sdpa（修复 xformers 无 CUDA 支持报错） 功能调整： - hair_grow_service 端口改 8899、preview 路由修复（send_file） - list_hairstyles 增加发型白名单，测试页只展示当前5个发型 新增脚本： - train_lora_parallel.py：直接调 kohya 并行训练 LoRA（绕过 photo_service 串行限制） - train_hairstyles_parallel.py / train_batch_stepC.py：批量训练辅助脚本 - scripts/sync_data_to_server.sh：大文件断点续传到云服务器 文档： - docs/换发型集成文档.md：换发型完整流程、服务架构、资源依赖、训练方法、集成步骤
This commit is contained in:
@@ -0,0 +1,292 @@
|
||||
# ------------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft
|
||||
# Licensed under the MIT License.
|
||||
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
|
||||
# ------------------------------------------------------------------------------
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import copy
|
||||
import logging
|
||||
import random
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.utils.data import Dataset
|
||||
|
||||
from utils.transforms import get_affine_transform
|
||||
from utils.transforms import affine_transform
|
||||
from utils.transforms import fliplr_joints
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class JointsDataset(Dataset):
|
||||
def __init__(self, cfg, root, image_set, is_train, transform=None):
|
||||
self.num_joints = 0
|
||||
self.pixel_std = 200
|
||||
self.flip_pairs = []
|
||||
self.parent_ids = []
|
||||
|
||||
self.is_train = is_train
|
||||
self.root = root
|
||||
self.image_set = image_set
|
||||
|
||||
self.output_path = cfg.OUTPUT_DIR
|
||||
self.data_format = cfg.DATASET.DATA_FORMAT
|
||||
|
||||
self.scale_factor = cfg.DATASET.SCALE_FACTOR
|
||||
self.rotation_factor = cfg.DATASET.ROT_FACTOR
|
||||
self.flip = cfg.DATASET.FLIP
|
||||
self.num_joints_half_body = cfg.DATASET.NUM_JOINTS_HALF_BODY
|
||||
self.prob_half_body = cfg.DATASET.PROB_HALF_BODY
|
||||
self.color_rgb = cfg.DATASET.COLOR_RGB
|
||||
|
||||
self.target_type = cfg.MODEL.TARGET_TYPE
|
||||
self.image_size = np.array(cfg.MODEL.IMAGE_SIZE)
|
||||
self.heatmap_size = np.array(cfg.MODEL.HEATMAP_SIZE)
|
||||
self.sigma = cfg.MODEL.SIGMA
|
||||
self.use_different_joints_weight = cfg.LOSS.USE_DIFFERENT_JOINTS_WEIGHT
|
||||
self.joints_weight = 1
|
||||
|
||||
self.transform = transform
|
||||
self.db = []
|
||||
|
||||
def _get_db(self):
|
||||
raise NotImplementedError
|
||||
|
||||
def evaluate(self, cfg, preds, output_dir, *args, **kwargs):
|
||||
raise NotImplementedError
|
||||
|
||||
def half_body_transform(self, joints, joints_vis):
|
||||
upper_joints = []
|
||||
lower_joints = []
|
||||
for joint_id in range(self.num_joints):
|
||||
if joints_vis[joint_id][0] > 0:
|
||||
if joint_id in self.upper_body_ids:
|
||||
upper_joints.append(joints[joint_id])
|
||||
else:
|
||||
lower_joints.append(joints[joint_id])
|
||||
|
||||
if np.random.randn() < 0.5 and len(upper_joints) > 2:
|
||||
selected_joints = upper_joints
|
||||
else:
|
||||
selected_joints = lower_joints \
|
||||
if len(lower_joints) > 2 else upper_joints
|
||||
|
||||
if len(selected_joints) < 2:
|
||||
return None, None
|
||||
|
||||
selected_joints = np.array(selected_joints, dtype=np.float32)
|
||||
center = selected_joints.mean(axis=0)[:2]
|
||||
|
||||
left_top = np.amin(selected_joints, axis=0)
|
||||
right_bottom = np.amax(selected_joints, axis=0)
|
||||
|
||||
w = right_bottom[0] - left_top[0]
|
||||
h = right_bottom[1] - left_top[1]
|
||||
|
||||
if w > self.aspect_ratio * 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 = scale * 1.5
|
||||
|
||||
return center, scale
|
||||
|
||||
def __len__(self,):
|
||||
return len(self.db)
|
||||
|
||||
def __getitem__(self, idx):
|
||||
db_rec = copy.deepcopy(self.db[idx])
|
||||
|
||||
image_file = db_rec['image']
|
||||
filename = db_rec['filename'] if 'filename' in db_rec else ''
|
||||
imgnum = db_rec['imgnum'] if 'imgnum' in db_rec else ''
|
||||
|
||||
if self.data_format == 'zip':
|
||||
from utils import zipreader
|
||||
data_numpy = zipreader.imread(
|
||||
image_file, cv2.IMREAD_COLOR | cv2.IMREAD_IGNORE_ORIENTATION
|
||||
)
|
||||
else:
|
||||
data_numpy = cv2.imread(
|
||||
image_file, cv2.IMREAD_COLOR | cv2.IMREAD_IGNORE_ORIENTATION
|
||||
)
|
||||
|
||||
if self.color_rgb:
|
||||
data_numpy = cv2.cvtColor(data_numpy, cv2.COLOR_BGR2RGB)
|
||||
|
||||
if data_numpy is None:
|
||||
logger.error('=> fail to read {}'.format(image_file))
|
||||
raise ValueError('Fail to read {}'.format(image_file))
|
||||
|
||||
joints = db_rec['joints_3d']
|
||||
joints_vis = db_rec['joints_3d_vis']
|
||||
|
||||
c = db_rec['center']
|
||||
s = db_rec['scale']
|
||||
score = db_rec['score'] if 'score' in db_rec else 1
|
||||
r = 0
|
||||
|
||||
if self.is_train:
|
||||
if (np.sum(joints_vis[:, 0]) > self.num_joints_half_body
|
||||
and np.random.rand() < self.prob_half_body):
|
||||
c_half_body, s_half_body = self.half_body_transform(
|
||||
joints, joints_vis
|
||||
)
|
||||
|
||||
if c_half_body is not None and s_half_body is not None:
|
||||
c, s = c_half_body, s_half_body
|
||||
|
||||
sf = self.scale_factor
|
||||
rf = self.rotation_factor
|
||||
s = s * np.clip(np.random.randn()*sf + 1, 1 - sf, 1 + sf)
|
||||
r = np.clip(np.random.randn()*rf, -rf*2, rf*2) \
|
||||
if random.random() <= 0.6 else 0
|
||||
|
||||
if self.flip and random.random() <= 0.5:
|
||||
data_numpy = data_numpy[:, ::-1, :]
|
||||
joints, joints_vis = fliplr_joints(
|
||||
joints, joints_vis, data_numpy.shape[1], self.flip_pairs)
|
||||
c[0] = data_numpy.shape[1] - c[0] - 1
|
||||
|
||||
trans = get_affine_transform(c, s, r, self.image_size)
|
||||
input = cv2.warpAffine(
|
||||
data_numpy,
|
||||
trans,
|
||||
(int(self.image_size[0]), int(self.image_size[1])),
|
||||
flags=cv2.INTER_LINEAR)
|
||||
|
||||
cv2.imshow('input', input)
|
||||
cv2.waitKey()
|
||||
|
||||
if self.transform:
|
||||
input = self.transform(input)
|
||||
|
||||
for i in range(self.num_joints):
|
||||
if joints_vis[i, 0] > 0.0:
|
||||
joints[i, 0:2] = affine_transform(joints[i, 0:2], trans)
|
||||
|
||||
target, target_weight = self.generate_target(joints, joints_vis)
|
||||
|
||||
target = torch.from_numpy(target)
|
||||
target_weight = torch.from_numpy(target_weight)
|
||||
|
||||
meta = {
|
||||
'image': image_file,
|
||||
'filename': filename,
|
||||
'imgnum': imgnum,
|
||||
'joints': joints,
|
||||
'joints_vis': joints_vis,
|
||||
'center': c,
|
||||
'scale': s,
|
||||
'rotation': r,
|
||||
'score': score
|
||||
}
|
||||
|
||||
return input, target, target_weight, meta
|
||||
|
||||
def select_data(self, db):
|
||||
db_selected = []
|
||||
for rec in db:
|
||||
num_vis = 0
|
||||
joints_x = 0.0
|
||||
joints_y = 0.0
|
||||
for joint, joint_vis in zip(
|
||||
rec['joints_3d'], rec['joints_3d_vis']):
|
||||
if joint_vis[0] <= 0:
|
||||
continue
|
||||
num_vis += 1
|
||||
|
||||
joints_x += joint[0]
|
||||
joints_y += joint[1]
|
||||
if num_vis == 0:
|
||||
continue
|
||||
|
||||
joints_x, joints_y = joints_x / num_vis, joints_y / num_vis
|
||||
|
||||
area = rec['scale'][0] * rec['scale'][1] * (self.pixel_std**2)
|
||||
joints_center = np.array([joints_x, joints_y])
|
||||
bbox_center = np.array(rec['center'])
|
||||
diff_norm2 = np.linalg.norm((joints_center-bbox_center), 2)
|
||||
ks = np.exp(-1.0*(diff_norm2**2) / ((0.2)**2*2.0*area))
|
||||
|
||||
metric = (0.2 / 16) * num_vis + 0.45 - 0.2 / 16
|
||||
if ks > metric:
|
||||
db_selected.append(rec)
|
||||
|
||||
logger.info('=> num db: {}'.format(len(db)))
|
||||
logger.info('=> num selected db: {}'.format(len(db_selected)))
|
||||
return db_selected
|
||||
|
||||
def generate_target(self, joints, joints_vis):
|
||||
'''
|
||||
:param joints: [num_joints, 3]
|
||||
:param joints_vis: [num_joints, 3]
|
||||
:return: target, target_weight(1: visible, 0: invisible)
|
||||
'''
|
||||
target_weight = np.ones((self.num_joints, 1), dtype=np.float32)
|
||||
target_weight[:, 0] = joints_vis[:, 0]
|
||||
|
||||
assert self.target_type == 'gaussian', \
|
||||
'Only support gaussian map now!'
|
||||
|
||||
if self.target_type == 'gaussian':
|
||||
target = np.zeros((self.num_joints,
|
||||
self.heatmap_size[1],
|
||||
self.heatmap_size[0]),
|
||||
dtype=np.float32)
|
||||
|
||||
tmp_size = self.sigma * 3
|
||||
|
||||
for joint_id in range(self.num_joints):
|
||||
feat_stride = self.image_size / self.heatmap_size
|
||||
mu_x = int(joints[joint_id][0] / feat_stride[0] + 0.5)
|
||||
mu_y = int(joints[joint_id][1] / feat_stride[1] + 0.5)
|
||||
# Check that any part of the gaussian is in-bounds
|
||||
ul = [int(mu_x - tmp_size), int(mu_y - tmp_size)]
|
||||
br = [int(mu_x + tmp_size + 1), int(mu_y + tmp_size + 1)]
|
||||
if ul[0] >= self.heatmap_size[0] or ul[1] >= self.heatmap_size[1] \
|
||||
or br[0] < 0 or br[1] < 0:
|
||||
# If not, just return the image as is
|
||||
target_weight[joint_id] = 0
|
||||
continue
|
||||
|
||||
# # Generate gaussian
|
||||
size = 2 * tmp_size + 1
|
||||
x = np.arange(0, size, 1, np.float32)
|
||||
y = x[:, np.newaxis]
|
||||
x0 = y0 = size // 2
|
||||
# The gaussian is not normalized, we want the center value to equal 1
|
||||
g = np.exp(- ((x - x0) ** 2 + (y - y0) ** 2) / (2 * self.sigma ** 2))
|
||||
|
||||
# Usable gaussian range
|
||||
g_x = max(0, -ul[0]), min(br[0], self.heatmap_size[0]) - ul[0]
|
||||
g_y = max(0, -ul[1]), min(br[1], self.heatmap_size[1]) - ul[1]
|
||||
# Image range
|
||||
img_x = max(0, ul[0]), min(br[0], self.heatmap_size[0])
|
||||
img_y = max(0, ul[1]), min(br[1], self.heatmap_size[1])
|
||||
|
||||
v = target_weight[joint_id]
|
||||
if v > 0.5:
|
||||
target[joint_id][img_y[0]:img_y[1], img_x[0]:img_x[1]] = \
|
||||
g[g_y[0]:g_y[1], g_x[0]:g_x[1]]
|
||||
|
||||
if self.use_different_joints_weight:
|
||||
target_weight = np.multiply(target_weight, self.joints_weight)
|
||||
|
||||
return target, target_weight
|
||||
@@ -0,0 +1,12 @@
|
||||
# ------------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft
|
||||
# Licensed under the MIT License.
|
||||
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
|
||||
# ------------------------------------------------------------------------------
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
from .mpii import MPIIDataset as mpii
|
||||
from .coco import COCODataset as coco
|
||||
@@ -0,0 +1,445 @@
|
||||
# ------------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft
|
||||
# Licensed under the MIT License.
|
||||
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
|
||||
# ------------------------------------------------------------------------------
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
from collections import defaultdict
|
||||
from collections import OrderedDict
|
||||
import logging
|
||||
import os
|
||||
|
||||
from pycocotools.coco import COCO
|
||||
from pycocotools.cocoeval import COCOeval
|
||||
import json_tricks as json
|
||||
import numpy as np
|
||||
|
||||
from dataset.JointsDataset import JointsDataset
|
||||
from nms.nms import oks_nms
|
||||
from nms.nms import soft_oks_nms
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class COCODataset(JointsDataset):
|
||||
'''
|
||||
"keypoints": {
|
||||
0: "nose",
|
||||
1: "left_eye",
|
||||
2: "right_eye",
|
||||
3: "left_ear",
|
||||
4: "right_ear",
|
||||
5: "left_shoulder",
|
||||
6: "right_shoulder",
|
||||
7: "left_elbow",
|
||||
8: "right_elbow",
|
||||
9: "left_wrist",
|
||||
10: "right_wrist",
|
||||
11: "left_hip",
|
||||
12: "right_hip",
|
||||
13: "left_knee",
|
||||
14: "right_knee",
|
||||
15: "left_ankle",
|
||||
16: "right_ankle"
|
||||
},
|
||||
"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
|
||||
self.image_thre = cfg.TEST.IMAGE_THRE
|
||||
self.soft_nms = cfg.TEST.SOFT_NMS
|
||||
self.oks_thre = cfg.TEST.OKS_THRE
|
||||
self.in_vis_thre = cfg.TEST.IN_VIS_THRE
|
||||
self.bbox_file = cfg.TEST.COCO_BBOX_FILE
|
||||
self.use_gt_bbox = cfg.TEST.USE_GT_BBOX
|
||||
self.image_width = cfg.MODEL.IMAGE_SIZE[0]
|
||||
self.image_height = cfg.MODEL.IMAGE_SIZE[1]
|
||||
self.aspect_ratio = self.image_width * 1.0 / self.image_height
|
||||
self.pixel_std = 200
|
||||
|
||||
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))
|
||||
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:]
|
||||
]
|
||||
)
|
||||
|
||||
# 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))
|
||||
|
||||
self.num_joints = 17
|
||||
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
|
||||
).reshape((self.num_joints, 1))
|
||||
|
||||
self.db = self._get_db()
|
||||
|
||||
if is_train and cfg.DATASET.SELECT_DATA:
|
||||
self.db = self.select_data(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'
|
||||
)
|
||||
|
||||
def _load_image_set_index(self):
|
||||
""" image id: int """
|
||||
image_ids = self.coco.getImgIds()
|
||||
return image_ids
|
||||
|
||||
def _get_db(self):
|
||||
if self.is_train or self.use_gt_bbox:
|
||||
# use ground truth bbox
|
||||
gt_db = self._load_coco_keypoint_annotations()
|
||||
else:
|
||||
# use bbox from detection
|
||||
gt_db = self._load_coco_person_detection_results()
|
||||
return gt_db
|
||||
|
||||
def _load_coco_keypoint_annotations(self):
|
||||
""" ground truth bbox and keypoints """
|
||||
gt_db = []
|
||||
for index in self.image_set_index:
|
||||
gt_db.extend(self._load_coco_keypoint_annotation_kernal(index))
|
||||
return gt_db
|
||||
|
||||
def _load_coco_keypoint_annotation_kernal(self, index):
|
||||
"""
|
||||
coco ann: [u'segmentation', u'area', u'iscrowd', u'image_id', u'bbox', u'category_id', u'id']
|
||||
iscrowd:
|
||||
crowd instances are handled by marking their overlaps with all categories to -1
|
||||
and later excluded in training
|
||||
bbox:
|
||||
[x1, y1, w, h]
|
||||
:param index: coco image id
|
||||
:return: db entry
|
||||
"""
|
||||
im_ann = self.coco.loadImgs(index)[0]
|
||||
width = im_ann['width']
|
||||
height = im_ann['height']
|
||||
|
||||
annIds = self.coco.getAnnIds(imgIds=index, iscrowd=False)
|
||||
objs = self.coco.loadAnns(annIds)
|
||||
|
||||
# sanitize bboxes
|
||||
valid_objs = []
|
||||
for obj in objs:
|
||||
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]
|
||||
valid_objs.append(obj)
|
||||
objs = valid_objs
|
||||
|
||||
rec = []
|
||||
for obj in objs:
|
||||
cls = self._coco_ind_to_class_ind[obj['category_id']]
|
||||
if cls != 1:
|
||||
continue
|
||||
|
||||
# ignore objs without keypoints annotation
|
||||
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)
|
||||
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, 2] = 0
|
||||
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,
|
||||
})
|
||||
|
||||
return rec
|
||||
|
||||
def _box2cs(self, box):
|
||||
x, y, w, h = box[:4]
|
||||
return self._xywh2cs(x, y, w, h)
|
||||
|
||||
def _xywh2cs(self, x, y, w, h):
|
||||
center = np.zeros((2), dtype=np.float32)
|
||||
center[0] = x + w * 0.5
|
||||
center[1] = y + h * 0.5
|
||||
|
||||
if w > self.aspect_ratio * 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)
|
||||
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
|
||||
|
||||
prefix = 'test2017' if 'test' in self.image_set else self.image_set
|
||||
|
||||
data_name = prefix + '.zip@' if self.data_format == 'zip' else prefix
|
||||
|
||||
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:
|
||||
all_boxes = json.load(f)[:10]
|
||||
|
||||
if not all_boxes:
|
||||
logger.error('=> Load %s fail!' % self.bbox_file)
|
||||
return None
|
||||
|
||||
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:
|
||||
continue
|
||||
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
|
||||
|
||||
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))
|
||||
return kpt_db
|
||||
|
||||
def evaluate(self, cfg, preds, output_dir, all_boxes, img_path,
|
||||
*args, **kwargs):
|
||||
rank = cfg.RANK
|
||||
|
||||
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))
|
||||
|
||||
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])
|
||||
})
|
||||
# image x person x (keypoints)
|
||||
kpts = defaultdict(list)
|
||||
for kpt in _kpts:
|
||||
kpts[kpt['image']].append(kpt)
|
||||
|
||||
# rescoring and oks nms
|
||||
num_joints = self.num_joints
|
||||
in_vis_thre = self.in_vis_thre
|
||||
oks_thre = self.oks_thre
|
||||
oks_nmsed_kpts = []
|
||||
for img in kpts.keys():
|
||||
img_kpts = kpts[img]
|
||||
for n_p in img_kpts:
|
||||
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]
|
||||
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
|
||||
|
||||
if self.soft_nms:
|
||||
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
|
||||
)
|
||||
|
||||
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)
|
||||
name_value = OrderedDict(info_str)
|
||||
return name_value, name_value['AP']
|
||||
else:
|
||||
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
|
||||
}
|
||||
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:
|
||||
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:
|
||||
for line in f:
|
||||
content.append(line)
|
||||
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_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
|
||||
)
|
||||
|
||||
for ipt in range(self.num_joints):
|
||||
key_points[:, ipt * 3 + 0] = _key_points[:, ipt, 0]
|
||||
key_points[:, ipt * 3 + 1] = _key_points[:, ipt, 1]
|
||||
key_points[:, ipt * 3 + 2] = _key_points[:, ipt, 2] # keypoints score.
|
||||
|
||||
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'])
|
||||
}
|
||||
for k in range(len(img_kpts))
|
||||
]
|
||||
cat_results.extend(result)
|
||||
|
||||
return cat_results
|
||||
|
||||
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.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)']
|
||||
|
||||
info_str = []
|
||||
for ind, name in enumerate(stats_names):
|
||||
info_str.append((name, coco_eval.stats[ind]))
|
||||
|
||||
return info_str
|
||||
@@ -0,0 +1,181 @@
|
||||
# ------------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft
|
||||
# Licensed under the MIT License.
|
||||
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
|
||||
# ------------------------------------------------------------------------------
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import logging
|
||||
import os
|
||||
import json_tricks as json
|
||||
from collections import OrderedDict
|
||||
|
||||
import numpy as np
|
||||
from scipy.io import loadmat, savemat
|
||||
|
||||
from dataset.JointsDataset import JointsDataset
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class MPIIDataset(JointsDataset):
|
||||
def __init__(self, cfg, root, image_set, is_train, transform=None):
|
||||
super().__init__(cfg, root, image_set, is_train, transform)
|
||||
|
||||
self.num_joints = 16
|
||||
self.flip_pairs = [[0, 5], [1, 4], [2, 3], [10, 15], [11, 14], [12, 13]]
|
||||
self.parent_ids = [1, 2, 6, 6, 3, 4, 6, 6, 7, 8, 11, 12, 7, 7, 13, 14]
|
||||
|
||||
self.upper_body_ids = (7, 8, 9, 10, 11, 12, 13, 14, 15)
|
||||
self.lower_body_ids = (0, 1, 2, 3, 4, 5, 6)
|
||||
|
||||
self.db = self._get_db()
|
||||
|
||||
if is_train and cfg.DATASET.SELECT_DATA:
|
||||
self.db = self.select_data(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'
|
||||
)
|
||||
with open(file_name) as anno_file:
|
||||
anno = json.load(anno_file)
|
||||
|
||||
gt_db = []
|
||||
for a in anno:
|
||||
image_name = a['image']
|
||||
|
||||
c = np.array(a['center'], dtype=np.float)
|
||||
s = np.array([a['scale'], a['scale']], dtype=np.float)
|
||||
|
||||
# Adjust center/scale slightly to avoid cropping limbs
|
||||
if c[0] != -1:
|
||||
c[1] = c[1] + 15 * s[1]
|
||||
s = s * 1.25
|
||||
|
||||
# MPII uses matlab format, index is based 1,
|
||||
# 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[:, 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_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'
|
||||
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,
|
||||
}
|
||||
)
|
||||
|
||||
return gt_db
|
||||
|
||||
def evaluate(self, cfg, preds, output_dir, *args, **kwargs):
|
||||
# convert 0-based index to 1-based index
|
||||
preds = preds[:, :, 0:2] + 1.0
|
||||
|
||||
if output_dir:
|
||||
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
|
||||
|
||||
SC_BIAS = 0.6
|
||||
threshold = 0.5
|
||||
|
||||
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']
|
||||
|
||||
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]
|
||||
|
||||
jnt_visible = 1 - jnt_missing
|
||||
uv_error = pos_pred_src - pos_gt_src
|
||||
uv_err = np.linalg.norm(uv_error, axis=1)
|
||||
headsizes = headboxes_src[1, :, :] - headboxes_src[0, :, :]
|
||||
headsizes = np.linalg.norm(headsizes, axis=0)
|
||||
headsizes *= SC_BIAS
|
||||
scale = np.multiply(headsizes, np.ones((len(uv_err), 1)))
|
||||
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)
|
||||
|
||||
# save
|
||||
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)
|
||||
|
||||
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)
|
||||
|
||||
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))
|
||||
]
|
||||
name_value = OrderedDict(name_value)
|
||||
|
||||
return name_value, name_value['Mean']
|
||||
Reference in New Issue
Block a user