完善部署并训练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:
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,3 @@
|
||||
from .wider_face import WiderFaceDetection, detection_collate
|
||||
from .data_augment import *
|
||||
from .config import *
|
||||
@@ -0,0 +1,42 @@
|
||||
# config.py
|
||||
|
||||
cfg_mnet = {
|
||||
'name': 'mobilenet0.25',
|
||||
'min_sizes': [[16, 32], [64, 128], [256, 512]],
|
||||
'steps': [8, 16, 32],
|
||||
'variance': [0.1, 0.2],
|
||||
'clip': False,
|
||||
'loc_weight': 2.0,
|
||||
'gpu_train': True,
|
||||
'batch_size': 32,
|
||||
'ngpu': 1,
|
||||
'epoch': 250,
|
||||
'decay1': 190,
|
||||
'decay2': 220,
|
||||
'image_size': 640,
|
||||
'pretrain': True,
|
||||
'return_layers': {'stage1': 1, 'stage2': 2, 'stage3': 3},
|
||||
'in_channel': 32,
|
||||
'out_channel': 64
|
||||
}
|
||||
|
||||
cfg_re50 = {
|
||||
'name': 'Resnet50',
|
||||
'min_sizes': [[16, 32], [64, 128], [256, 512]],
|
||||
'steps': [8, 16, 32],
|
||||
'variance': [0.1, 0.2],
|
||||
'clip': False,
|
||||
'loc_weight': 2.0,
|
||||
'gpu_train': True,
|
||||
'batch_size': 24,
|
||||
'ngpu': 4,
|
||||
'epoch': 100,
|
||||
'decay1': 70,
|
||||
'decay2': 90,
|
||||
'image_size': 840,
|
||||
'pretrain': True,
|
||||
'return_layers': {'layer2': 1, 'layer3': 2, 'layer4': 3},
|
||||
'in_channel': 256,
|
||||
'out_channel': 256
|
||||
}
|
||||
|
||||
@@ -0,0 +1,237 @@
|
||||
import cv2
|
||||
import numpy as np
|
||||
import random
|
||||
from core.utils.box_utils_Retina import matrix_iof
|
||||
|
||||
|
||||
def _crop(image, boxes, labels, landm, img_dim):
|
||||
height, width, _ = image.shape
|
||||
pad_image_flag = True
|
||||
|
||||
for _ in range(250):
|
||||
"""
|
||||
if random.uniform(0, 1) <= 0.2:
|
||||
scale = 1.0
|
||||
else:
|
||||
scale = random.uniform(0.3, 1.0)
|
||||
"""
|
||||
PRE_SCALES = [0.3, 0.45, 0.6, 0.8, 1.0]
|
||||
scale = random.choice(PRE_SCALES)
|
||||
short_side = min(width, height)
|
||||
w = int(scale * short_side)
|
||||
h = w
|
||||
|
||||
if width == w:
|
||||
l = 0
|
||||
else:
|
||||
l = random.randrange(width - w)
|
||||
if height == h:
|
||||
t = 0
|
||||
else:
|
||||
t = random.randrange(height - h)
|
||||
roi = np.array((l, t, l + w, t + h))
|
||||
|
||||
value = matrix_iof(boxes, roi[np.newaxis])
|
||||
flag = (value >= 1)
|
||||
if not flag.any():
|
||||
continue
|
||||
|
||||
centers = (boxes[:, :2] + boxes[:, 2:]) / 2
|
||||
mask_a = np.logical_and(roi[:2] < centers, centers < roi[2:]).all(axis=1)
|
||||
boxes_t = boxes[mask_a].copy()
|
||||
labels_t = labels[mask_a].copy()
|
||||
landms_t = landm[mask_a].copy()
|
||||
landms_t = landms_t.reshape([-1, 5, 2])
|
||||
|
||||
if boxes_t.shape[0] == 0:
|
||||
continue
|
||||
|
||||
image_t = image[roi[1]:roi[3], roi[0]:roi[2]]
|
||||
|
||||
boxes_t[:, :2] = np.maximum(boxes_t[:, :2], roi[:2])
|
||||
boxes_t[:, :2] -= roi[:2]
|
||||
boxes_t[:, 2:] = np.minimum(boxes_t[:, 2:], roi[2:])
|
||||
boxes_t[:, 2:] -= roi[:2]
|
||||
|
||||
# landm
|
||||
landms_t[:, :, :2] = landms_t[:, :, :2] - roi[:2]
|
||||
landms_t[:, :, :2] = np.maximum(landms_t[:, :, :2], np.array([0, 0]))
|
||||
landms_t[:, :, :2] = np.minimum(landms_t[:, :, :2], roi[2:] - roi[:2])
|
||||
landms_t = landms_t.reshape([-1, 10])
|
||||
|
||||
|
||||
# make sure that the cropped image contains at least one face > 16 pixel at training image scale
|
||||
b_w_t = (boxes_t[:, 2] - boxes_t[:, 0] + 1) / w * img_dim
|
||||
b_h_t = (boxes_t[:, 3] - boxes_t[:, 1] + 1) / h * img_dim
|
||||
mask_b = np.minimum(b_w_t, b_h_t) > 0.0
|
||||
boxes_t = boxes_t[mask_b]
|
||||
labels_t = labels_t[mask_b]
|
||||
landms_t = landms_t[mask_b]
|
||||
|
||||
if boxes_t.shape[0] == 0:
|
||||
continue
|
||||
|
||||
pad_image_flag = False
|
||||
|
||||
return image_t, boxes_t, labels_t, landms_t, pad_image_flag
|
||||
return image, boxes, labels, landm, pad_image_flag
|
||||
|
||||
|
||||
def _distort(image):
|
||||
|
||||
def _convert(image, alpha=1, beta=0):
|
||||
tmp = image.astype(float) * alpha + beta
|
||||
tmp[tmp < 0] = 0
|
||||
tmp[tmp > 255] = 255
|
||||
image[:] = tmp
|
||||
|
||||
image = image.copy()
|
||||
|
||||
if random.randrange(2):
|
||||
|
||||
#brightness distortion
|
||||
if random.randrange(2):
|
||||
_convert(image, beta=random.uniform(-32, 32))
|
||||
|
||||
#contrast distortion
|
||||
if random.randrange(2):
|
||||
_convert(image, alpha=random.uniform(0.5, 1.5))
|
||||
|
||||
image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
|
||||
|
||||
#saturation distortion
|
||||
if random.randrange(2):
|
||||
_convert(image[:, :, 1], alpha=random.uniform(0.5, 1.5))
|
||||
|
||||
#hue distortion
|
||||
if random.randrange(2):
|
||||
tmp = image[:, :, 0].astype(int) + random.randint(-18, 18)
|
||||
tmp %= 180
|
||||
image[:, :, 0] = tmp
|
||||
|
||||
image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR)
|
||||
|
||||
else:
|
||||
|
||||
#brightness distortion
|
||||
if random.randrange(2):
|
||||
_convert(image, beta=random.uniform(-32, 32))
|
||||
|
||||
image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
|
||||
|
||||
#saturation distortion
|
||||
if random.randrange(2):
|
||||
_convert(image[:, :, 1], alpha=random.uniform(0.5, 1.5))
|
||||
|
||||
#hue distortion
|
||||
if random.randrange(2):
|
||||
tmp = image[:, :, 0].astype(int) + random.randint(-18, 18)
|
||||
tmp %= 180
|
||||
image[:, :, 0] = tmp
|
||||
|
||||
image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR)
|
||||
|
||||
#contrast distortion
|
||||
if random.randrange(2):
|
||||
_convert(image, alpha=random.uniform(0.5, 1.5))
|
||||
|
||||
return image
|
||||
|
||||
|
||||
def _expand(image, boxes, fill, p):
|
||||
if random.randrange(2):
|
||||
return image, boxes
|
||||
|
||||
height, width, depth = image.shape
|
||||
|
||||
scale = random.uniform(1, p)
|
||||
w = int(scale * width)
|
||||
h = int(scale * height)
|
||||
|
||||
left = random.randint(0, w - width)
|
||||
top = random.randint(0, h - height)
|
||||
|
||||
boxes_t = boxes.copy()
|
||||
boxes_t[:, :2] += (left, top)
|
||||
boxes_t[:, 2:] += (left, top)
|
||||
expand_image = np.empty(
|
||||
(h, w, depth),
|
||||
dtype=image.dtype)
|
||||
expand_image[:, :] = fill
|
||||
expand_image[top:top + height, left:left + width] = image
|
||||
image = expand_image
|
||||
|
||||
return image, boxes_t
|
||||
|
||||
|
||||
def _mirror(image, boxes, landms):
|
||||
_, width, _ = image.shape
|
||||
if random.randrange(2):
|
||||
image = image[:, ::-1]
|
||||
boxes = boxes.copy()
|
||||
boxes[:, 0::2] = width - boxes[:, 2::-2]
|
||||
|
||||
# landm
|
||||
landms = landms.copy()
|
||||
landms = landms.reshape([-1, 5, 2])
|
||||
landms[:, :, 0] = width - landms[:, :, 0]
|
||||
tmp = landms[:, 1, :].copy()
|
||||
landms[:, 1, :] = landms[:, 0, :]
|
||||
landms[:, 0, :] = tmp
|
||||
tmp1 = landms[:, 4, :].copy()
|
||||
landms[:, 4, :] = landms[:, 3, :]
|
||||
landms[:, 3, :] = tmp1
|
||||
landms = landms.reshape([-1, 10])
|
||||
|
||||
return image, boxes, landms
|
||||
|
||||
|
||||
def _pad_to_square(image, rgb_mean, pad_image_flag):
|
||||
if not pad_image_flag:
|
||||
return image
|
||||
height, width, _ = image.shape
|
||||
long_side = max(width, height)
|
||||
image_t = np.empty((long_side, long_side, 3), dtype=image.dtype)
|
||||
image_t[:, :] = rgb_mean
|
||||
image_t[0:0 + height, 0:0 + width] = image
|
||||
return image_t
|
||||
|
||||
|
||||
def _resize_subtract_mean(image, insize, rgb_mean):
|
||||
interp_methods = [cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_AREA, cv2.INTER_NEAREST, cv2.INTER_LANCZOS4]
|
||||
interp_method = interp_methods[random.randrange(5)]
|
||||
image = cv2.resize(image, (insize, insize), interpolation=interp_method)
|
||||
image = image.astype(np.float32)
|
||||
image -= rgb_mean
|
||||
return image.transpose(2, 0, 1)
|
||||
|
||||
|
||||
class preproc(object):
|
||||
|
||||
def __init__(self, img_dim, rgb_means):
|
||||
self.img_dim = img_dim
|
||||
self.rgb_means = rgb_means
|
||||
|
||||
def __call__(self, image, targets):
|
||||
assert targets.shape[0] > 0, "this image does not have gt"
|
||||
|
||||
boxes = targets[:, :4].copy()
|
||||
labels = targets[:, -1].copy()
|
||||
landm = targets[:, 4:-1].copy()
|
||||
|
||||
image_t, boxes_t, labels_t, landm_t, pad_image_flag = _crop(image, boxes, labels, landm, self.img_dim)
|
||||
image_t = _distort(image_t)
|
||||
image_t = _pad_to_square(image_t,self.rgb_means, pad_image_flag)
|
||||
image_t, boxes_t, landm_t = _mirror(image_t, boxes_t, landm_t)
|
||||
height, width, _ = image_t.shape
|
||||
image_t = _resize_subtract_mean(image_t, self.img_dim, self.rgb_means)
|
||||
boxes_t[:, 0::2] /= width
|
||||
boxes_t[:, 1::2] /= height
|
||||
|
||||
landm_t[:, 0::2] /= width
|
||||
landm_t[:, 1::2] /= height
|
||||
|
||||
labels_t = np.expand_dims(labels_t, 1)
|
||||
targets_t = np.hstack((boxes_t, landm_t, labels_t))
|
||||
|
||||
return image_t, targets_t
|
||||
@@ -0,0 +1,101 @@
|
||||
import os
|
||||
import os.path
|
||||
import sys
|
||||
import torch
|
||||
import torch.utils.data as data
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
class WiderFaceDetection(data.Dataset):
|
||||
def __init__(self, txt_path, preproc=None):
|
||||
self.preproc = preproc
|
||||
self.imgs_path = []
|
||||
self.words = []
|
||||
f = open(txt_path,'r')
|
||||
lines = f.readlines()
|
||||
isFirst = True
|
||||
labels = []
|
||||
for line in lines:
|
||||
line = line.rstrip()
|
||||
if line.startswith('#'):
|
||||
if isFirst is True:
|
||||
isFirst = False
|
||||
else:
|
||||
labels_copy = labels.copy()
|
||||
self.words.append(labels_copy)
|
||||
labels.clear()
|
||||
path = line[2:]
|
||||
path = txt_path.replace('label.txt','images/') + path
|
||||
self.imgs_path.append(path)
|
||||
else:
|
||||
line = line.split(' ')
|
||||
label = [float(x) for x in line]
|
||||
labels.append(label)
|
||||
|
||||
self.words.append(labels)
|
||||
|
||||
def __len__(self):
|
||||
return len(self.imgs_path)
|
||||
|
||||
def __getitem__(self, index):
|
||||
img = cv2.imread(self.imgs_path[index])
|
||||
height, width, _ = img.shape
|
||||
|
||||
labels = self.words[index]
|
||||
annotations = np.zeros((0, 15))
|
||||
if len(labels) == 0:
|
||||
return annotations
|
||||
for idx, label in enumerate(labels):
|
||||
annotation = np.zeros((1, 15))
|
||||
# bbox
|
||||
annotation[0, 0] = label[0] # x1
|
||||
annotation[0, 1] = label[1] # y1
|
||||
annotation[0, 2] = label[0] + label[2] # x2
|
||||
annotation[0, 3] = label[1] + label[3] # y2
|
||||
|
||||
# landmarks
|
||||
annotation[0, 4] = label[4] # l0_x
|
||||
annotation[0, 5] = label[5] # l0_y
|
||||
annotation[0, 6] = label[7] # l1_x
|
||||
annotation[0, 7] = label[8] # l1_y
|
||||
annotation[0, 8] = label[10] # l2_x
|
||||
annotation[0, 9] = label[11] # l2_y
|
||||
annotation[0, 10] = label[13] # l3_x
|
||||
annotation[0, 11] = label[14] # l3_y
|
||||
annotation[0, 12] = label[16] # l4_x
|
||||
annotation[0, 13] = label[17] # l4_y
|
||||
if (annotation[0, 4]<0):
|
||||
annotation[0, 14] = -1
|
||||
else:
|
||||
annotation[0, 14] = 1
|
||||
|
||||
annotations = np.append(annotations, annotation, axis=0)
|
||||
target = np.array(annotations)
|
||||
if self.preproc is not None:
|
||||
img, target = self.preproc(img, target)
|
||||
|
||||
return torch.from_numpy(img), target
|
||||
|
||||
def detection_collate(batch):
|
||||
"""Custom collate fn for dealing with batches of images that have a different
|
||||
number of associated object annotations (bounding boxes).
|
||||
|
||||
Arguments:
|
||||
batch: (tuple) A tuple of tensor images and lists of annotations
|
||||
|
||||
Return:
|
||||
A tuple containing:
|
||||
1) (tensor) batch of images stacked on their 0 dim
|
||||
2) (list of tensors) annotations for a given image are stacked on 0 dim
|
||||
"""
|
||||
targets = []
|
||||
imgs = []
|
||||
for _, sample in enumerate(batch):
|
||||
for _, tup in enumerate(sample):
|
||||
if torch.is_tensor(tup):
|
||||
imgs.append(tup)
|
||||
elif isinstance(tup, type(np.empty(0))):
|
||||
annos = torch.from_numpy(tup).float()
|
||||
targets.append(annos)
|
||||
|
||||
return (torch.stack(imgs, 0), targets)
|
||||
@@ -5,16 +5,24 @@ import os
|
||||
|
||||
class OSS_object():
|
||||
def __init__(self):
|
||||
access_key_id = os.getenv('OSS_TEST_ACCESS_KEY_ID', '<your-access-key-id>')
|
||||
access_key_secret = os.getenv('OSS_TEST_ACCESS_KEY_SECRET', '<your-access-key-secret>')
|
||||
bucket_name = os.getenv('OSS_TEST_BUCKET', '<your-bucket-name>')
|
||||
endpoint = os.getenv('OSS_TEST_ENDPOINT', '<your-endpoint>')
|
||||
for param in (access_key_id, access_key_secret, bucket_name, endpoint):
|
||||
assert '<' not in param, '请设置环境变量 OSS_TEST_ACCESS_KEY_ID / OSS_TEST_ACCESS_KEY_SECRET / OSS_TEST_BUCKET / OSS_TEST_ENDPOINT'
|
||||
# 懒加载：只读取凭证，不立即连接 OSS。
|
||||
# 这样本地用 output_format=base64 时不配 OSS 凭证也能启动服务；
|
||||
# 仅在真正调用 upload_file 时才校验并连接。
|
||||
self.access_key_id = os.getenv('OSS_TEST_ACCESS_KEY_ID', '<your-access-key-id>')
|
||||
self.access_key_secret = os.getenv('OSS_TEST_ACCESS_KEY_SECRET', '<your-access-key-secret>')
|
||||
self.bucket_name = os.getenv('OSS_TEST_BUCKET', '<your-bucket-name>')
|
||||
self.endpoint = os.getenv('OSS_TEST_ENDPOINT', '<your-endpoint>')
|
||||
self.bucket = None
|
||||
|
||||
self.bucket = oss2.Bucket(oss2.Auth(access_key_id, access_key_secret), endpoint, bucket_name)
|
||||
def _ensure_bucket(self):
|
||||
if self.bucket is not None:
|
||||
return
|
||||
for param in (self.access_key_id, self.access_key_secret, self.bucket_name, self.endpoint):
|
||||
assert '<' not in param, '请设置环境变量 OSS_TEST_ACCESS_KEY_ID / OSS_TEST_ACCESS_KEY_SECRET / OSS_TEST_BUCKET / OSS_TEST_ENDPOINT'
|
||||
self.bucket = oss2.Bucket(oss2.Auth(self.access_key_id, self.access_key_secret), self.endpoint, self.bucket_name)
|
||||
|
||||
def upload_file(self, file, target_name):
|
||||
self._ensure_bucket()
|
||||
t0 = time.time()
|
||||
with open(oss2.to_unicode(file), 'rb') as f:
|
||||
ret = self.bucket.put_object(target_name, f)
|
||||
|
||||
@@ -177,10 +177,17 @@ def list_hairstyles(hairstyle_dir, train_dir, upload_dir, limit=None):
|
||||
返回:
|
||||
[{hair_id, gender, has_lora}, ...]
|
||||
"""
|
||||
# 仅展示这些发型（其余数据保留在磁盘，但不在列表中显示）
|
||||
_VISIBLE_HAIRSTYLES = {
|
||||
"chang_tuoyuan", "chang_bolang", "chang_zhixian",
|
||||
"chang_huaban", "chang_xinxing",
|
||||
}
|
||||
styles = []
|
||||
if not os.path.isdir(hairstyle_dir):
|
||||
return styles
|
||||
for hair_id in os.listdir(hairstyle_dir):
|
||||
if hair_id not in _VISIBLE_HAIRSTYLES:
|
||||
continue
|
||||
cfg_path = os.path.join(hairstyle_dir, hair_id, "config.json")
|
||||
ref_path = os.path.join(hairstyle_dir, hair_id, "ref_rgb_8uc3_768.png")
|
||||
lora_path = os.path.join(train_dir, hair_id, "model", "hairstyle_hd_lora.safetensors")
|
||||
|
||||
@@ -0,0 +1,9 @@
|
||||
# ------------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft
|
||||
# Licensed under the MIT License.
|
||||
# Written by Bin Xiao (Bin.Xiao@microsoft.com)
|
||||
# ------------------------------------------------------------------------------
|
||||
|
||||
from .default import _C as cfg
|
||||
from .default import update_config
|
||||
from .models import MODEL_EXTRAS
|
||||
@@ -0,0 +1,160 @@
|
||||
|
||||
# ------------------------------------------------------------------------------
|
||||
# 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 os
|
||||
|
||||
from yacs.config import CfgNode as CN
|
||||
|
||||
|
||||
_C = CN()
|
||||
|
||||
_C.OUTPUT_DIR = ''
|
||||
_C.LOG_DIR = ''
|
||||
_C.DATA_DIR = ''
|
||||
_C.GPUS = (0,)
|
||||
_C.WORKERS = 4
|
||||
_C.PRINT_FREQ = 20
|
||||
_C.AUTO_RESUME = False
|
||||
_C.PIN_MEMORY = True
|
||||
_C.RANK = 0
|
||||
|
||||
# Cudnn related params
|
||||
_C.CUDNN = CN()
|
||||
_C.CUDNN.BENCHMARK = True
|
||||
_C.CUDNN.DETERMINISTIC = False
|
||||
_C.CUDNN.ENABLED = True
|
||||
|
||||
# common params for NETWORK
|
||||
_C.MODEL = CN()
|
||||
_C.MODEL.NAME = 'pose_hrnet'
|
||||
_C.MODEL.INIT_WEIGHTS = True
|
||||
_C.MODEL.PRETRAINED = ''
|
||||
_C.MODEL.NUM_JOINTS = 17
|
||||
_C.MODEL.TAG_PER_JOINT = True
|
||||
_C.MODEL.TARGET_TYPE = 'gaussian'
|
||||
_C.MODEL.IMAGE_SIZE = [256, 256] # width * height, ex: 192 * 256
|
||||
_C.MODEL.HEATMAP_SIZE = [64, 64] # width * height, ex: 24 * 32
|
||||
_C.MODEL.SIGMA = 2
|
||||
_C.MODEL.EXTRA = CN(new_allowed=True)
|
||||
|
||||
_C.LOSS = CN()
|
||||
_C.LOSS.USE_OHKM = False
|
||||
_C.LOSS.TOPK = 8
|
||||
_C.LOSS.USE_TARGET_WEIGHT = True
|
||||
_C.LOSS.USE_DIFFERENT_JOINTS_WEIGHT = False
|
||||
|
||||
# DATASET related params
|
||||
_C.DATASET = CN()
|
||||
_C.DATASET.ROOT = ''
|
||||
_C.DATASET.DATASET = 'mpii'
|
||||
_C.DATASET.TRAIN_SET = 'train'
|
||||
_C.DATASET.TEST_SET = 'valid'
|
||||
_C.DATASET.DATA_FORMAT = 'jpg'
|
||||
_C.DATASET.HYBRID_JOINTS_TYPE = ''
|
||||
_C.DATASET.SELECT_DATA = False
|
||||
|
||||
# training data augmentation
|
||||
_C.DATASET.FLIP = True
|
||||
_C.DATASET.SCALE_FACTOR = 0.25
|
||||
_C.DATASET.ROT_FACTOR = 30
|
||||
_C.DATASET.PROB_HALF_BODY = 0.0
|
||||
_C.DATASET.NUM_JOINTS_HALF_BODY = 8
|
||||
_C.DATASET.COLOR_RGB = False
|
||||
|
||||
# train
|
||||
_C.TRAIN = CN()
|
||||
|
||||
_C.TRAIN.LR_FACTOR = 0.1
|
||||
_C.TRAIN.LR_STEP = [90, 110]
|
||||
_C.TRAIN.LR = 0.001
|
||||
|
||||
_C.TRAIN.OPTIMIZER = 'adam'
|
||||
_C.TRAIN.MOMENTUM = 0.9
|
||||
_C.TRAIN.WD = 0.0001
|
||||
_C.TRAIN.NESTEROV = False
|
||||
_C.TRAIN.GAMMA1 = 0.99
|
||||
_C.TRAIN.GAMMA2 = 0.0
|
||||
|
||||
_C.TRAIN.BEGIN_EPOCH = 0
|
||||
_C.TRAIN.END_EPOCH = 140
|
||||
|
||||
_C.TRAIN.RESUME = False
|
||||
_C.TRAIN.CHECKPOINT = ''
|
||||
|
||||
_C.TRAIN.BATCH_SIZE_PER_GPU = 32
|
||||
_C.TRAIN.SHUFFLE = True
|
||||
|
||||
# testing
|
||||
_C.TEST = CN()
|
||||
|
||||
# size of images for each device
|
||||
_C.TEST.BATCH_SIZE_PER_GPU = 32
|
||||
# Test Model Epoch
|
||||
_C.TEST.FLIP_TEST = False
|
||||
_C.TEST.POST_PROCESS = False
|
||||
_C.TEST.SHIFT_HEATMAP = False
|
||||
|
||||
_C.TEST.USE_GT_BBOX = False
|
||||
|
||||
# nms
|
||||
_C.TEST.IMAGE_THRE = 0.1
|
||||
_C.TEST.NMS_THRE = 0.6
|
||||
_C.TEST.SOFT_NMS = False
|
||||
_C.TEST.OKS_THRE = 0.5
|
||||
_C.TEST.IN_VIS_THRE = 0.0
|
||||
_C.TEST.COCO_BBOX_FILE = ''
|
||||
_C.TEST.BBOX_THRE = 1.0
|
||||
_C.TEST.MODEL_FILE = ''
|
||||
|
||||
# debug
|
||||
_C.DEBUG = CN()
|
||||
_C.DEBUG.DEBUG = False
|
||||
_C.DEBUG.SAVE_BATCH_IMAGES_GT = False
|
||||
_C.DEBUG.SAVE_BATCH_IMAGES_PRED = False
|
||||
_C.DEBUG.SAVE_HEATMAPS_GT = False
|
||||
_C.DEBUG.SAVE_HEATMAPS_PRED = False
|
||||
|
||||
|
||||
def update_config(cfg, args):
|
||||
cfg.defrost()
|
||||
cfg.merge_from_file(args.cfg)
|
||||
cfg.merge_from_list(args.opts)
|
||||
|
||||
if args.modelDir:
|
||||
cfg.OUTPUT_DIR = args.modelDir
|
||||
|
||||
if args.logDir:
|
||||
cfg.LOG_DIR = args.logDir
|
||||
|
||||
if args.dataDir:
|
||||
cfg.DATA_DIR = args.dataDir
|
||||
|
||||
cfg.DATASET.ROOT = os.path.join(
|
||||
cfg.DATA_DIR, cfg.DATASET.ROOT
|
||||
)
|
||||
|
||||
cfg.MODEL.PRETRAINED = os.path.join(
|
||||
cfg.DATA_DIR, cfg.MODEL.PRETRAINED
|
||||
)
|
||||
|
||||
if cfg.TEST.MODEL_FILE:
|
||||
cfg.TEST.MODEL_FILE = os.path.join(
|
||||
cfg.DATA_DIR, cfg.TEST.MODEL_FILE
|
||||
)
|
||||
|
||||
cfg.freeze()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
import sys
|
||||
with open(sys.argv[1], 'w') as f:
|
||||
print(_C, file=f)
|
||||
|
||||
@@ -0,0 +1,58 @@
|
||||
# ------------------------------------------------------------------------------
|
||||
# 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 yacs.config import CfgNode as CN
|
||||
|
||||
|
||||
# pose_resnet related params
|
||||
POSE_RESNET = CN()
|
||||
POSE_RESNET.NUM_LAYERS = 50
|
||||
POSE_RESNET.DECONV_WITH_BIAS = False
|
||||
POSE_RESNET.NUM_DECONV_LAYERS = 3
|
||||
POSE_RESNET.NUM_DECONV_FILTERS = [256, 256, 256]
|
||||
POSE_RESNET.NUM_DECONV_KERNELS = [4, 4, 4]
|
||||
POSE_RESNET.FINAL_CONV_KERNEL = 1
|
||||
POSE_RESNET.PRETRAINED_LAYERS = ['*']
|
||||
|
||||
# pose_multi_resoluton_net related params
|
||||
POSE_HIGH_RESOLUTION_NET = CN()
|
||||
POSE_HIGH_RESOLUTION_NET.PRETRAINED_LAYERS = ['*']
|
||||
POSE_HIGH_RESOLUTION_NET.STEM_INPLANES = 64
|
||||
POSE_HIGH_RESOLUTION_NET.FINAL_CONV_KERNEL = 1
|
||||
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE2 = CN()
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE2.NUM_MODULES = 1
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE2.NUM_BRANCHES = 2
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE2.NUM_BLOCKS = [4, 4]
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE2.NUM_CHANNELS = [32, 64]
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE2.BLOCK = 'BASIC'
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE2.FUSE_METHOD = 'SUM'
|
||||
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE3 = CN()
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE3.NUM_MODULES = 1
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE3.NUM_BRANCHES = 3
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE3.NUM_BLOCKS = [4, 4, 4]
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE3.NUM_CHANNELS = [32, 64, 128]
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE3.BLOCK = 'BASIC'
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE3.FUSE_METHOD = 'SUM'
|
||||
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE4 = CN()
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE4.NUM_MODULES = 1
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE4.NUM_BRANCHES = 4
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE4.NUM_BLOCKS = [4, 4, 4, 4]
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE4.NUM_CHANNELS = [32, 64, 128, 256]
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE4.BLOCK = 'BASIC'
|
||||
POSE_HIGH_RESOLUTION_NET.STAGE4.FUSE_METHOD = 'SUM'
|
||||
|
||||
|
||||
MODEL_EXTRAS = {
|
||||
'pose_resnet': POSE_RESNET,
|
||||
'pose_high_resolution_net': POSE_HIGH_RESOLUTION_NET,
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
# ------------------------------------------------------------------------------
|
||||
# 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 numpy as np
|
||||
|
||||
from core.inference import get_max_preds
|
||||
|
||||
|
||||
def calc_dists(preds, target, normalize):
|
||||
preds = preds.astype(np.float32)
|
||||
target = target.astype(np.float32)
|
||||
dists = np.zeros((preds.shape[1], preds.shape[0]))
|
||||
for n in range(preds.shape[0]):
|
||||
for c in range(preds.shape[1]):
|
||||
if target[n, c, 0] > 1 and target[n, c, 1] > 1:
|
||||
normed_preds = preds[n, c, :] / normalize[n]
|
||||
normed_targets = target[n, c, :] / normalize[n]
|
||||
dists[c, n] = np.linalg.norm(normed_preds - normed_targets)
|
||||
else:
|
||||
dists[c, n] = -1
|
||||
return dists
|
||||
|
||||
|
||||
def dist_acc(dists, thr=0.5):
|
||||
''' Return percentage below threshold while ignoring values with a -1 '''
|
||||
dist_cal = np.not_equal(dists, -1)
|
||||
num_dist_cal = dist_cal.sum()
|
||||
if num_dist_cal > 0:
|
||||
return np.less(dists[dist_cal], thr).sum() * 1.0 / num_dist_cal
|
||||
else:
|
||||
return -1
|
||||
|
||||
|
||||
def accuracy(output, target, hm_type='gaussian', thr=0.5):
|
||||
'''
|
||||
Calculate accuracy according to PCK,
|
||||
but uses ground truth heatmap rather than x,y locations
|
||||
First value to be returned is average accuracy across 'idxs',
|
||||
followed by individual accuracies
|
||||
'''
|
||||
idx = list(range(output.shape[1]))
|
||||
norm = 1.0
|
||||
if hm_type == 'gaussian':
|
||||
pred, _ = get_max_preds(output)
|
||||
target, _ = get_max_preds(target)
|
||||
h = output.shape[2]
|
||||
w = output.shape[3]
|
||||
norm = np.ones((pred.shape[0], 2)) * np.array([h, w]) / 10
|
||||
dists = calc_dists(pred, target, norm)
|
||||
|
||||
acc = np.zeros((len(idx) + 1))
|
||||
avg_acc = 0
|
||||
cnt = 0
|
||||
|
||||
for i in range(len(idx)):
|
||||
acc[i + 1] = dist_acc(dists[idx[i]])
|
||||
if acc[i + 1] >= 0:
|
||||
avg_acc = avg_acc + acc[i + 1]
|
||||
cnt += 1
|
||||
|
||||
avg_acc = avg_acc / cnt if cnt != 0 else 0
|
||||
if cnt != 0:
|
||||
acc[0] = avg_acc
|
||||
return acc, avg_acc, cnt, pred
|
||||
|
||||
|
||||
@@ -0,0 +1,279 @@
|
||||
# ------------------------------------------------------------------------------
|
||||
# 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 time
|
||||
import logging
|
||||
import os
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from core.evaluate import accuracy
|
||||
from core.inference import get_final_preds
|
||||
from utils.transforms import flip_back
|
||||
from utils.vis import save_debug_images
|
||||
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def train(config, train_loader, model, criterion, optimizer, epoch,
|
||||
output_dir, tb_log_dir, writer_dict):
|
||||
batch_time = AverageMeter()
|
||||
data_time = AverageMeter()
|
||||
losses = AverageMeter()
|
||||
acc = AverageMeter()
|
||||
|
||||
# switch to train mode
|
||||
model.train()
|
||||
|
||||
end = time.time()
|
||||
for i, (input, target, target_weight, meta) in enumerate(train_loader):
|
||||
# measure data loading time
|
||||
data_time.update(time.time() - end)
|
||||
|
||||
# compute output
|
||||
outputs = model(input)
|
||||
|
||||
target = target.cuda(non_blocking=True)
|
||||
target_weight = target_weight.cuda(non_blocking=True)
|
||||
|
||||
if isinstance(outputs, list):
|
||||
loss = criterion(outputs[0], target, target_weight)
|
||||
for output in outputs[1:]:
|
||||
loss += criterion(output, target, target_weight)
|
||||
else:
|
||||
output = outputs
|
||||
loss = criterion(output, target, target_weight)
|
||||
|
||||
# loss = criterion(output, target, target_weight)
|
||||
|
||||
# compute gradient and do update step
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
|
||||
# measure accuracy and record loss
|
||||
losses.update(loss.item(), input.size(0))
|
||||
|
||||
_, avg_acc, cnt, pred = accuracy(output.detach().cpu().numpy(),
|
||||
target.detach().cpu().numpy())
|
||||
acc.update(avg_acc, cnt)
|
||||
|
||||
# measure elapsed time
|
||||
batch_time.update(time.time() - end)
|
||||
end = time.time()
|
||||
|
||||
if i % config.PRINT_FREQ == 0:
|
||||
msg = 'Epoch: [{0}][{1}/{2}]\t' \
|
||||
'Time {batch_time.val:.3f}s ({batch_time.avg:.3f}s)\t' \
|
||||
'Speed {speed:.1f} samples/s\t' \
|
||||
'Data {data_time.val:.3f}s ({data_time.avg:.3f}s)\t' \
|
||||
'Loss {loss.val:.5f} ({loss.avg:.5f})\t' \
|
||||
'Accuracy {acc.val:.3f} ({acc.avg:.3f})'.format(
|
||||
epoch, i, len(train_loader), batch_time=batch_time,
|
||||
speed=input.size(0)/batch_time.val,
|
||||
data_time=data_time, loss=losses, acc=acc)
|
||||
logger.info(msg)
|
||||
|
||||
writer = writer_dict['writer']
|
||||
global_steps = writer_dict['train_global_steps']
|
||||
writer.add_scalar('train_loss', losses.val, global_steps)
|
||||
writer.add_scalar('train_acc', acc.val, global_steps)
|
||||
writer_dict['train_global_steps'] = global_steps + 1
|
||||
|
||||
prefix = '{}_{}'.format(os.path.join(output_dir, 'train'), i)
|
||||
save_debug_images(config, input, meta, target, pred*4, output,
|
||||
prefix)
|
||||
|
||||
|
||||
def validate(config, val_loader, val_dataset, model, criterion, output_dir,
|
||||
tb_log_dir, writer_dict=None):
|
||||
batch_time = AverageMeter()
|
||||
losses = AverageMeter()
|
||||
acc = AverageMeter()
|
||||
|
||||
# switch to evaluate mode
|
||||
model.eval()
|
||||
|
||||
num_samples = len(val_dataset)
|
||||
all_preds = np.zeros(
|
||||
(num_samples, config.MODEL.NUM_JOINTS, 3),
|
||||
dtype=np.float32
|
||||
)
|
||||
all_boxes = np.zeros((num_samples, 6))
|
||||
image_path = []
|
||||
filenames = []
|
||||
imgnums = []
|
||||
idx = 0
|
||||
with torch.no_grad():
|
||||
end = time.time()
|
||||
for i, (input, target, target_weight, meta) in enumerate(val_loader):
|
||||
# compute output
|
||||
outputs = model(input)
|
||||
if isinstance(outputs, list):
|
||||
output = outputs[-1]
|
||||
else:
|
||||
output = outputs
|
||||
|
||||
if config.TEST.FLIP_TEST:
|
||||
# this part is ugly, because pytorch has not supported negative index
|
||||
# input_flipped = model(input[:, :, :, ::-1])
|
||||
input_flipped = np.flip(input.cpu().numpy(), 3).copy()
|
||||
input_flipped = torch.from_numpy(input_flipped).cuda()
|
||||
outputs_flipped = model(input_flipped)
|
||||
|
||||
if isinstance(outputs_flipped, list):
|
||||
output_flipped = outputs_flipped[-1]
|
||||
else:
|
||||
output_flipped = outputs_flipped
|
||||
|
||||
output_flipped = flip_back(output_flipped.cpu().numpy(),
|
||||
val_dataset.flip_pairs)
|
||||
output_flipped = torch.from_numpy(output_flipped.copy()).cuda()
|
||||
|
||||
|
||||
# feature is not aligned, shift flipped heatmap for higher accuracy
|
||||
if config.TEST.SHIFT_HEATMAP:
|
||||
output_flipped[:, :, :, 1:] = \
|
||||
output_flipped.clone()[:, :, :, 0:-1]
|
||||
|
||||
output = (output + output_flipped) * 0.5
|
||||
|
||||
target = target.cuda(non_blocking=True)
|
||||
target_weight = target_weight.cuda(non_blocking=True)
|
||||
|
||||
loss = criterion(output, target, target_weight)
|
||||
|
||||
num_images = input.size(0)
|
||||
# measure accuracy and record loss
|
||||
losses.update(loss.item(), num_images)
|
||||
_, avg_acc, cnt, pred = accuracy(output.cpu().numpy(),
|
||||
target.cpu().numpy())
|
||||
|
||||
acc.update(avg_acc, cnt)
|
||||
|
||||
# measure elapsed time
|
||||
batch_time.update(time.time() - end)
|
||||
end = time.time()
|
||||
|
||||
c = meta['center'].numpy()
|
||||
s = meta['scale'].numpy()
|
||||
score = meta['score'].numpy()
|
||||
|
||||
preds, maxvals = get_final_preds(
|
||||
config, output.clone().cpu().numpy(), c, s)
|
||||
|
||||
all_preds[idx:idx + num_images, :, 0:2] = preds[:, :, 0:2]
|
||||
all_preds[idx:idx + num_images, :, 2:3] = maxvals
|
||||
# double check this all_boxes parts
|
||||
all_boxes[idx:idx + num_images, 0:2] = c[:, 0:2]
|
||||
all_boxes[idx:idx + num_images, 2:4] = s[:, 0:2]
|
||||
all_boxes[idx:idx + num_images, 4] = np.prod(s*200, 1)
|
||||
all_boxes[idx:idx + num_images, 5] = score
|
||||
image_path.extend(meta['image'])
|
||||
|
||||
idx += num_images
|
||||
|
||||
if i % config.PRINT_FREQ == 0:
|
||||
msg = 'Test: [{0}/{1}]\t' \
|
||||
'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t' \
|
||||
'Loss {loss.val:.4f} ({loss.avg:.4f})\t' \
|
||||
'Accuracy {acc.val:.3f} ({acc.avg:.3f})'.format(
|
||||
i, len(val_loader), batch_time=batch_time,
|
||||
loss=losses, acc=acc)
|
||||
logger.info(msg)
|
||||
|
||||
prefix = '{}_{}'.format(
|
||||
os.path.join(output_dir, 'val'), i
|
||||
)
|
||||
save_debug_images(config, input, meta, target, pred*4, output,
|
||||
prefix)
|
||||
|
||||
name_values, perf_indicator = val_dataset.evaluate(
|
||||
config, all_preds, output_dir, all_boxes, image_path,
|
||||
filenames, imgnums
|
||||
)
|
||||
|
||||
model_name = config.MODEL.NAME
|
||||
if isinstance(name_values, list):
|
||||
for name_value in name_values:
|
||||
_print_name_value(name_value, model_name)
|
||||
else:
|
||||
_print_name_value(name_values, model_name)
|
||||
|
||||
if writer_dict:
|
||||
writer = writer_dict['writer']
|
||||
global_steps = writer_dict['valid_global_steps']
|
||||
writer.add_scalar(
|
||||
'valid_loss',
|
||||
losses.avg,
|
||||
global_steps
|
||||
)
|
||||
writer.add_scalar(
|
||||
'valid_acc',
|
||||
acc.avg,
|
||||
global_steps
|
||||
)
|
||||
if isinstance(name_values, list):
|
||||
for name_value in name_values:
|
||||
writer.add_scalars(
|
||||
'valid',
|
||||
dict(name_value),
|
||||
global_steps
|
||||
)
|
||||
else:
|
||||
writer.add_scalars(
|
||||
'valid',
|
||||
dict(name_values),
|
||||
global_steps
|
||||
)
|
||||
writer_dict['valid_global_steps'] = global_steps + 1
|
||||
|
||||
return perf_indicator
|
||||
|
||||
|
||||
# markdown format output
|
||||
def _print_name_value(name_value, full_arch_name):
|
||||
names = name_value.keys()
|
||||
values = name_value.values()
|
||||
num_values = len(name_value)
|
||||
logger.info(
|
||||
'| Arch ' +
|
||||
' '.join(['| {}'.format(name) for name in names]) +
|
||||
' |'
|
||||
)
|
||||
logger.info('|---' * (num_values+1) + '|')
|
||||
|
||||
if len(full_arch_name) > 15:
|
||||
full_arch_name = full_arch_name[:8] + '...'
|
||||
logger.info(
|
||||
'| ' + full_arch_name + ' ' +
|
||||
' '.join(['| {:.3f}'.format(value) for value in values]) +
|
||||
' |'
|
||||
)
|
||||
|
||||
|
||||
class AverageMeter(object):
|
||||
"""Computes and stores the average and current value"""
|
||||
def __init__(self):
|
||||
self.reset()
|
||||
|
||||
def reset(self):
|
||||
self.val = 0
|
||||
self.avg = 0
|
||||
self.sum = 0
|
||||
self.count = 0
|
||||
|
||||
def update(self, val, n=1):
|
||||
self.val = val
|
||||
self.sum += val * n
|
||||
self.count += n
|
||||
self.avg = self.sum / self.count if self.count != 0 else 0
|
||||
@@ -0,0 +1,79 @@
|
||||
# ------------------------------------------------------------------------------
|
||||
# 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 math
|
||||
|
||||
import numpy as np
|
||||
|
||||
from keypoints.lib.utils.transforms import transform_preds
|
||||
|
||||
|
||||
def get_max_preds(batch_heatmaps):
|
||||
'''
|
||||
get predictions from score maps
|
||||
heatmaps: numpy.ndarray([batch_size, num_joints, height, width])
|
||||
'''
|
||||
assert isinstance(batch_heatmaps, np.ndarray), \
|
||||
'batch_heatmaps should be numpy.ndarray'
|
||||
assert batch_heatmaps.ndim == 4, 'batch_images should be 4-ndim'
|
||||
|
||||
batch_size = batch_heatmaps.shape[0]
|
||||
num_joints = batch_heatmaps.shape[1]
|
||||
width = batch_heatmaps.shape[3]
|
||||
heatmaps_reshaped = batch_heatmaps.reshape((batch_size, num_joints, -1))
|
||||
idx = np.argmax(heatmaps_reshaped, 2)
|
||||
maxvals = np.amax(heatmaps_reshaped, 2)
|
||||
|
||||
maxvals = maxvals.reshape((batch_size, num_joints, 1))
|
||||
idx = idx.reshape((batch_size, num_joints, 1))
|
||||
|
||||
preds = np.tile(idx, (1, 1, 2)).astype(np.float32)
|
||||
|
||||
preds[:, :, 0] = (preds[:, :, 0]) % width
|
||||
preds[:, :, 1] = np.floor((preds[:, :, 1]) / width)
|
||||
|
||||
pred_mask = np.tile(np.greater(maxvals, 0.0), (1, 1, 2))
|
||||
pred_mask = pred_mask.astype(np.float32)
|
||||
|
||||
preds *= pred_mask
|
||||
return preds, maxvals
|
||||
|
||||
|
||||
def get_final_preds(config, batch_heatmaps, center, scale):
|
||||
coords, maxvals = get_max_preds(batch_heatmaps)
|
||||
|
||||
heatmap_height = batch_heatmaps.shape[2]
|
||||
heatmap_width = batch_heatmaps.shape[3]
|
||||
|
||||
# post-processing
|
||||
if config.TEST.POST_PROCESS:
|
||||
for n in range(coords.shape[0]):
|
||||
for p in range(coords.shape[1]):
|
||||
hm = batch_heatmaps[n][p]
|
||||
px = int(math.floor(coords[n][p][0] + 0.5))
|
||||
py = int(math.floor(coords[n][p][1] + 0.5))
|
||||
if 1 < px < heatmap_width-1 and 1 < py < heatmap_height-1:
|
||||
diff = np.array(
|
||||
[
|
||||
hm[py][px+1] - hm[py][px-1],
|
||||
hm[py+1][px]-hm[py-1][px]
|
||||
]
|
||||
)
|
||||
coords[n][p] += np.sign(diff) * .25
|
||||
|
||||
preds = coords.copy()
|
||||
|
||||
# Transform back
|
||||
for i in range(coords.shape[0]):
|
||||
preds[i] = transform_preds(
|
||||
coords[i], center[i], scale[i], [heatmap_width, heatmap_height]
|
||||
)
|
||||
|
||||
return preds, maxvals
|
||||
@@ -0,0 +1,84 @@
|
||||
# ------------------------------------------------------------------------------
|
||||
# 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 torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class JointsMSELoss(nn.Module):
|
||||
def __init__(self, use_target_weight):
|
||||
super(JointsMSELoss, self).__init__()
|
||||
self.criterion = nn.MSELoss(reduction='mean')
|
||||
self.use_target_weight = use_target_weight
|
||||
|
||||
def forward(self, output, target, target_weight):
|
||||
batch_size = output.size(0)
|
||||
num_joints = output.size(1)
|
||||
heatmaps_pred = output.reshape((batch_size, num_joints, -1)).split(1, 1)
|
||||
heatmaps_gt = target.reshape((batch_size, num_joints, -1)).split(1, 1)
|
||||
loss = 0
|
||||
|
||||
for idx in range(num_joints):
|
||||
heatmap_pred = heatmaps_pred[idx].squeeze()
|
||||
heatmap_gt = heatmaps_gt[idx].squeeze()
|
||||
if self.use_target_weight:
|
||||
loss += 0.5 * self.criterion(
|
||||
heatmap_pred.mul(target_weight[:, idx]),
|
||||
heatmap_gt.mul(target_weight[:, idx])
|
||||
)
|
||||
else:
|
||||
loss += 0.5 * self.criterion(heatmap_pred, heatmap_gt)
|
||||
|
||||
return loss / num_joints
|
||||
|
||||
|
||||
class JointsOHKMMSELoss(nn.Module):
|
||||
def __init__(self, use_target_weight, topk=8):
|
||||
super(JointsOHKMMSELoss, self).__init__()
|
||||
self.criterion = nn.MSELoss(reduction='none')
|
||||
self.use_target_weight = use_target_weight
|
||||
self.topk = topk
|
||||
|
||||
def ohkm(self, loss):
|
||||
ohkm_loss = 0.
|
||||
for i in range(loss.size()[0]):
|
||||
sub_loss = loss[i]
|
||||
topk_val, topk_idx = torch.topk(
|
||||
sub_loss, k=self.topk, dim=0, sorted=False
|
||||
)
|
||||
tmp_loss = torch.gather(sub_loss, 0, topk_idx)
|
||||
ohkm_loss += torch.sum(tmp_loss) / self.topk
|
||||
ohkm_loss /= loss.size()[0]
|
||||
return ohkm_loss
|
||||
|
||||
def forward(self, output, target, target_weight):
|
||||
batch_size = output.size(0)
|
||||
num_joints = output.size(1)
|
||||
heatmaps_pred = output.reshape((batch_size, num_joints, -1)).split(1, 1)
|
||||
heatmaps_gt = target.reshape((batch_size, num_joints, -1)).split(1, 1)
|
||||
|
||||
loss = []
|
||||
for idx in range(num_joints):
|
||||
heatmap_pred = heatmaps_pred[idx].squeeze()
|
||||
heatmap_gt = heatmaps_gt[idx].squeeze()
|
||||
if self.use_target_weight:
|
||||
loss.append(0.5 * self.criterion(
|
||||
heatmap_pred.mul(target_weight[:, idx]),
|
||||
heatmap_gt.mul(target_weight[:, idx])
|
||||
))
|
||||
else:
|
||||
loss.append(
|
||||
0.5 * self.criterion(heatmap_pred, heatmap_gt)
|
||||
)
|
||||
|
||||
loss = [l.mean(dim=1).unsqueeze(dim=1) for l in loss]
|
||||
loss = torch.cat(loss, dim=1)
|
||||
|
||||
return self.ohkm(loss)
|
||||
@@ -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']
|
||||
@@ -0,0 +1,13 @@
|
||||
# ------------------------------------------------------------------------------
|
||||
# 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 __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
@@ -0,0 +1,501 @@
|
||||
# ------------------------------------------------------------------------------
|
||||
# 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 os
|
||||
import logging
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
BN_MOMENTUM = 0.1
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def conv3x3(in_planes, out_planes, stride=1):
|
||||
"""3x3 convolution with padding"""
|
||||
return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
|
||||
padding=1, bias=False)
|
||||
|
||||
|
||||
class BasicBlock(nn.Module):
|
||||
expansion = 1
|
||||
|
||||
def __init__(self, inplanes, planes, stride=1, downsample=None):
|
||||
super(BasicBlock, self).__init__()
|
||||
self.conv1 = conv3x3(inplanes, planes, stride)
|
||||
self.bn1 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.conv2 = conv3x3(planes, planes)
|
||||
self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
|
||||
self.downsample = downsample
|
||||
self.stride = stride
|
||||
|
||||
def forward(self, x):
|
||||
residual = x
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.conv2(out)
|
||||
out = self.bn2(out)
|
||||
|
||||
if self.downsample is not None:
|
||||
residual = self.downsample(x)
|
||||
|
||||
out += residual
|
||||
out = self.relu(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class Bottleneck(nn.Module):
|
||||
expansion = 4
|
||||
|
||||
def __init__(self, inplanes, planes, stride=1, downsample=None):
|
||||
super(Bottleneck, self).__init__()
|
||||
self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
|
||||
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
|
||||
padding=1, bias=False)
|
||||
self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
|
||||
self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1,
|
||||
bias=False)
|
||||
self.bn3 = nn.BatchNorm2d(planes * self.expansion,
|
||||
momentum=BN_MOMENTUM)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.downsample = downsample
|
||||
self.stride = stride
|
||||
|
||||
def forward(self, x):
|
||||
residual = x
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.conv2(out)
|
||||
out = self.bn2(out)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.conv3(out)
|
||||
out = self.bn3(out)
|
||||
|
||||
if self.downsample is not None:
|
||||
residual = self.downsample(x)
|
||||
|
||||
out += residual
|
||||
out = self.relu(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class HighResolutionModule(nn.Module):
|
||||
def __init__(self, num_branches, blocks, num_blocks, num_inchannels,
|
||||
num_channels, fuse_method, multi_scale_output=True):
|
||||
super(HighResolutionModule, self).__init__()
|
||||
self._check_branches(
|
||||
num_branches, blocks, num_blocks, num_inchannels, num_channels)
|
||||
|
||||
self.num_inchannels = num_inchannels
|
||||
self.fuse_method = fuse_method
|
||||
self.num_branches = num_branches
|
||||
|
||||
self.multi_scale_output = multi_scale_output
|
||||
|
||||
self.branches = self._make_branches(
|
||||
num_branches, blocks, num_blocks, num_channels)
|
||||
self.fuse_layers = self._make_fuse_layers()
|
||||
self.relu = nn.ReLU(True)
|
||||
|
||||
def _check_branches(self, num_branches, blocks, num_blocks,
|
||||
num_inchannels, num_channels):
|
||||
if num_branches != len(num_blocks):
|
||||
error_msg = 'NUM_BRANCHES({}) <> NUM_BLOCKS({})'.format(
|
||||
num_branches, len(num_blocks))
|
||||
logger.error(error_msg)
|
||||
raise ValueError(error_msg)
|
||||
|
||||
if num_branches != len(num_channels):
|
||||
error_msg = 'NUM_BRANCHES({}) <> NUM_CHANNELS({})'.format(
|
||||
num_branches, len(num_channels))
|
||||
logger.error(error_msg)
|
||||
raise ValueError(error_msg)
|
||||
|
||||
if num_branches != len(num_inchannels):
|
||||
error_msg = 'NUM_BRANCHES({}) <> NUM_INCHANNELS({})'.format(
|
||||
num_branches, len(num_inchannels))
|
||||
logger.error(error_msg)
|
||||
raise ValueError(error_msg)
|
||||
|
||||
def _make_one_branch(self, branch_index, block, num_blocks, num_channels,
|
||||
stride=1):
|
||||
downsample = None
|
||||
if stride != 1 or \
|
||||
self.num_inchannels[branch_index] != num_channels[branch_index] * block.expansion:
|
||||
downsample = nn.Sequential(
|
||||
nn.Conv2d(
|
||||
self.num_inchannels[branch_index],
|
||||
num_channels[branch_index] * block.expansion,
|
||||
kernel_size=1, stride=stride, bias=False
|
||||
),
|
||||
nn.BatchNorm2d(
|
||||
num_channels[branch_index] * block.expansion,
|
||||
momentum=BN_MOMENTUM
|
||||
),
|
||||
)
|
||||
|
||||
layers = []
|
||||
layers.append(
|
||||
block(
|
||||
self.num_inchannels[branch_index],
|
||||
num_channels[branch_index],
|
||||
stride,
|
||||
downsample
|
||||
)
|
||||
)
|
||||
self.num_inchannels[branch_index] = \
|
||||
num_channels[branch_index] * block.expansion
|
||||
for i in range(1, num_blocks[branch_index]):
|
||||
layers.append(
|
||||
block(
|
||||
self.num_inchannels[branch_index],
|
||||
num_channels[branch_index]
|
||||
)
|
||||
)
|
||||
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def _make_branches(self, num_branches, block, num_blocks, num_channels):
|
||||
branches = []
|
||||
|
||||
for i in range(num_branches):
|
||||
branches.append(
|
||||
self._make_one_branch(i, block, num_blocks, num_channels)
|
||||
)
|
||||
|
||||
return nn.ModuleList(branches)
|
||||
|
||||
def _make_fuse_layers(self):
|
||||
if self.num_branches == 1:
|
||||
return None
|
||||
|
||||
num_branches = self.num_branches
|
||||
num_inchannels = self.num_inchannels
|
||||
fuse_layers = []
|
||||
for i in range(num_branches if self.multi_scale_output else 1):
|
||||
fuse_layer = []
|
||||
for j in range(num_branches):
|
||||
if j > i:
|
||||
fuse_layer.append(
|
||||
nn.Sequential(
|
||||
nn.Conv2d(
|
||||
num_inchannels[j],
|
||||
num_inchannels[i],
|
||||
1, 1, 0, bias=False
|
||||
),
|
||||
nn.BatchNorm2d(num_inchannels[i]),
|
||||
nn.Upsample(scale_factor=2**(j-i), mode='nearest')
|
||||
)
|
||||
)
|
||||
elif j == i:
|
||||
fuse_layer.append(None)
|
||||
else:
|
||||
conv3x3s = []
|
||||
for k in range(i-j):
|
||||
if k == i - j - 1:
|
||||
num_outchannels_conv3x3 = num_inchannels[i]
|
||||
conv3x3s.append(
|
||||
nn.Sequential(
|
||||
nn.Conv2d(
|
||||
num_inchannels[j],
|
||||
num_outchannels_conv3x3,
|
||||
3, 2, 1, bias=False
|
||||
),
|
||||
nn.BatchNorm2d(num_outchannels_conv3x3)
|
||||
)
|
||||
)
|
||||
else:
|
||||
num_outchannels_conv3x3 = num_inchannels[j]
|
||||
conv3x3s.append(
|
||||
nn.Sequential(
|
||||
nn.Conv2d(
|
||||
num_inchannels[j],
|
||||
num_outchannels_conv3x3,
|
||||
3, 2, 1, bias=False
|
||||
),
|
||||
nn.BatchNorm2d(num_outchannels_conv3x3),
|
||||
nn.ReLU(True)
|
||||
)
|
||||
)
|
||||
fuse_layer.append(nn.Sequential(*conv3x3s))
|
||||
fuse_layers.append(nn.ModuleList(fuse_layer))
|
||||
|
||||
return nn.ModuleList(fuse_layers)
|
||||
|
||||
def get_num_inchannels(self):
|
||||
return self.num_inchannels
|
||||
|
||||
def forward(self, x):
|
||||
if self.num_branches == 1:
|
||||
return [self.branches[0](x[0])]
|
||||
|
||||
for i in range(self.num_branches):
|
||||
x[i] = self.branches[i](x[i])
|
||||
|
||||
x_fuse = []
|
||||
|
||||
for i in range(len(self.fuse_layers)):
|
||||
y = x[0] if i == 0 else self.fuse_layers[i][0](x[0])
|
||||
for j in range(1, self.num_branches):
|
||||
if i == j:
|
||||
y = y + x[j]
|
||||
else:
|
||||
y = y + self.fuse_layers[i][j](x[j])
|
||||
x_fuse.append(self.relu(y))
|
||||
|
||||
return x_fuse
|
||||
|
||||
|
||||
blocks_dict = {
|
||||
'BASIC': BasicBlock,
|
||||
'BOTTLENECK': Bottleneck
|
||||
}
|
||||
|
||||
|
||||
class PoseHighResolutionNet(nn.Module):
|
||||
|
||||
def __init__(self, cfg, **kwargs):
|
||||
self.inplanes = 64
|
||||
extra = cfg.MODEL.EXTRA
|
||||
super(PoseHighResolutionNet, self).__init__()
|
||||
|
||||
# stem net
|
||||
self.conv1 = nn.Conv2d(3, 64, kernel_size=3, stride=2, padding=1,
|
||||
bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM)
|
||||
self.conv2 = nn.Conv2d(64, 64, kernel_size=3, stride=2, padding=1,
|
||||
bias=False)
|
||||
self.bn2 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.layer1 = self._make_layer(Bottleneck, 64, 4)
|
||||
|
||||
self.stage2_cfg = cfg['MODEL']['EXTRA']['STAGE2']
|
||||
num_channels = self.stage2_cfg['NUM_CHANNELS']
|
||||
block = blocks_dict[self.stage2_cfg['BLOCK']]
|
||||
num_channels = [
|
||||
num_channels[i] * block.expansion for i in range(len(num_channels))
|
||||
]
|
||||
self.transition1 = self._make_transition_layer([256], num_channels)
|
||||
self.stage2, pre_stage_channels = self._make_stage(
|
||||
self.stage2_cfg, num_channels)
|
||||
|
||||
self.stage3_cfg = cfg['MODEL']['EXTRA']['STAGE3']
|
||||
num_channels = self.stage3_cfg['NUM_CHANNELS']
|
||||
block = blocks_dict[self.stage3_cfg['BLOCK']]
|
||||
num_channels = [
|
||||
num_channels[i] * block.expansion for i in range(len(num_channels))
|
||||
]
|
||||
self.transition2 = self._make_transition_layer(
|
||||
pre_stage_channels, num_channels)
|
||||
self.stage3, pre_stage_channels = self._make_stage(
|
||||
self.stage3_cfg, num_channels)
|
||||
|
||||
self.stage4_cfg = cfg['MODEL']['EXTRA']['STAGE4']
|
||||
num_channels = self.stage4_cfg['NUM_CHANNELS']
|
||||
block = blocks_dict[self.stage4_cfg['BLOCK']]
|
||||
num_channels = [
|
||||
num_channels[i] * block.expansion for i in range(len(num_channels))
|
||||
]
|
||||
self.transition3 = self._make_transition_layer(
|
||||
pre_stage_channels, num_channels)
|
||||
self.stage4, pre_stage_channels = self._make_stage(
|
||||
self.stage4_cfg, num_channels, multi_scale_output=False)
|
||||
|
||||
self.final_layer = nn.Conv2d(
|
||||
in_channels=pre_stage_channels[0],
|
||||
out_channels=cfg.MODEL.NUM_JOINTS,
|
||||
kernel_size=extra.FINAL_CONV_KERNEL,
|
||||
stride=1,
|
||||
padding=1 if extra.FINAL_CONV_KERNEL == 3 else 0
|
||||
)
|
||||
|
||||
self.pretrained_layers = cfg['MODEL']['EXTRA']['PRETRAINED_LAYERS']
|
||||
|
||||
def _make_transition_layer(
|
||||
self, num_channels_pre_layer, num_channels_cur_layer):
|
||||
num_branches_cur = len(num_channels_cur_layer)
|
||||
num_branches_pre = len(num_channels_pre_layer)
|
||||
|
||||
transition_layers = []
|
||||
for i in range(num_branches_cur):
|
||||
if i < num_branches_pre:
|
||||
if num_channels_cur_layer[i] != num_channels_pre_layer[i]:
|
||||
transition_layers.append(
|
||||
nn.Sequential(
|
||||
nn.Conv2d(
|
||||
num_channels_pre_layer[i],
|
||||
num_channels_cur_layer[i],
|
||||
3, 1, 1, bias=False
|
||||
),
|
||||
nn.BatchNorm2d(num_channels_cur_layer[i]),
|
||||
nn.ReLU(inplace=True)
|
||||
)
|
||||
)
|
||||
else:
|
||||
transition_layers.append(None)
|
||||
else:
|
||||
conv3x3s = []
|
||||
for j in range(i+1-num_branches_pre):
|
||||
inchannels = num_channels_pre_layer[-1]
|
||||
outchannels = num_channels_cur_layer[i] \
|
||||
if j == i-num_branches_pre else inchannels
|
||||
conv3x3s.append(
|
||||
nn.Sequential(
|
||||
nn.Conv2d(
|
||||
inchannels, outchannels, 3, 2, 1, bias=False
|
||||
),
|
||||
nn.BatchNorm2d(outchannels),
|
||||
nn.ReLU(inplace=True)
|
||||
)
|
||||
)
|
||||
transition_layers.append(nn.Sequential(*conv3x3s))
|
||||
|
||||
return nn.ModuleList(transition_layers)
|
||||
|
||||
def _make_layer(self, block, planes, blocks, stride=1):
|
||||
downsample = None
|
||||
if stride != 1 or self.inplanes != planes * block.expansion:
|
||||
downsample = nn.Sequential(
|
||||
nn.Conv2d(
|
||||
self.inplanes, planes * block.expansion,
|
||||
kernel_size=1, stride=stride, bias=False
|
||||
),
|
||||
nn.BatchNorm2d(planes * block.expansion, momentum=BN_MOMENTUM),
|
||||
)
|
||||
|
||||
layers = []
|
||||
layers.append(block(self.inplanes, planes, stride, downsample))
|
||||
self.inplanes = planes * block.expansion
|
||||
for i in range(1, blocks):
|
||||
layers.append(block(self.inplanes, planes))
|
||||
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def _make_stage(self, layer_config, num_inchannels,
|
||||
multi_scale_output=True):
|
||||
num_modules = layer_config['NUM_MODULES']
|
||||
num_branches = layer_config['NUM_BRANCHES']
|
||||
num_blocks = layer_config['NUM_BLOCKS']
|
||||
num_channels = layer_config['NUM_CHANNELS']
|
||||
block = blocks_dict[layer_config['BLOCK']]
|
||||
fuse_method = layer_config['FUSE_METHOD']
|
||||
|
||||
modules = []
|
||||
for i in range(num_modules):
|
||||
# multi_scale_output is only used last module
|
||||
if not multi_scale_output and i == num_modules - 1:
|
||||
reset_multi_scale_output = False
|
||||
else:
|
||||
reset_multi_scale_output = True
|
||||
|
||||
modules.append(
|
||||
HighResolutionModule(
|
||||
num_branches,
|
||||
block,
|
||||
num_blocks,
|
||||
num_inchannels,
|
||||
num_channels,
|
||||
fuse_method,
|
||||
reset_multi_scale_output
|
||||
)
|
||||
)
|
||||
num_inchannels = modules[-1].get_num_inchannels()
|
||||
|
||||
return nn.Sequential(*modules), num_inchannels
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv1(x)
|
||||
x = self.bn1(x)
|
||||
x = self.relu(x)
|
||||
x = self.conv2(x)
|
||||
x = self.bn2(x)
|
||||
x = self.relu(x)
|
||||
x = self.layer1(x)
|
||||
|
||||
x_list = []
|
||||
for i in range(self.stage2_cfg['NUM_BRANCHES']):
|
||||
if self.transition1[i] is not None:
|
||||
x_list.append(self.transition1[i](x))
|
||||
else:
|
||||
x_list.append(x)
|
||||
y_list = self.stage2(x_list)
|
||||
|
||||
x_list = []
|
||||
for i in range(self.stage3_cfg['NUM_BRANCHES']):
|
||||
if self.transition2[i] is not None:
|
||||
x_list.append(self.transition2[i](y_list[-1]))
|
||||
else:
|
||||
x_list.append(y_list[i])
|
||||
y_list = self.stage3(x_list)
|
||||
|
||||
x_list = []
|
||||
for i in range(self.stage4_cfg['NUM_BRANCHES']):
|
||||
if self.transition3[i] is not None:
|
||||
x_list.append(self.transition3[i](y_list[-1]))
|
||||
else:
|
||||
x_list.append(y_list[i])
|
||||
y_list = self.stage4(x_list)
|
||||
|
||||
x = self.final_layer(y_list[0])
|
||||
|
||||
return x
|
||||
|
||||
def init_weights(self, pretrained=''):
|
||||
logger.info('=> init weights from normal distribution')
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
# nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
|
||||
nn.init.normal_(m.weight, std=0.001)
|
||||
for name, _ in m.named_parameters():
|
||||
if name in ['bias']:
|
||||
nn.init.constant_(m.bias, 0)
|
||||
elif isinstance(m, nn.BatchNorm2d):
|
||||
nn.init.constant_(m.weight, 1)
|
||||
nn.init.constant_(m.bias, 0)
|
||||
elif isinstance(m, nn.ConvTranspose2d):
|
||||
nn.init.normal_(m.weight, std=0.001)
|
||||
for name, _ in m.named_parameters():
|
||||
if name in ['bias']:
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
if os.path.isfile(pretrained):
|
||||
pretrained_state_dict = torch.load(pretrained)
|
||||
logger.info('=> loading pretrained model {}'.format(pretrained))
|
||||
|
||||
need_init_state_dict = {}
|
||||
for name, m in pretrained_state_dict.items():
|
||||
if name.split('.')[0] in self.pretrained_layers \
|
||||
or self.pretrained_layers[0] is '*':
|
||||
need_init_state_dict[name] = m
|
||||
self.load_state_dict(need_init_state_dict, strict=False)
|
||||
elif pretrained:
|
||||
logger.error('=> please download pre-trained models first!')
|
||||
raise ValueError('{} is not exist!'.format(pretrained))
|
||||
|
||||
|
||||
def get_pose_net(cfg, is_train, **kwargs):
|
||||
model = PoseHighResolutionNet(cfg, **kwargs)
|
||||
|
||||
if is_train and cfg.MODEL.INIT_WEIGHTS:
|
||||
model.init_weights(cfg.MODEL.PRETRAINED)
|
||||
|
||||
return model
|
||||
@@ -0,0 +1,271 @@
|
||||
# ------------------------------------------------------------------------------
|
||||
# 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 os
|
||||
import logging
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
BN_MOMENTUM = 0.1
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def conv3x3(in_planes, out_planes, stride=1):
|
||||
"""3x3 convolution with padding"""
|
||||
return nn.Conv2d(
|
||||
in_planes, out_planes, kernel_size=3, stride=stride,
|
||||
padding=1, bias=False
|
||||
)
|
||||
|
||||
|
||||
class BasicBlock(nn.Module):
|
||||
expansion = 1
|
||||
|
||||
def __init__(self, inplanes, planes, stride=1, downsample=None):
|
||||
super(BasicBlock, self).__init__()
|
||||
self.conv1 = conv3x3(inplanes, planes, stride)
|
||||
self.bn1 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.conv2 = conv3x3(planes, planes)
|
||||
self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
|
||||
self.downsample = downsample
|
||||
self.stride = stride
|
||||
|
||||
def forward(self, x):
|
||||
residual = x
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.conv2(out)
|
||||
out = self.bn2(out)
|
||||
|
||||
if self.downsample is not None:
|
||||
residual = self.downsample(x)
|
||||
|
||||
out += residual
|
||||
out = self.relu(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class Bottleneck(nn.Module):
|
||||
expansion = 4
|
||||
|
||||
def __init__(self, inplanes, planes, stride=1, downsample=None):
|
||||
super(Bottleneck, self).__init__()
|
||||
self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
|
||||
self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
|
||||
padding=1, bias=False)
|
||||
self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
|
||||
self.conv3 = nn.Conv2d(planes, planes * self.expansion, kernel_size=1,
|
||||
bias=False)
|
||||
self.bn3 = nn.BatchNorm2d(planes * self.expansion,
|
||||
momentum=BN_MOMENTUM)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.downsample = downsample
|
||||
self.stride = stride
|
||||
|
||||
def forward(self, x):
|
||||
residual = x
|
||||
|
||||
out = self.conv1(x)
|
||||
out = self.bn1(out)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.conv2(out)
|
||||
out = self.bn2(out)
|
||||
out = self.relu(out)
|
||||
|
||||
out = self.conv3(out)
|
||||
out = self.bn3(out)
|
||||
|
||||
if self.downsample is not None:
|
||||
residual = self.downsample(x)
|
||||
|
||||
out += residual
|
||||
out = self.relu(out)
|
||||
|
||||
return out
|
||||
|
||||
|
||||
class PoseResNet(nn.Module):
|
||||
|
||||
def __init__(self, block, layers, cfg, **kwargs):
|
||||
self.inplanes = 64
|
||||
extra = cfg.MODEL.EXTRA
|
||||
self.deconv_with_bias = extra.DECONV_WITH_BIAS
|
||||
|
||||
super(PoseResNet, self).__init__()
|
||||
self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
|
||||
bias=False)
|
||||
self.bn1 = nn.BatchNorm2d(64, momentum=BN_MOMENTUM)
|
||||
self.relu = nn.ReLU(inplace=True)
|
||||
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
|
||||
self.layer1 = self._make_layer(block, 64, layers[0])
|
||||
self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
|
||||
self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
|
||||
self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
|
||||
|
||||
# used for deconv layers
|
||||
self.deconv_layers = self._make_deconv_layer(
|
||||
extra.NUM_DECONV_LAYERS,
|
||||
extra.NUM_DECONV_FILTERS,
|
||||
extra.NUM_DECONV_KERNELS,
|
||||
)
|
||||
|
||||
self.final_layer = nn.Conv2d(
|
||||
in_channels=extra.NUM_DECONV_FILTERS[-1],
|
||||
out_channels=cfg.MODEL.NUM_JOINTS,
|
||||
kernel_size=extra.FINAL_CONV_KERNEL,
|
||||
stride=1,
|
||||
padding=1 if extra.FINAL_CONV_KERNEL == 3 else 0
|
||||
)
|
||||
|
||||
def _make_layer(self, block, planes, blocks, stride=1):
|
||||
downsample = None
|
||||
if stride != 1 or self.inplanes != planes * block.expansion:
|
||||
downsample = nn.Sequential(
|
||||
nn.Conv2d(self.inplanes, planes * block.expansion,
|
||||
kernel_size=1, stride=stride, bias=False),
|
||||
nn.BatchNorm2d(planes * block.expansion, momentum=BN_MOMENTUM),
|
||||
)
|
||||
|
||||
layers = []
|
||||
layers.append(block(self.inplanes, planes, stride, downsample))
|
||||
self.inplanes = planes * block.expansion
|
||||
for i in range(1, blocks):
|
||||
layers.append(block(self.inplanes, planes))
|
||||
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def _get_deconv_cfg(self, deconv_kernel, index):
|
||||
if deconv_kernel == 4:
|
||||
padding = 1
|
||||
output_padding = 0
|
||||
elif deconv_kernel == 3:
|
||||
padding = 1
|
||||
output_padding = 1
|
||||
elif deconv_kernel == 2:
|
||||
padding = 0
|
||||
output_padding = 0
|
||||
|
||||
return deconv_kernel, padding, output_padding
|
||||
|
||||
def _make_deconv_layer(self, num_layers, num_filters, num_kernels):
|
||||
assert num_layers == len(num_filters), \
|
||||
'ERROR: num_deconv_layers is different len(num_deconv_filters)'
|
||||
assert num_layers == len(num_kernels), \
|
||||
'ERROR: num_deconv_layers is different len(num_deconv_filters)'
|
||||
|
||||
layers = []
|
||||
for i in range(num_layers):
|
||||
kernel, padding, output_padding = \
|
||||
self._get_deconv_cfg(num_kernels[i], i)
|
||||
|
||||
planes = num_filters[i]
|
||||
layers.append(
|
||||
nn.ConvTranspose2d(
|
||||
in_channels=self.inplanes,
|
||||
out_channels=planes,
|
||||
kernel_size=kernel,
|
||||
stride=2,
|
||||
padding=padding,
|
||||
output_padding=output_padding,
|
||||
bias=self.deconv_with_bias))
|
||||
layers.append(nn.BatchNorm2d(planes, momentum=BN_MOMENTUM))
|
||||
layers.append(nn.ReLU(inplace=True))
|
||||
self.inplanes = planes
|
||||
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
def forward(self, x):
|
||||
x = self.conv1(x)
|
||||
x = self.bn1(x)
|
||||
x = self.relu(x)
|
||||
x = self.maxpool(x)
|
||||
|
||||
x = self.layer1(x)
|
||||
x = self.layer2(x)
|
||||
x = self.layer3(x)
|
||||
x = self.layer4(x)
|
||||
|
||||
x = self.deconv_layers(x)
|
||||
x = self.final_layer(x)
|
||||
|
||||
return x
|
||||
|
||||
def init_weights(self, pretrained=''):
|
||||
if os.path.isfile(pretrained):
|
||||
logger.info('=> init deconv weights from normal distribution')
|
||||
for name, m in self.deconv_layers.named_modules():
|
||||
if isinstance(m, nn.ConvTranspose2d):
|
||||
logger.info('=> init {}.weight as normal(0, 0.001)'.format(name))
|
||||
logger.info('=> init {}.bias as 0'.format(name))
|
||||
nn.init.normal_(m.weight, std=0.001)
|
||||
if self.deconv_with_bias:
|
||||
nn.init.constant_(m.bias, 0)
|
||||
elif isinstance(m, nn.BatchNorm2d):
|
||||
logger.info('=> init {}.weight as 1'.format(name))
|
||||
logger.info('=> init {}.bias as 0'.format(name))
|
||||
nn.init.constant_(m.weight, 1)
|
||||
nn.init.constant_(m.bias, 0)
|
||||
logger.info('=> init final conv weights from normal distribution')
|
||||
for m in self.final_layer.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
# nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
|
||||
logger.info('=> init {}.weight as normal(0, 0.001)'.format(name))
|
||||
logger.info('=> init {}.bias as 0'.format(name))
|
||||
nn.init.normal_(m.weight, std=0.001)
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
pretrained_state_dict = torch.load(pretrained)
|
||||
logger.info('=> loading pretrained model {}'.format(pretrained))
|
||||
self.load_state_dict(pretrained_state_dict, strict=False)
|
||||
else:
|
||||
logger.info('=> init weights from normal distribution')
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Conv2d):
|
||||
# nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
|
||||
nn.init.normal_(m.weight, std=0.001)
|
||||
# nn.init.constant_(m.bias, 0)
|
||||
elif isinstance(m, nn.BatchNorm2d):
|
||||
nn.init.constant_(m.weight, 1)
|
||||
nn.init.constant_(m.bias, 0)
|
||||
elif isinstance(m, nn.ConvTranspose2d):
|
||||
nn.init.normal_(m.weight, std=0.001)
|
||||
if self.deconv_with_bias:
|
||||
nn.init.constant_(m.bias, 0)
|
||||
|
||||
|
||||
resnet_spec = {
|
||||
18: (BasicBlock, [2, 2, 2, 2]),
|
||||
34: (BasicBlock, [3, 4, 6, 3]),
|
||||
50: (Bottleneck, [3, 4, 6, 3]),
|
||||
101: (Bottleneck, [3, 4, 23, 3]),
|
||||
152: (Bottleneck, [3, 8, 36, 3])
|
||||
}
|
||||
|
||||
|
||||
def get_pose_net(cfg, is_train, **kwargs):
|
||||
num_layers = cfg.MODEL.EXTRA.NUM_LAYERS
|
||||
|
||||
block_class, layers = resnet_spec[num_layers]
|
||||
|
||||
model = PoseResNet(block_class, layers, cfg, **kwargs)
|
||||
|
||||
if is_train and cfg.MODEL.INIT_WEIGHTS:
|
||||
model.init_weights(cfg.MODEL.PRETRAINED)
|
||||
|
||||
return model
|
||||
@@ -0,0 +1,164 @@
|
||||
# ------------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft
|
||||
# Licensed under the MIT License.
|
||||
# Modified from py-faster-rcnn (https://github.com/rbgirshick/py-faster-rcnn)
|
||||
# ------------------------------------------------------------------------------
|
||||
|
||||
from __future__ import absolute_import
|
||||
from __future__ import division
|
||||
from __future__ import print_function
|
||||
|
||||
import numpy as np
|
||||
|
||||
def py_nms_wrapper(thresh):
|
||||
def _nms(dets):
|
||||
return nms(dets, thresh)
|
||||
return _nms
|
||||
|
||||
|
||||
def nms(dets, thresh):
|
||||
"""
|
||||
greedily select boxes with high confidence and overlap with current maximum <= thresh
|
||||
rule out overlap >= thresh
|
||||
:param dets: [[x1, y1, x2, y2 score]]
|
||||
:param thresh: retain overlap < thresh
|
||||
:return: indexes to keep
|
||||
"""
|
||||
if dets.shape[0] == 0:
|
||||
return []
|
||||
|
||||
x1 = dets[:, 0]
|
||||
y1 = dets[:, 1]
|
||||
x2 = dets[:, 2]
|
||||
y2 = dets[:, 3]
|
||||
scores = dets[:, 4]
|
||||
|
||||
areas = (x2 - x1 + 1) * (y2 - y1 + 1)
|
||||
order = scores.argsort()[::-1]
|
||||
|
||||
keep = []
|
||||
while order.size > 0:
|
||||
i = order[0]
|
||||
keep.append(i)
|
||||
xx1 = np.maximum(x1[i], x1[order[1:]])
|
||||
yy1 = np.maximum(y1[i], y1[order[1:]])
|
||||
xx2 = np.minimum(x2[i], x2[order[1:]])
|
||||
yy2 = np.minimum(y2[i], y2[order[1:]])
|
||||
|
||||
w = np.maximum(0.0, xx2 - xx1 + 1)
|
||||
h = np.maximum(0.0, yy2 - yy1 + 1)
|
||||
inter = w * h
|
||||
ovr = inter / (areas[i] + areas[order[1:]] - inter)
|
||||
|
||||
inds = np.where(ovr <= thresh)[0]
|
||||
order = order[inds + 1]
|
||||
|
||||
return keep
|
||||
|
||||
|
||||
def oks_iou(g, d, a_g, a_d, sigmas=None, in_vis_thre=None):
|
||||
if not isinstance(sigmas, np.ndarray):
|
||||
sigmas = np.array([.26, .25, .25, .35, .35, .79, .79, .72, .72, .62, .62, 1.07, 1.07, .87, .87, .89, .89]) / 10.0
|
||||
vars = (sigmas * 2) ** 2
|
||||
xg = g[0::3]
|
||||
yg = g[1::3]
|
||||
vg = g[2::3]
|
||||
ious = np.zeros((d.shape[0]))
|
||||
for n_d in range(0, d.shape[0]):
|
||||
xd = d[n_d, 0::3]
|
||||
yd = d[n_d, 1::3]
|
||||
vd = d[n_d, 2::3]
|
||||
dx = xd - xg
|
||||
dy = yd - yg
|
||||
e = (dx ** 2 + dy ** 2) / vars / ((a_g + a_d[n_d]) / 2 + np.spacing(1)) / 2
|
||||
if in_vis_thre is not None:
|
||||
ind = list(vg > in_vis_thre) and list(vd > in_vis_thre)
|
||||
e = e[ind]
|
||||
ious[n_d] = np.sum(np.exp(-e)) / e.shape[0] if e.shape[0] != 0 else 0.0
|
||||
return ious
|
||||
|
||||
|
||||
def oks_nms(kpts_db, thresh, sigmas=None, in_vis_thre=None):
|
||||
"""
|
||||
greedily select boxes with high confidence and overlap with current maximum <= thresh
|
||||
rule out overlap >= thresh, overlap = oks
|
||||
:param kpts_db
|
||||
:param thresh: retain overlap < thresh
|
||||
:return: indexes to keep
|
||||
"""
|
||||
if len(kpts_db) == 0:
|
||||
return []
|
||||
|
||||
scores = np.array([kpts_db[i]['score'] for i in range(len(kpts_db))])
|
||||
kpts = np.array([kpts_db[i]['keypoints'].flatten() for i in range(len(kpts_db))])
|
||||
areas = np.array([kpts_db[i]['area'] for i in range(len(kpts_db))])
|
||||
|
||||
order = scores.argsort()[::-1]
|
||||
|
||||
keep = []
|
||||
while order.size > 0:
|
||||
i = order[0]
|
||||
keep.append(i)
|
||||
|
||||
oks_ovr = oks_iou(kpts[i], kpts[order[1:]], areas[i], areas[order[1:]], sigmas, in_vis_thre)
|
||||
|
||||
inds = np.where(oks_ovr <= thresh)[0]
|
||||
order = order[inds + 1]
|
||||
|
||||
return keep
|
||||
|
||||
|
||||
def rescore(overlap, scores, thresh, type='gaussian'):
|
||||
assert overlap.shape[0] == scores.shape[0]
|
||||
if type == 'linear':
|
||||
inds = np.where(overlap >= thresh)[0]
|
||||
scores[inds] = scores[inds] * (1 - overlap[inds])
|
||||
else:
|
||||
scores = scores * np.exp(- overlap**2 / thresh)
|
||||
|
||||
return scores
|
||||
|
||||
|
||||
def soft_oks_nms(kpts_db, thresh, sigmas=None, in_vis_thre=None):
|
||||
"""
|
||||
greedily select boxes with high confidence and overlap with current maximum <= thresh
|
||||
rule out overlap >= thresh, overlap = oks
|
||||
:param kpts_db
|
||||
:param thresh: retain overlap < thresh
|
||||
:return: indexes to keep
|
||||
"""
|
||||
if len(kpts_db) == 0:
|
||||
return []
|
||||
|
||||
scores = np.array([kpts_db[i]['score'] for i in range(len(kpts_db))])
|
||||
kpts = np.array([kpts_db[i]['keypoints'].flatten() for i in range(len(kpts_db))])
|
||||
areas = np.array([kpts_db[i]['area'] for i in range(len(kpts_db))])
|
||||
|
||||
order = scores.argsort()[::-1]
|
||||
scores = scores[order]
|
||||
|
||||
# max_dets = order.size
|
||||
max_dets = 20
|
||||
keep = np.zeros(max_dets, dtype=np.intp)
|
||||
keep_cnt = 0
|
||||
while order.size > 0 and keep_cnt < max_dets:
|
||||
i = order[0]
|
||||
|
||||
oks_ovr = oks_iou(kpts[i], kpts[order[1:]], areas[i], areas[order[1:]], sigmas, in_vis_thre)
|
||||
|
||||
order = order[1:]
|
||||
scores = rescore(oks_ovr, scores[1:], thresh)
|
||||
|
||||
tmp = scores.argsort()[::-1]
|
||||
order = order[tmp]
|
||||
scores = scores[tmp]
|
||||
|
||||
keep[keep_cnt] = i
|
||||
keep_cnt += 1
|
||||
|
||||
keep = keep[:keep_cnt]
|
||||
|
||||
return keep
|
||||
# kpts_db = kpts_db[:keep_cnt]
|
||||
|
||||
# return kpts_db
|
||||
@@ -0,0 +1,121 @@
|
||||
# ------------------------------------------------------------------------------
|
||||
# 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 numpy as np
|
||||
import cv2
|
||||
|
||||
|
||||
def flip_back(output_flipped, matched_parts):
|
||||
'''
|
||||
ouput_flipped: numpy.ndarray(batch_size, num_joints, height, width)
|
||||
'''
|
||||
assert output_flipped.ndim == 4,\
|
||||
'output_flipped should be [batch_size, num_joints, height, width]'
|
||||
|
||||
output_flipped = output_flipped[:, :, :, ::-1]
|
||||
|
||||
for pair in matched_parts:
|
||||
tmp = output_flipped[:, pair[0], :, :].copy()
|
||||
output_flipped[:, pair[0], :, :] = output_flipped[:, pair[1], :, :]
|
||||
output_flipped[:, pair[1], :, :] = tmp
|
||||
|
||||
return output_flipped
|
||||
|
||||
|
||||
def fliplr_joints(joints, joints_vis, width, matched_parts):
|
||||
"""
|
||||
flip coords
|
||||
"""
|
||||
# Flip horizontal
|
||||
joints[:, 0] = width - joints[:, 0] - 1
|
||||
|
||||
# Change left-right parts
|
||||
for pair in matched_parts:
|
||||
joints[pair[0], :], joints[pair[1], :] = \
|
||||
joints[pair[1], :], joints[pair[0], :].copy()
|
||||
joints_vis[pair[0], :], joints_vis[pair[1], :] = \
|
||||
joints_vis[pair[1], :], joints_vis[pair[0], :].copy()
|
||||
|
||||
return joints*joints_vis, joints_vis
|
||||
|
||||
|
||||
def transform_preds(coords, center, scale, output_size):
|
||||
target_coords = np.zeros(coords.shape)
|
||||
trans = get_affine_transform(center, scale, 0, output_size, inv=1)
|
||||
for p in range(coords.shape[0]):
|
||||
target_coords[p, 0:2] = affine_transform(coords[p, 0:2], trans)
|
||||
return target_coords
|
||||
|
||||
|
||||
def get_affine_transform(
|
||||
center, scale, rot, output_size,
|
||||
shift=np.array([0, 0], dtype=np.float32), inv=0
|
||||
):
|
||||
if not isinstance(scale, np.ndarray) and not isinstance(scale, list):
|
||||
print(scale)
|
||||
scale = np.array([scale, scale])
|
||||
|
||||
scale_tmp = scale * 200.0
|
||||
src_w = scale_tmp[0]
|
||||
dst_w = output_size[0]
|
||||
dst_h = output_size[1]
|
||||
|
||||
rot_rad = np.pi * rot / 180
|
||||
src_dir = get_dir([0, src_w * -0.5], rot_rad)
|
||||
dst_dir = np.array([0, dst_w * -0.5], np.float32)
|
||||
|
||||
src = np.zeros((3, 2), dtype=np.float32)
|
||||
dst = np.zeros((3, 2), dtype=np.float32)
|
||||
src[0, :] = center + scale_tmp * shift
|
||||
src[1, :] = center + src_dir + scale_tmp * shift
|
||||
dst[0, :] = [dst_w * 0.5, dst_h * 0.5]
|
||||
dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir
|
||||
|
||||
src[2:, :] = get_3rd_point(src[0, :], src[1, :])
|
||||
dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :])
|
||||
|
||||
if inv:
|
||||
trans = cv2.getAffineTransform(np.float32(dst), np.float32(src))
|
||||
else:
|
||||
trans = cv2.getAffineTransform(np.float32(src), np.float32(dst))
|
||||
|
||||
return trans
|
||||
|
||||
|
||||
def affine_transform(pt, t):
|
||||
new_pt = np.array([pt[0], pt[1], 1.]).T
|
||||
new_pt = np.dot(t, new_pt)
|
||||
return new_pt[:2]
|
||||
|
||||
|
||||
def get_3rd_point(a, b):
|
||||
direct = a - b
|
||||
return b + np.array([-direct[1], direct[0]], dtype=np.float32)
|
||||
|
||||
|
||||
def get_dir(src_point, rot_rad):
|
||||
sn, cs = np.sin(rot_rad), np.cos(rot_rad)
|
||||
|
||||
src_result = [0, 0]
|
||||
src_result[0] = src_point[0] * cs - src_point[1] * sn
|
||||
src_result[1] = src_point[0] * sn + src_point[1] * cs
|
||||
|
||||
return src_result
|
||||
|
||||
|
||||
def crop(img, center, scale, output_size, rot=0):
|
||||
trans = get_affine_transform(center, scale, rot, output_size)
|
||||
|
||||
dst_img = cv2.warpAffine(
|
||||
img, trans, (int(output_size[0]), int(output_size[1])),
|
||||
flags=cv2.INTER_LINEAR
|
||||
)
|
||||
|
||||
return dst_img
|
||||
@@ -0,0 +1,203 @@
|
||||
# ------------------------------------------------------------------------------
|
||||
# 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 os
|
||||
import logging
|
||||
import time
|
||||
from collections import namedtuple
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
import torch.optim as optim
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
def create_logger(cfg, cfg_name, phase='train'):
|
||||
root_output_dir = Path(cfg.OUTPUT_DIR)
|
||||
# set up logger
|
||||
if not root_output_dir.exists():
|
||||
print('=> creating {}'.format(root_output_dir))
|
||||
root_output_dir.mkdir()
|
||||
|
||||
dataset = cfg.DATASET.DATASET + '_' + cfg.DATASET.HYBRID_JOINTS_TYPE \
|
||||
if cfg.DATASET.HYBRID_JOINTS_TYPE else cfg.DATASET.DATASET
|
||||
dataset = dataset.replace(':', '_')
|
||||
model = cfg.MODEL.NAME
|
||||
cfg_name = os.path.basename(cfg_name).split('.')[0]
|
||||
|
||||
final_output_dir = root_output_dir / dataset / model / cfg_name
|
||||
|
||||
print('=> creating {}'.format(final_output_dir))
|
||||
final_output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
time_str = time.strftime('%Y-%m-%d-%H-%M')
|
||||
log_file = '{}_{}_{}.log'.format(cfg_name, time_str, phase)
|
||||
final_log_file = final_output_dir / log_file
|
||||
head = '%(asctime)-15s %(message)s'
|
||||
logging.basicConfig(filename=str(final_log_file),
|
||||
format=head)
|
||||
logger = logging.getLogger()
|
||||
logger.setLevel(logging.INFO)
|
||||
console = logging.StreamHandler()
|
||||
logging.getLogger('').addHandler(console)
|
||||
|
||||
tensorboard_log_dir = Path(cfg.LOG_DIR) / dataset / model / \
|
||||
(cfg_name + '_' + time_str)
|
||||
|
||||
print('=> creating {}'.format(tensorboard_log_dir))
|
||||
tensorboard_log_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
return logger, str(final_output_dir), str(tensorboard_log_dir)
|
||||
|
||||
|
||||
def get_optimizer(cfg, model):
|
||||
optimizer = None
|
||||
if cfg.TRAIN.OPTIMIZER == 'sgd':
|
||||
optimizer = optim.SGD(
|
||||
model.parameters(),
|
||||
lr=cfg.TRAIN.LR,
|
||||
momentum=cfg.TRAIN.MOMENTUM,
|
||||
weight_decay=cfg.TRAIN.WD,
|
||||
nesterov=cfg.TRAIN.NESTEROV
|
||||
)
|
||||
elif cfg.TRAIN.OPTIMIZER == 'adam':
|
||||
optimizer = optim.Adam(
|
||||
model.parameters(),
|
||||
lr=cfg.TRAIN.LR
|
||||
)
|
||||
|
||||
return optimizer
|
||||
|
||||
|
||||
def save_checkpoint(states, is_best, output_dir,
|
||||
filename='checkpoint.pth'):
|
||||
torch.save(states, os.path.join(output_dir, filename))
|
||||
if is_best and 'state_dict' in states:
|
||||
torch.save(states['best_state_dict'],
|
||||
os.path.join(output_dir, 'model_best.pth'))
|
||||
|
||||
|
||||
def get_model_summary(model, *input_tensors, item_length=26, verbose=False):
|
||||
"""
|
||||
:param model:
|
||||
:param input_tensors:
|
||||
:param item_length:
|
||||
:return:
|
||||
"""
|
||||
|
||||
summary = []
|
||||
|
||||
ModuleDetails = namedtuple(
|
||||
"Layer", ["name", "input_size", "output_size", "num_parameters", "multiply_adds"])
|
||||
hooks = []
|
||||
layer_instances = {}
|
||||
|
||||
def add_hooks(module):
|
||||
|
||||
def hook(module, input, output):
|
||||
class_name = str(module.__class__.__name__)
|
||||
|
||||
instance_index = 1
|
||||
if class_name not in layer_instances:
|
||||
layer_instances[class_name] = instance_index
|
||||
else:
|
||||
instance_index = layer_instances[class_name] + 1
|
||||
layer_instances[class_name] = instance_index
|
||||
|
||||
layer_name = class_name + "_" + str(instance_index)
|
||||
|
||||
params = 0
|
||||
|
||||
if class_name.find("Conv") != -1 or class_name.find("BatchNorm") != -1 or \
|
||||
class_name.find("Linear") != -1:
|
||||
for param_ in module.parameters():
|
||||
params += param_.view(-1).size(0)
|
||||
|
||||
flops = "Not Available"
|
||||
if class_name.find("Conv") != -1 and hasattr(module, "weight"):
|
||||
flops = (
|
||||
torch.prod(
|
||||
torch.LongTensor(list(module.weight.data.size()))) *
|
||||
torch.prod(
|
||||
torch.LongTensor(list(output.size())[2:]))).item()
|
||||
elif isinstance(module, nn.Linear):
|
||||
flops = (torch.prod(torch.LongTensor(list(output.size()))) \
|
||||
* input[0].size(1)).item()
|
||||
|
||||
if isinstance(input[0], list):
|
||||
input = input[0]
|
||||
if isinstance(output, list):
|
||||
output = output[0]
|
||||
|
||||
summary.append(
|
||||
ModuleDetails(
|
||||
name=layer_name,
|
||||
input_size=list(input[0].size()),
|
||||
output_size=list(output.size()),
|
||||
num_parameters=params,
|
||||
multiply_adds=flops)
|
||||
)
|
||||
|
||||
if not isinstance(module, nn.ModuleList) \
|
||||
and not isinstance(module, nn.Sequential) \
|
||||
and module != model:
|
||||
hooks.append(module.register_forward_hook(hook))
|
||||
|
||||
model.eval()
|
||||
model.apply(add_hooks)
|
||||
|
||||
space_len = item_length
|
||||
|
||||
model(*input_tensors)
|
||||
for hook in hooks:
|
||||
hook.remove()
|
||||
|
||||
details = ''
|
||||
if verbose:
|
||||
details = "Model Summary" + \
|
||||
os.linesep + \
|
||||
"Name{}Input Size{}Output Size{}Parameters{}Multiply Adds (Flops){}".format(
|
||||
' ' * (space_len - len("Name")),
|
||||
' ' * (space_len - len("Input Size")),
|
||||
' ' * (space_len - len("Output Size")),
|
||||
' ' * (space_len - len("Parameters")),
|
||||
' ' * (space_len - len("Multiply Adds (Flops)"))) \
|
||||
+ os.linesep + '-' * space_len * 5 + os.linesep
|
||||
|
||||
params_sum = 0
|
||||
flops_sum = 0
|
||||
for layer in summary:
|
||||
params_sum += layer.num_parameters
|
||||
if layer.multiply_adds != "Not Available":
|
||||
flops_sum += layer.multiply_adds
|
||||
if verbose:
|
||||
details += "{}{}{}{}{}{}{}{}{}{}".format(
|
||||
layer.name,
|
||||
' ' * (space_len - len(layer.name)),
|
||||
layer.input_size,
|
||||
' ' * (space_len - len(str(layer.input_size))),
|
||||
layer.output_size,
|
||||
' ' * (space_len - len(str(layer.output_size))),
|
||||
layer.num_parameters,
|
||||
' ' * (space_len - len(str(layer.num_parameters))),
|
||||
layer.multiply_adds,
|
||||
' ' * (space_len - len(str(layer.multiply_adds)))) \
|
||||
+ os.linesep + '-' * space_len * 5 + os.linesep
|
||||
|
||||
details += os.linesep \
|
||||
+ "Total Parameters: {:,}".format(params_sum) \
|
||||
+ os.linesep + '-' * space_len * 5 + os.linesep
|
||||
details += "Total Multiply Adds (For Convolution and Linear Layers only): {:,} GFLOPs".format(flops_sum/(1024**3)) \
|
||||
+ os.linesep + '-' * space_len * 5 + os.linesep
|
||||
details += "Number of Layers" + os.linesep
|
||||
for layer in layer_instances:
|
||||
details += "{} : {} layers ".format(layer, layer_instances[layer])
|
||||
|
||||
return details
|
||||
@@ -0,0 +1,141 @@
|
||||
# ------------------------------------------------------------------------------
|
||||
# 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 math
|
||||
|
||||
import numpy as np
|
||||
import torchvision
|
||||
import cv2
|
||||
|
||||
from core.inference import get_max_preds
|
||||
|
||||
|
||||
def save_batch_image_with_joints(batch_image, batch_joints, batch_joints_vis,
|
||||
file_name, nrow=8, padding=2):
|
||||
'''
|
||||
batch_image: [batch_size, channel, height, width]
|
||||
batch_joints: [batch_size, num_joints, 3],
|
||||
batch_joints_vis: [batch_size, num_joints, 1],
|
||||
}
|
||||
'''
|
||||
grid = torchvision.utils.make_grid(batch_image, nrow, padding, True)
|
||||
ndarr = grid.mul(255).clamp(0, 255).byte().permute(1, 2, 0).cpu().numpy()
|
||||
ndarr = ndarr.copy()
|
||||
|
||||
nmaps = batch_image.size(0)
|
||||
xmaps = min(nrow, nmaps)
|
||||
ymaps = int(math.ceil(float(nmaps) / xmaps))
|
||||
height = int(batch_image.size(2) + padding)
|
||||
width = int(batch_image.size(3) + padding)
|
||||
k = 0
|
||||
for y in range(ymaps):
|
||||
for x in range(xmaps):
|
||||
if k >= nmaps:
|
||||
break
|
||||
joints = batch_joints[k]
|
||||
joints_vis = batch_joints_vis[k]
|
||||
|
||||
for joint, joint_vis in zip(joints, joints_vis):
|
||||
joint[0] = x * width + padding + joint[0]
|
||||
joint[1] = y * height + padding + joint[1]
|
||||
if joint_vis[0]:
|
||||
cv2.circle(ndarr, (int(joint[0]), int(joint[1])), 2, [255, 0, 0], 2)
|
||||
k = k + 1
|
||||
cv2.imwrite(file_name, ndarr)
|
||||
|
||||
|
||||
def save_batch_heatmaps(batch_image, batch_heatmaps, file_name,
|
||||
normalize=True):
|
||||
'''
|
||||
batch_image: [batch_size, channel, height, width]
|
||||
batch_heatmaps: ['batch_size, num_joints, height, width]
|
||||
file_name: saved file name
|
||||
'''
|
||||
if normalize:
|
||||
batch_image = batch_image.clone()
|
||||
min = float(batch_image.min())
|
||||
max = float(batch_image.max())
|
||||
|
||||
batch_image.add_(-min).div_(max - min + 1e-5)
|
||||
|
||||
batch_size = batch_heatmaps.size(0)
|
||||
num_joints = batch_heatmaps.size(1)
|
||||
heatmap_height = batch_heatmaps.size(2)
|
||||
heatmap_width = batch_heatmaps.size(3)
|
||||
|
||||
grid_image = np.zeros((batch_size*heatmap_height,
|
||||
(num_joints+1)*heatmap_width,
|
||||
3),
|
||||
dtype=np.uint8)
|
||||
|
||||
preds, maxvals = get_max_preds(batch_heatmaps.detach().cpu().numpy())
|
||||
|
||||
for i in range(batch_size):
|
||||
image = batch_image[i].mul(255)\
|
||||
.clamp(0, 255)\
|
||||
.byte()\
|
||||
.permute(1, 2, 0)\
|
||||
.cpu().numpy()
|
||||
heatmaps = batch_heatmaps[i].mul(255)\
|
||||
.clamp(0, 255)\
|
||||
.byte()\
|
||||
.cpu().numpy()
|
||||
|
||||
resized_image = cv2.resize(image,
|
||||
(int(heatmap_width), int(heatmap_height)))
|
||||
|
||||
height_begin = heatmap_height * i
|
||||
height_end = heatmap_height * (i + 1)
|
||||
for j in range(num_joints):
|
||||
cv2.circle(resized_image,
|
||||
(int(preds[i][j][0]), int(preds[i][j][1])),
|
||||
1, [0, 0, 255], 1)
|
||||
heatmap = heatmaps[j, :, :]
|
||||
colored_heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
|
||||
masked_image = colored_heatmap*0.7 + resized_image*0.3
|
||||
cv2.circle(masked_image,
|
||||
(int(preds[i][j][0]), int(preds[i][j][1])),
|
||||
1, [0, 0, 255], 1)
|
||||
|
||||
width_begin = heatmap_width * (j+1)
|
||||
width_end = heatmap_width * (j+2)
|
||||
grid_image[height_begin:height_end, width_begin:width_end, :] = \
|
||||
masked_image
|
||||
# grid_image[height_begin:height_end, width_begin:width_end, :] = \
|
||||
# colored_heatmap*0.7 + resized_image*0.3
|
||||
|
||||
grid_image[height_begin:height_end, 0:heatmap_width, :] = resized_image
|
||||
|
||||
cv2.imwrite(file_name, grid_image)
|
||||
|
||||
|
||||
def save_debug_images(config, input, meta, target, joints_pred, output,
|
||||
prefix):
|
||||
if not config.DEBUG.DEBUG:
|
||||
return
|
||||
|
||||
if config.DEBUG.SAVE_BATCH_IMAGES_GT:
|
||||
save_batch_image_with_joints(
|
||||
input, meta['joints'], meta['joints_vis'],
|
||||
'{}_gt.jpg'.format(prefix)
|
||||
)
|
||||
if config.DEBUG.SAVE_BATCH_IMAGES_PRED:
|
||||
save_batch_image_with_joints(
|
||||
input, joints_pred, meta['joints_vis'],
|
||||
'{}_pred.jpg'.format(prefix)
|
||||
)
|
||||
if config.DEBUG.SAVE_HEATMAPS_GT:
|
||||
save_batch_heatmaps(
|
||||
input, target, '{}_hm_gt.jpg'.format(prefix)
|
||||
)
|
||||
if config.DEBUG.SAVE_HEATMAPS_PRED:
|
||||
save_batch_heatmaps(
|
||||
input, output, '{}_hm_pred.jpg'.format(prefix)
|
||||
)
|
||||
@@ -0,0 +1,70 @@
|
||||
# ------------------------------------------------------------------------------
|
||||
# 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 os
|
||||
import zipfile
|
||||
import xml.etree.ElementTree as ET
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
_im_zfile = []
|
||||
_xml_path_zip = []
|
||||
_xml_zfile = []
|
||||
|
||||
|
||||
def imread(filename, flags=cv2.IMREAD_COLOR):
|
||||
global _im_zfile
|
||||
path = filename
|
||||
pos_at = path.index('@')
|
||||
if pos_at == -1:
|
||||
print("character '@' is not found from the given path '%s'"%(path))
|
||||
assert 0
|
||||
path_zip = path[0: pos_at]
|
||||
path_img = path[pos_at + 2:]
|
||||
if not os.path.isfile(path_zip):
|
||||
print("zip file '%s' is not found"%(path_zip))
|
||||
assert 0
|
||||
for i in range(len(_im_zfile)):
|
||||
if _im_zfile[i]['path'] == path_zip:
|
||||
data = _im_zfile[i]['zipfile'].read(path_img)
|
||||
return cv2.imdecode(np.frombuffer(data, np.uint8), flags)
|
||||
|
||||
_im_zfile.append({
|
||||
'path': path_zip,
|
||||
'zipfile': zipfile.ZipFile(path_zip, 'r')
|
||||
})
|
||||
data = _im_zfile[-1]['zipfile'].read(path_img)
|
||||
|
||||
return cv2.imdecode(np.frombuffer(data, np.uint8), flags)
|
||||
|
||||
|
||||
def xmlread(filename):
|
||||
global _xml_path_zip
|
||||
global _xml_zfile
|
||||
path = filename
|
||||
pos_at = path.index('@')
|
||||
if pos_at == -1:
|
||||
print("character '@' is not found from the given path '%s'"%(path))
|
||||
assert 0
|
||||
path_zip = path[0: pos_at]
|
||||
path_xml = path[pos_at + 2:]
|
||||
if not os.path.isfile(path_zip):
|
||||
print("zip file '%s' is not found"%(path_zip))
|
||||
assert 0
|
||||
for i in xrange(len(_xml_path_zip)):
|
||||
if _xml_path_zip[i] == path_zip:
|
||||
data = _xml_zfile[i].open(path_xml)
|
||||
return ET.fromstring(data.read())
|
||||
_xml_path_zip.append(path_zip)
|
||||
print("read new xml file '%s'"%(path_zip))
|
||||
_xml_zfile.append(zipfile.ZipFile(path_zip, 'r'))
|
||||
data = _xml_zfile[-1].open(path_xml)
|
||||
return ET.fromstring(data.read())
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,3 @@
|
||||
from .wider_face import WiderFaceDetection, detection_collate
|
||||
from .data_augment import *
|
||||
from .config import *
|
||||
@@ -0,0 +1,42 @@
|
||||
# config.py
|
||||
|
||||
cfg_mnet = {
|
||||
'name': 'mobilenet0.25',
|
||||
'min_sizes': [[16, 32], [64, 128], [256, 512]],
|
||||
'steps': [8, 16, 32],
|
||||
'variance': [0.1, 0.2],
|
||||
'clip': False,
|
||||
'loc_weight': 2.0,
|
||||
'gpu_train': True,
|
||||
'batch_size': 32,
|
||||
'ngpu': 1,
|
||||
'epoch': 250,
|
||||
'decay1': 190,
|
||||
'decay2': 220,
|
||||
'image_size': 640,
|
||||
'pretrain': True,
|
||||
'return_layers': {'stage1': 1, 'stage2': 2, 'stage3': 3},
|
||||
'in_channel': 32,
|
||||
'out_channel': 64
|
||||
}
|
||||
|
||||
cfg_re50 = {
|
||||
'name': 'Resnet50',
|
||||
'min_sizes': [[16, 32], [64, 128], [256, 512]],
|
||||
'steps': [8, 16, 32],
|
||||
'variance': [0.1, 0.2],
|
||||
'clip': False,
|
||||
'loc_weight': 2.0,
|
||||
'gpu_train': True,
|
||||
'batch_size': 24,
|
||||
'ngpu': 4,
|
||||
'epoch': 100,
|
||||
'decay1': 70,
|
||||
'decay2': 90,
|
||||
'image_size': 840,
|
||||
'pretrain': True,
|
||||
'return_layers': {'layer2': 1, 'layer3': 2, 'layer4': 3},
|
||||
'in_channel': 256,
|
||||
'out_channel': 256
|
||||
}
|
||||
|
||||
@@ -0,0 +1,237 @@
|
||||
import cv2
|
||||
import numpy as np
|
||||
import random
|
||||
from utils.box_utils_Retina import matrix_iof
|
||||
|
||||
|
||||
def _crop(image, boxes, labels, landm, img_dim):
|
||||
height, width, _ = image.shape
|
||||
pad_image_flag = True
|
||||
|
||||
for _ in range(250):
|
||||
"""
|
||||
if random.uniform(0, 1) <= 0.2:
|
||||
scale = 1.0
|
||||
else:
|
||||
scale = random.uniform(0.3, 1.0)
|
||||
"""
|
||||
PRE_SCALES = [0.3, 0.45, 0.6, 0.8, 1.0]
|
||||
scale = random.choice(PRE_SCALES)
|
||||
short_side = min(width, height)
|
||||
w = int(scale * short_side)
|
||||
h = w
|
||||
|
||||
if width == w:
|
||||
l = 0
|
||||
else:
|
||||
l = random.randrange(width - w)
|
||||
if height == h:
|
||||
t = 0
|
||||
else:
|
||||
t = random.randrange(height - h)
|
||||
roi = np.array((l, t, l + w, t + h))
|
||||
|
||||
value = matrix_iof(boxes, roi[np.newaxis])
|
||||
flag = (value >= 1)
|
||||
if not flag.any():
|
||||
continue
|
||||
|
||||
centers = (boxes[:, :2] + boxes[:, 2:]) / 2
|
||||
mask_a = np.logical_and(roi[:2] < centers, centers < roi[2:]).all(axis=1)
|
||||
boxes_t = boxes[mask_a].copy()
|
||||
labels_t = labels[mask_a].copy()
|
||||
landms_t = landm[mask_a].copy()
|
||||
landms_t = landms_t.reshape([-1, 5, 2])
|
||||
|
||||
if boxes_t.shape[0] == 0:
|
||||
continue
|
||||
|
||||
image_t = image[roi[1]:roi[3], roi[0]:roi[2]]
|
||||
|
||||
boxes_t[:, :2] = np.maximum(boxes_t[:, :2], roi[:2])
|
||||
boxes_t[:, :2] -= roi[:2]
|
||||
boxes_t[:, 2:] = np.minimum(boxes_t[:, 2:], roi[2:])
|
||||
boxes_t[:, 2:] -= roi[:2]
|
||||
|
||||
# landm
|
||||
landms_t[:, :, :2] = landms_t[:, :, :2] - roi[:2]
|
||||
landms_t[:, :, :2] = np.maximum(landms_t[:, :, :2], np.array([0, 0]))
|
||||
landms_t[:, :, :2] = np.minimum(landms_t[:, :, :2], roi[2:] - roi[:2])
|
||||
landms_t = landms_t.reshape([-1, 10])
|
||||
|
||||
|
||||
# make sure that the cropped image contains at least one face > 16 pixel at training image scale
|
||||
b_w_t = (boxes_t[:, 2] - boxes_t[:, 0] + 1) / w * img_dim
|
||||
b_h_t = (boxes_t[:, 3] - boxes_t[:, 1] + 1) / h * img_dim
|
||||
mask_b = np.minimum(b_w_t, b_h_t) > 0.0
|
||||
boxes_t = boxes_t[mask_b]
|
||||
labels_t = labels_t[mask_b]
|
||||
landms_t = landms_t[mask_b]
|
||||
|
||||
if boxes_t.shape[0] == 0:
|
||||
continue
|
||||
|
||||
pad_image_flag = False
|
||||
|
||||
return image_t, boxes_t, labels_t, landms_t, pad_image_flag
|
||||
return image, boxes, labels, landm, pad_image_flag
|
||||
|
||||
|
||||
def _distort(image):
|
||||
|
||||
def _convert(image, alpha=1, beta=0):
|
||||
tmp = image.astype(float) * alpha + beta
|
||||
tmp[tmp < 0] = 0
|
||||
tmp[tmp > 255] = 255
|
||||
image[:] = tmp
|
||||
|
||||
image = image.copy()
|
||||
|
||||
if random.randrange(2):
|
||||
|
||||
#brightness distortion
|
||||
if random.randrange(2):
|
||||
_convert(image, beta=random.uniform(-32, 32))
|
||||
|
||||
#contrast distortion
|
||||
if random.randrange(2):
|
||||
_convert(image, alpha=random.uniform(0.5, 1.5))
|
||||
|
||||
image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
|
||||
|
||||
#saturation distortion
|
||||
if random.randrange(2):
|
||||
_convert(image[:, :, 1], alpha=random.uniform(0.5, 1.5))
|
||||
|
||||
#hue distortion
|
||||
if random.randrange(2):
|
||||
tmp = image[:, :, 0].astype(int) + random.randint(-18, 18)
|
||||
tmp %= 180
|
||||
image[:, :, 0] = tmp
|
||||
|
||||
image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR)
|
||||
|
||||
else:
|
||||
|
||||
#brightness distortion
|
||||
if random.randrange(2):
|
||||
_convert(image, beta=random.uniform(-32, 32))
|
||||
|
||||
image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)
|
||||
|
||||
#saturation distortion
|
||||
if random.randrange(2):
|
||||
_convert(image[:, :, 1], alpha=random.uniform(0.5, 1.5))
|
||||
|
||||
#hue distortion
|
||||
if random.randrange(2):
|
||||
tmp = image[:, :, 0].astype(int) + random.randint(-18, 18)
|
||||
tmp %= 180
|
||||
image[:, :, 0] = tmp
|
||||
|
||||
image = cv2.cvtColor(image, cv2.COLOR_HSV2BGR)
|
||||
|
||||
#contrast distortion
|
||||
if random.randrange(2):
|
||||
_convert(image, alpha=random.uniform(0.5, 1.5))
|
||||
|
||||
return image
|
||||
|
||||
|
||||
def _expand(image, boxes, fill, p):
|
||||
if random.randrange(2):
|
||||
return image, boxes
|
||||
|
||||
height, width, depth = image.shape
|
||||
|
||||
scale = random.uniform(1, p)
|
||||
w = int(scale * width)
|
||||
h = int(scale * height)
|
||||
|
||||
left = random.randint(0, w - width)
|
||||
top = random.randint(0, h - height)
|
||||
|
||||
boxes_t = boxes.copy()
|
||||
boxes_t[:, :2] += (left, top)
|
||||
boxes_t[:, 2:] += (left, top)
|
||||
expand_image = np.empty(
|
||||
(h, w, depth),
|
||||
dtype=image.dtype)
|
||||
expand_image[:, :] = fill
|
||||
expand_image[top:top + height, left:left + width] = image
|
||||
image = expand_image
|
||||
|
||||
return image, boxes_t
|
||||
|
||||
|
||||
def _mirror(image, boxes, landms):
|
||||
_, width, _ = image.shape
|
||||
if random.randrange(2):
|
||||
image = image[:, ::-1]
|
||||
boxes = boxes.copy()
|
||||
boxes[:, 0::2] = width - boxes[:, 2::-2]
|
||||
|
||||
# landm
|
||||
landms = landms.copy()
|
||||
landms = landms.reshape([-1, 5, 2])
|
||||
landms[:, :, 0] = width - landms[:, :, 0]
|
||||
tmp = landms[:, 1, :].copy()
|
||||
landms[:, 1, :] = landms[:, 0, :]
|
||||
landms[:, 0, :] = tmp
|
||||
tmp1 = landms[:, 4, :].copy()
|
||||
landms[:, 4, :] = landms[:, 3, :]
|
||||
landms[:, 3, :] = tmp1
|
||||
landms = landms.reshape([-1, 10])
|
||||
|
||||
return image, boxes, landms
|
||||
|
||||
|
||||
def _pad_to_square(image, rgb_mean, pad_image_flag):
|
||||
if not pad_image_flag:
|
||||
return image
|
||||
height, width, _ = image.shape
|
||||
long_side = max(width, height)
|
||||
image_t = np.empty((long_side, long_side, 3), dtype=image.dtype)
|
||||
image_t[:, :] = rgb_mean
|
||||
image_t[0:0 + height, 0:0 + width] = image
|
||||
return image_t
|
||||
|
||||
|
||||
def _resize_subtract_mean(image, insize, rgb_mean):
|
||||
interp_methods = [cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_AREA, cv2.INTER_NEAREST, cv2.INTER_LANCZOS4]
|
||||
interp_method = interp_methods[random.randrange(5)]
|
||||
image = cv2.resize(image, (insize, insize), interpolation=interp_method)
|
||||
image = image.astype(np.float32)
|
||||
image -= rgb_mean
|
||||
return image.transpose(2, 0, 1)
|
||||
|
||||
|
||||
class preproc(object):
|
||||
|
||||
def __init__(self, img_dim, rgb_means):
|
||||
self.img_dim = img_dim
|
||||
self.rgb_means = rgb_means
|
||||
|
||||
def __call__(self, image, targets):
|
||||
assert targets.shape[0] > 0, "this image does not have gt"
|
||||
|
||||
boxes = targets[:, :4].copy()
|
||||
labels = targets[:, -1].copy()
|
||||
landm = targets[:, 4:-1].copy()
|
||||
|
||||
image_t, boxes_t, labels_t, landm_t, pad_image_flag = _crop(image, boxes, labels, landm, self.img_dim)
|
||||
image_t = _distort(image_t)
|
||||
image_t = _pad_to_square(image_t,self.rgb_means, pad_image_flag)
|
||||
image_t, boxes_t, landm_t = _mirror(image_t, boxes_t, landm_t)
|
||||
height, width, _ = image_t.shape
|
||||
image_t = _resize_subtract_mean(image_t, self.img_dim, self.rgb_means)
|
||||
boxes_t[:, 0::2] /= width
|
||||
boxes_t[:, 1::2] /= height
|
||||
|
||||
landm_t[:, 0::2] /= width
|
||||
landm_t[:, 1::2] /= height
|
||||
|
||||
labels_t = np.expand_dims(labels_t, 1)
|
||||
targets_t = np.hstack((boxes_t, landm_t, labels_t))
|
||||
|
||||
return image_t, targets_t
|
||||
@@ -0,0 +1,101 @@
|
||||
import os
|
||||
import os.path
|
||||
import sys
|
||||
import torch
|
||||
import torch.utils.data as data
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
class WiderFaceDetection(data.Dataset):
|
||||
def __init__(self, txt_path, preproc=None):
|
||||
self.preproc = preproc
|
||||
self.imgs_path = []
|
||||
self.words = []
|
||||
f = open(txt_path,'r')
|
||||
lines = f.readlines()
|
||||
isFirst = True
|
||||
labels = []
|
||||
for line in lines:
|
||||
line = line.rstrip()
|
||||
if line.startswith('#'):
|
||||
if isFirst is True:
|
||||
isFirst = False
|
||||
else:
|
||||
labels_copy = labels.copy()
|
||||
self.words.append(labels_copy)
|
||||
labels.clear()
|
||||
path = line[2:]
|
||||
path = txt_path.replace('label.txt','images/') + path
|
||||
self.imgs_path.append(path)
|
||||
else:
|
||||
line = line.split(' ')
|
||||
label = [float(x) for x in line]
|
||||
labels.append(label)
|
||||
|
||||
self.words.append(labels)
|
||||
|
||||
def __len__(self):
|
||||
return len(self.imgs_path)
|
||||
|
||||
def __getitem__(self, index):
|
||||
img = cv2.imread(self.imgs_path[index])
|
||||
height, width, _ = img.shape
|
||||
|
||||
labels = self.words[index]
|
||||
annotations = np.zeros((0, 15))
|
||||
if len(labels) == 0:
|
||||
return annotations
|
||||
for idx, label in enumerate(labels):
|
||||
annotation = np.zeros((1, 15))
|
||||
# bbox
|
||||
annotation[0, 0] = label[0] # x1
|
||||
annotation[0, 1] = label[1] # y1
|
||||
annotation[0, 2] = label[0] + label[2] # x2
|
||||
annotation[0, 3] = label[1] + label[3] # y2
|
||||
|
||||
# landmarks
|
||||
annotation[0, 4] = label[4] # l0_x
|
||||
annotation[0, 5] = label[5] # l0_y
|
||||
annotation[0, 6] = label[7] # l1_x
|
||||
annotation[0, 7] = label[8] # l1_y
|
||||
annotation[0, 8] = label[10] # l2_x
|
||||
annotation[0, 9] = label[11] # l2_y
|
||||
annotation[0, 10] = label[13] # l3_x
|
||||
annotation[0, 11] = label[14] # l3_y
|
||||
annotation[0, 12] = label[16] # l4_x
|
||||
annotation[0, 13] = label[17] # l4_y
|
||||
if (annotation[0, 4]<0):
|
||||
annotation[0, 14] = -1
|
||||
else:
|
||||
annotation[0, 14] = 1
|
||||
|
||||
annotations = np.append(annotations, annotation, axis=0)
|
||||
target = np.array(annotations)
|
||||
if self.preproc is not None:
|
||||
img, target = self.preproc(img, target)
|
||||
|
||||
return torch.from_numpy(img), target
|
||||
|
||||
def detection_collate(batch):
|
||||
"""Custom collate fn for dealing with batches of images that have a different
|
||||
number of associated object annotations (bounding boxes).
|
||||
|
||||
Arguments:
|
||||
batch: (tuple) A tuple of tensor images and lists of annotations
|
||||
|
||||
Return:
|
||||
A tuple containing:
|
||||
1) (tensor) batch of images stacked on their 0 dim
|
||||
2) (list of tensors) annotations for a given image are stacked on 0 dim
|
||||
"""
|
||||
targets = []
|
||||
imgs = []
|
||||
for _, sample in enumerate(batch):
|
||||
for _, tup in enumerate(sample):
|
||||
if torch.is_tensor(tup):
|
||||
imgs.append(tup)
|
||||
elif isinstance(tup, type(np.empty(0))):
|
||||
annos = torch.from_numpy(tup).float()
|
||||
targets.append(annos)
|
||||
|
||||
return (torch.stack(imgs, 0), targets)
|
||||
@@ -4,6 +4,14 @@ from uuid import uuid4
|
||||
import imghdr
|
||||
import traceback
|
||||
import torch
|
||||
# PyTorch 2.6+ 默认 weights_only=True，无法加载含 numpy 对象的旧权重
|
||||
# （yolov5l.pt、各 *.pth 等）。统一改回 False（权重均为本机自有可信文件）。
|
||||
_orig_torch_load = torch.load
|
||||
def _torch_load(*args, **kwargs):
|
||||
kwargs.setdefault("weights_only", False)
|
||||
return _orig_torch_load(*args, **kwargs)
|
||||
torch.load = _torch_load
|
||||
|
||||
from gevent import monkey
|
||||
monkey.patch_all()
|
||||
|
||||
@@ -527,21 +535,41 @@ def change_hairstyle_v4():
|
||||
crop_result = landmark_processor.high_quality_warpAffine(img_res, M, dst_size)
|
||||
# crop_result = landmark_processor.high_quality_warpAffine(img_res, M, (w, h))
|
||||
|
||||
origin_matting = cv2.imread(user_orig_mask_path, cv2.IMREAD_GRAYSCALE)
|
||||
new_matting = cv2.imread(hair_matting_path, cv2.IMREAD_GRAYSCALE)
|
||||
left_point = user_landmarks_origin_img_137[14] # (x1, y1) 获取刘海区域的圆心和半径
|
||||
right_point = user_landmarks_origin_img_137[7] # (x2, y2)
|
||||
center_x = int((left_point[0] + right_point[0]) // 2)
|
||||
center_y = int((left_point[1] + right_point[1]) // 2)
|
||||
radius = int(np.sqrt((right_point[0] - left_point[0]) ** 2 + (right_point[1] - left_point[1]) ** 2) / 2)
|
||||
h, w = origin_matting.shape # 初始化刘海区域为空（全黑色）
|
||||
part_bangs = np.zeros((h, w), dtype=np.uint8)
|
||||
cv2.circle(part_bangs, (center_x, center_y), radius, 255, -1) # 在 part_bangs 上绘制圆作为刘海区域 圆形区域设为白色
|
||||
no_bang_result = cv2.subtract(origin_matting, part_bangs) # 去除刘海区域
|
||||
matting_merge = np.max( np.stack([no_bang_result, new_matting], axis=2), axis=2).astype(np.uint8)# 合并 origin_matting 和 new_matting
|
||||
crop_matting = cv2.warpAffine(matting_merge, M, dst_size)
|
||||
# crop_matting = cv2.warpAffine(matting_merge, M, (w, h))
|
||||
mask = (crop_matting > 10).astype(np.float32)
|
||||
# 接口11：可选外部遮罩 ext_mask（接口9 头发遮罩，原图坐标、白=生发区）。
|
||||
# 传入则经 M 变换到 crop 坐标后替换内部 matting 合成遮罩，webui 精确重绘该区域；
|
||||
# 不传则保持原 matting_merge 逻辑（向后兼容，其它功能不受影响）。
|
||||
ext_mask_b64 = input_info.get('ext_mask', '')
|
||||
if ext_mask_b64:
|
||||
if 'data:image/' in ext_mask_b64:
|
||||
ext_mask_b64 = ext_mask_b64.split(',')[1]
|
||||
ext_mask_np = cv2.imdecode(
|
||||
np.frombuffer(base64.b64decode(ext_mask_b64), np.uint8),
|
||||
cv2.IMREAD_GRAYSCALE)
|
||||
if ext_mask_np is None:
|
||||
raise RuntimeError("ext_mask 解析失败")
|
||||
if ext_mask_np.shape[:2] != origin_img.shape[:2]:
|
||||
ext_mask_np = cv2.resize(
|
||||
ext_mask_np, (origin_img.shape[1], origin_img.shape[0]),
|
||||
interpolation=cv2.INTER_NEAREST)
|
||||
print(f"[swapHair] 使用外部遮罩 ext_mask, shape={ext_mask_np.shape}, "
|
||||
f"白像素={int((ext_mask_np > 10).sum())}", flush=True)
|
||||
crop_matting = cv2.warpAffine(ext_mask_np, M, dst_size)
|
||||
mask = (crop_matting > 10).astype(np.float32)
|
||||
else:
|
||||
origin_matting = cv2.imread(user_orig_mask_path, cv2.IMREAD_GRAYSCALE)
|
||||
new_matting = cv2.imread(hair_matting_path, cv2.IMREAD_GRAYSCALE)
|
||||
left_point = user_landmarks_origin_img_137[14] # (x1, y1) 获取刘海区域的圆心和半径
|
||||
right_point = user_landmarks_origin_img_137[7] # (x2, y2)
|
||||
center_x = int((left_point[0] + right_point[0]) // 2)
|
||||
center_y = int((left_point[1] + right_point[1]) // 2)
|
||||
radius = int(np.sqrt((right_point[0] - left_point[0]) ** 2 + (right_point[1] - left_point[1]) ** 2) / 2)
|
||||
h, w = origin_matting.shape # 初始化刘海区域为空（全黑色）
|
||||
part_bangs = np.zeros((h, w), dtype=np.uint8)
|
||||
cv2.circle(part_bangs, (center_x, center_y), radius, 255, -1) # 在 part_bangs 上绘制圆作为刘海区域 圆形区域设为白色
|
||||
no_bang_result = cv2.subtract(origin_matting, part_bangs) # 去除刘海区域
|
||||
matting_merge = np.max( np.stack([no_bang_result, new_matting], axis=2), axis=2).astype(np.uint8)# 合并 origin_matting 和 new_matting
|
||||
crop_matting = cv2.warpAffine(matting_merge, M, dst_size)
|
||||
mask = (crop_matting > 10).astype(np.float32)
|
||||
if not is_hr:
|
||||
mask_dilate = cv2.dilate(mask, np.ones((3, 9), np.uint8))
|
||||
else:
|
||||
@@ -576,7 +604,7 @@ def change_hairstyle_v4():
|
||||
p_tag = p_tag[p_tag.find("simple background, ") + len("simple background, "):]
|
||||
else:
|
||||
p_tag = ""
|
||||
denoising_strength = 0.6
|
||||
denoising_strength = float(input_info.get('denoising_strength', 0.6)) # 接口11 可调；默认 0.6
|
||||
print(f"功能8：处理发型区域，耗时：{time.time() - start_time:.3f}s")
|
||||
# cv2.imwrite(f"{task_id}_mask_dilate.jpg", mask_dilate)
|
||||
# cv2.imwrite(f"{task_id}_final_img.jpg", final_img)
|
||||
|
||||
@@ -5,16 +5,24 @@ import os
|
||||
|
||||
class OSS_object():
|
||||
def __init__(self):
|
||||
access_key_id = os.getenv('OSS_TEST_ACCESS_KEY_ID', '<your-access-key-id>')
|
||||
access_key_secret = os.getenv('OSS_TEST_ACCESS_KEY_SECRET', '<your-access-key-secret>')
|
||||
bucket_name = os.getenv('OSS_TEST_BUCKET', '<your-bucket-name>')
|
||||
endpoint = os.getenv('OSS_TEST_ENDPOINT', '<your-endpoint>')
|
||||
for param in (access_key_id, access_key_secret, bucket_name, endpoint):
|
||||
assert '<' not in param, '请设置环境变量 OSS_TEST_ACCESS_KEY_ID / OSS_TEST_ACCESS_KEY_SECRET / OSS_TEST_BUCKET / OSS_TEST_ENDPOINT'
|
||||
# 懒加载：只读取凭证，不立即连接 OSS。
|
||||
# 本地用 output_format=base64 时不配 OSS 凭证也能启动服务；
|
||||
# 仅在真正调用 upload_file 时才校验并连接。
|
||||
self.access_key_id = os.getenv('OSS_TEST_ACCESS_KEY_ID', '<your-access-key-id>')
|
||||
self.access_key_secret = os.getenv('OSS_TEST_ACCESS_KEY_SECRET', '<your-access-key-secret>')
|
||||
self.bucket_name = os.getenv('OSS_TEST_BUCKET', '<your-bucket-name>')
|
||||
self.endpoint = os.getenv('OSS_TEST_ENDPOINT', '<your-endpoint>')
|
||||
self.bucket = None
|
||||
|
||||
self.bucket = oss2.Bucket(oss2.Auth(access_key_id, access_key_secret), endpoint, bucket_name)
|
||||
def _ensure_bucket(self):
|
||||
if self.bucket is not None:
|
||||
return
|
||||
for param in (self.access_key_id, self.access_key_secret, self.bucket_name, self.endpoint):
|
||||
assert '<' not in param, '请设置环境变量 OSS_TEST_ACCESS_KEY_ID / OSS_TEST_ACCESS_KEY_SECRET / OSS_TEST_BUCKET / OSS_TEST_ENDPOINT'
|
||||
self.bucket = oss2.Bucket(oss2.Auth(self.access_key_id, self.access_key_secret), self.endpoint, self.bucket_name)
|
||||
|
||||
def upload_file(self, file, target_name):
|
||||
self._ensure_bucket()
|
||||
t0 = time.time()
|
||||
with open(oss2.to_unicode(file), 'rb') as f:
|
||||
ret = self.bucket.put_object(target_name, f)
|
||||
|
||||
@@ -324,7 +324,7 @@ def train_thread(sq, gpu_id):
|
||||
'--caption_extension=".txt" --sample_sampler=ddim '
|
||||
f'--sample_prompts={sample_txt} --sample_every_n_epochs="1000" '
|
||||
'--seed="1234" --cache_latents --optimizer_type="AdamW" --max_data_loader_n_workers="0" --bucket_reso_steps=64 '
|
||||
'--xformers --bucket_no_upscale --noise_offset=0.0')
|
||||
'--sdpa --bucket_no_upscale --noise_offset=0.0')
|
||||
print("cmd_train:", cmd_train)
|
||||
os.system(cmd_train)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user