def __init__(self,
cfg,
opt,
save_video=False,
video_save_opt=DEFAULT_VIDEO_SAVE_OPT,
queueSize=1024):
self.cfg = cfg
self.opt = opt
self.video_save_opt = video_save_opt
self.eval_joints = EVAL_JOINTS
self.save_video = save_video
self.final_result = []
self.heatmap_to_coord = get_func_heatmap_to_coord(cfg)
# initialize the queue used to store frames read from
# the video file
if opt.sp:
self.result_queue = Queue(maxsize=queueSize)
self.final_result_queue = Queue(maxsize=queueSize)
else:
self.result_queue = mp.Queue(maxsize=queueSize)
self.final_result_queue = mp.Queue(maxsize=queueSize)
if opt.save_img:
if not os.path.exists(opt.outputpath + '/vis'):
os.mkdir(opt.outputpath + '/vis')
def __init__(self,
cfg,
opt,
save_video=False,
video_save_opt=DEFAULT_VIDEO_SAVE_OPT,
queueSize=1024):
self.cfg = cfg
self.opt = opt
self.video_save_opt = video_save_opt
self.eval_joints = EVAL_JOINTS
self.save_video = save_video
self.heatmap_to_coord = get_func_heatmap_to_coord(cfg)
# initialize the queue used to store frames read from
# the video file
if opt.sp:
self.result_queue = Queue(maxsize=queueSize)
else:
self.result_queue = mp.Queue(maxsize=queueSize)
if opt.save_img:
if not os.path.exists(opt.outputpath + '/vis'):
os.mkdir(opt.outputpath + '/vis')
if opt.pose_flow:
from trackers.PoseFlow.poseflow_infer import PoseFlowWrapper
self.pose_flow_wrapper = PoseFlowWrapper(
save_path=os.path.join(opt.outputpath, 'poseflow'))
def __init__(self, cfg, opt):
self.cfg = cfg
self.opt = opt
self.eval_joints = list(range(cfg.DATA_PRESET.NUM_JOINTS))
self.heatmap_to_coord = get_func_heatmap_to_coord(cfg)
self.item = (None, None, None, None, None, None, None)
def __init__(self, poser_type='AlphaPose', cfg_file=None):
self.poser_type = poser_type
self.cfg_file = cfg_file if cfg_file is not None else f'{os.getcwd()}/poser/config_files/{poser_type}_config_file.yaml'
self.model, self.cfg, self.device = self.init_poser(poser_type)
self.heatmap_to_coord = get_func_heatmap_to_coord(self.cfg)
self.eval_joints = [*range(0, self.cfg.DATA_PRESET.NUM_JOINTS)]
self.colors = [[random.randint(0, 255) for _ in range(3)]
for _ in range(100)] # 最大100种颜色，超过的重新取
def __init__(self, cfg, opt, save_video=False,
video_save_opt=DEFAULT_VIDEO_SAVE_OPT,
queueSize=1024):
self.cfg = cfg
self.opt = opt
self.video_save_opt = video_save_opt
self.eval_joints = EVAL_JOINTS
self.save_video = save_video
self.heatmap_to_coord = get_func_heatmap_to_coord(cfg)
# initialize the queue used to store frames read from
# the video file
if opt.sp:
self.result_queue = Queue(maxsize=queueSize)
else:
self.result_queue = mp.Queue(maxsize=queueSize)
if opt.save_img:
if not os.path.exists(opt.outputpath + '/vis'):
os.mkdir(opt.outputpath + '/vis')
if opt.pose_flow:
from trackers.PoseFlow.poseflow_infer import PoseFlowWrapper
self.pose_flow_wrapper = PoseFlowWrapper(save_path=os.path.join(opt.outputpath, 'poseflow'))
if self.opt.save_img or self.save_video or self.opt.vis:
loss_type = self.cfg.DATA_PRESET.get('LOSS_TYPE', 'MSELoss')
num_joints = self.cfg.DATA_PRESET.NUM_JOINTS
if loss_type == 'MSELoss':
self.vis_thres = [0.4] * num_joints
elif 'JointRegression' in loss_type:
self.vis_thres = [0.05] * num_joints
elif loss_type == 'Combined':
if num_joints == 68:
hand_face_num = 42
else:
hand_face_num = 110
self.vis_thres = [0.4] * (num_joints - hand_face_num) + [0.05] * hand_face_num
self.use_heatmap_loss = (self.cfg.DATA_PRESET.get('LOSS_TYPE', 'MSELoss') == 'MSELoss')
def __init__(self,
cfg,
opt,
save_video=False,
video_save_opt=DEFAULT_VIDEO_SAVE_OPT,
queueSize=1024):
self.cfg = cfg
self.opt = opt
self.video_save_opt = video_save_opt
self.head_pose = opt.head_pose # 是否开启头部姿态相关内容
self.eval_joints = EVAL_JOINTS
self.save_video = save_video
self.heatmap_to_coord = get_func_heatmap_to_coord(cfg)
# initialize the queue used to store frames read from
# the video file
if opt.sp:
self.result_queue = Queue(maxsize=queueSize)
else:
self.result_queue = mp.Queue(maxsize=queueSize)
if opt.save_img:
if not os.path.exists(opt.outputpath + '/vis'):
os.mkdir(opt.outputpath + '/vis')
if opt.pose_flow:
from trackers.PoseFlow.poseflow_infer import PoseFlowWrapper
self.pose_flow_wrapper = PoseFlowWrapper(
save_path=os.path.join(opt.outputpath, 'poseflow'))
if opt.tracking: # 实例状态
self.reid_states = ReIDStates()
# 是否使用场景蒙版
self.scene_mask = SceneMasker(
opt.scene_mask) if opt.scene_mask else None
if self.opt.analyse_cheating and self.opt.save_cheaters:
try:
os.mkdir(os.path.join(self.opt.outputpath, 'cheating'))
except FileExistsError:
pass
def __init__(self, cfg, opt):
self.cfg = cfg
self.opt = opt
self.eval_joints = list(range(cfg.DATA_PRESET.NUM_JOINTS))
self.heatmap_to_coord = get_func_heatmap_to_coord(cfg)
self.item = (None, None, None, None, None, None, None)
loss_type = self.cfg.DATA_PRESET.get('LOSS_TYPE', 'MSELoss')
num_joints = self.cfg.DATA_PRESET.NUM_JOINTS
if loss_type == 'MSELoss':
self.vis_thres = [0.4] * num_joints
elif 'JointRegression' in loss_type:
self.vis_thres = [0.05] * num_joints
elif loss_type == 'Combined':
if num_joints == 68:
hand_face_num = 42
else:
hand_face_num = 110
self.vis_thres = [0.4] * (num_joints -
hand_face_num) + [0.05] * hand_face_num
self.use_heatmap_loss = (self.cfg.DATA_PRESET.get(
'LOSS_TYPE', 'MSELoss') == 'MSELoss')
def main():
logger.info('******************************')
logger.info(opt)
logger.info('******************************')
logger.info(cfg)
logger.info('******************************')
# Model Initialize
m = preset_model(cfg)
m = nn.DataParallel(m).cuda()
criterion = builder.build_loss(cfg.LOSS).cuda()
if cfg.TRAIN.OPTIMIZER == 'adam':
optimizer = torch.optim.Adam(m.parameters(), lr=cfg.TRAIN.LR)
elif cfg.TRAIN.OPTIMIZER == 'rmsprop':
optimizer = torch.optim.RMSprop(m.parameters(), lr=cfg.TRAIN.LR)
lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer, milestones=cfg.TRAIN.LR_STEP, gamma=cfg.TRAIN.LR_FACTOR)
writer = SummaryWriter('.tensorboard/{}-{}'.format(opt.exp_id,
cfg.FILE_NAME))
train_dataset = builder.build_dataset(cfg.DATASET.TRAIN,
preset_cfg=cfg.DATA_PRESET,
train=True)
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=cfg.TRAIN.BATCH_SIZE * num_gpu,
shuffle=True,
num_workers=opt.nThreads)
heatmap_to_coord = get_func_heatmap_to_coord(cfg)
opt.trainIters = 0
for i in range(cfg.TRAIN.BEGIN_EPOCH, cfg.TRAIN.END_EPOCH):
opt.epoch = i
current_lr = optimizer.state_dict()['param_groups'][0]['lr']
logger.info(
f'############# Starting Epoch {opt.epoch} | LR: {current_lr} #############'
)
# Training
loss, miou = train(opt, train_loader, m, criterion, optimizer, writer)
logger.epochInfo('Train', opt.epoch, loss, miou)
lr_scheduler.step()
if (i + 1) % opt.snapshot == 0:
# Save checkpoint
torch.save(
m.module.state_dict(),
'./exp/{}-{}/model_{}.pth'.format(opt.exp_id, cfg.FILE_NAME,
opt.epoch))
# Prediction Test
with torch.no_grad():
gt_AP = validate_gt(m.module, opt, cfg, heatmap_to_coord)
rcnn_AP = validate(m.module, opt, heatmap_to_coord)
logger.info(
f'##### Epoch {opt.epoch} | gt mAP: {gt_AP} | rcnn mAP: {rcnn_AP} #####'
)
# Time to add DPG
if i == cfg.TRAIN.DPG_MILESTONE:
torch.save(
m.module.state_dict(),
'./exp/{}-{}/final.pth'.format(opt.exp_id, cfg.FILE_NAME))
# Adjust learning rate
for param_group in optimizer.param_groups:
param_group['lr'] = cfg.TRAIN.LR
lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer, milestones=cfg.TRAIN.DPG_STEP, gamma=0.1)
# Reset dataset
train_dataset = builder.build_dataset(cfg.DATASET.TRAIN,
preset_cfg=cfg.DATA_PRESET,
train=True,
dpg=True)
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=cfg.TRAIN.BATCH_SIZE * num_gpu,
shuffle=True,
num_workers=opt.nThreads)
torch.save(m.module.state_dict(),
'./exp/{}-{}/final_DPG.pth'.format(opt.exp_id, cfg.FILE_NAME))
def main():
logger.info('******************************')
logger.info(opt)
logger.info('******************************')
logger.info(cfg)
logger.info('******************************')
# Model Initialize
m = preset_model(cfg)
# todo: try to replace with distributedDataParallel to see if it is faster
m = nn.DataParallel(m)
if opt.device.type != 'cpu':
m = m.cuda()
criterion = builder.build_loss(cfg.LOSS)
if opt.device.type != 'cpu':
criterion = criterion.cuda()
if cfg.TRAIN.OPTIMIZER == 'adam':
optimizer = torch.optim.Adam(m.parameters(), lr=cfg.TRAIN.LR)
elif cfg.TRAIN.OPTIMIZER == 'rmsprop':
optimizer = torch.optim.RMSprop(m.parameters(), lr=cfg.TRAIN.LR)
lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer, milestones=cfg.TRAIN.LR_STEP, gamma=cfg.TRAIN.LR_FACTOR)
if opt.clean:
if opt.tensorboard_path.exists():
shutil.rmtree(opt.tensorboard_path)
if opt.experiment_path.exists():
shutil.rmtree(opt.experiment_path)
opt.tensorboard_path.mkdir(exist_ok=True, parents=True)
opt.experiment_path.mkdir(exist_ok=True, parents=True)
writer = SummaryWriter(str(opt.tensorboard_path))
train_dataset = builder.build_dataset(cfg.DATASET.TRAIN,
preset_cfg=cfg.DATA_PRESET,
train=True)
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=cfg.TRAIN.BATCH_SIZE * max(1, num_gpu),
shuffle=True,
num_workers=opt.nThreads)
heatmap_to_coord = get_func_heatmap_to_coord(cfg)
opt.trainIters = 0
scaler = GradScaler()
for i in range(cfg.TRAIN.BEGIN_EPOCH, cfg.TRAIN.END_EPOCH):
opt.epoch = i
current_lr = optimizer.state_dict()['param_groups'][0]['lr']
logger.info(
f'############# Starting Epoch {opt.epoch} | LR: {current_lr} #############'
)
# Training
loggers = train(opt, train_loader, m, criterion, optimizer, writer,
scaler)
logger.info(
f'Train-{opt.epoch:d} epoch | '
f'{" | ".join(f"{name}:{l.avg:.07f}" for name, l in loggers.items())}'
)
lr_scheduler.step()
if (i + 1) % opt.snapshot == 0:
# Save checkpoint
torch.save(m.module.state_dict(),
str(opt.experiment_path / f'model_{opt.epoch}.pth'))
# Prediction Test
with torch.no_grad():
metrics_on_true_box = validate_gt(m.module, opt, cfg,
heatmap_to_coord)
gt_AP = metrics_on_true_box["map"]
gt_radius_mse = metrics_on_true_box["radius_mse"]
rcnn_AP = validate(m.module, opt, heatmap_to_coord)
logger.info(f'##### Epoch {opt.epoch} | '
f'gt mAP: {gt_AP} | '
f'rcnn mAP: {rcnn_AP} | '
f'gt radius_mse {gt_radius_mse}'
f' #####')
writer.add_scalar(f'Validation/mAP_on_gt_box', gt_AP,
opt.trainIters)
writer.add_scalar(f'Validation/mAP_on_pred_box', rcnn_AP,
opt.trainIters)
writer.add_scalar(f'Validation/radius_mse_on_gt_box',
gt_radius_mse, opt.trainIters)
# Time to add DPG
if i == cfg.TRAIN.DPG_MILESTONE:
torch.save(m.module.state_dict(),
str(opt.experiment_path / "final.pth"))
# Adjust learning rate
for param_group in optimizer.param_groups:
param_group['lr'] = cfg.TRAIN.LR
lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
optimizer, milestones=cfg.TRAIN.DPG_STEP, gamma=0.1)
# Reset dataset
train_dataset = builder.build_dataset(cfg.DATASET.TRAIN,
preset_cfg=cfg.DATA_PRESET,
train=True,
dpg=True)
train_loader = torch.utils.data.DataLoader(
train_dataset,
batch_size=cfg.TRAIN.BATCH_SIZE * max(1, num_gpu),
shuffle=True,
num_workers=opt.nThreads)
torch.save(m.module.state_dict(),
str(opt.experiment_path / 'final_DPG.pth'))
data['image_id'] = int(img_ids[i])
data['score'] = float(np.mean(pose_scores) + np.max(pose_scores))
data['category_id'] = 1
data['keypoints'] = keypoints
kpt_json.append(data)
with open('./exp/json/validate_gt_kpt.json', 'w') as fid:
json.dump(kpt_json, fid)
res = evaluate_mAP('./exp/json/validate_gt_kpt.json',
ann_type='keypoints',
ann_file=os.path.join(cfg.DATASET.VAL.ROOT,
cfg.DATASET.VAL.ANN))
return res['AP']
if __name__ == "__main__":
m = builder.build_sppe(cfg.MODEL, preset_cfg=cfg.DATA_PRESET)
print(f'Loading model from {opt.checkpoint}...')
m.load_state_dict(torch.load(opt.checkpoint))
m = torch.nn.DataParallel(m, device_ids=gpus).cuda()
heatmap_to_coord = get_func_heatmap_to_coord(cfg)
with torch.no_grad():
gt_AP = validate_gt(m, cfg, heatmap_to_coord, opt.batch)
detbox_AP = validate(m, heatmap_to_coord, opt.batch)
print('##### gt box: {} mAP | det box: {} mAP #####'.format(
gt_AP, detbox_AP))
def __init__(self, args=None):
if args is None:
args = Namespace(
# Pose config
pose_cfg='configs/coco/resnet/256x192_res50_lr1e-3_1x.yaml',
# Pose checkpoint
pose_checkpoint='pretrained_models/fast_res50_256x192.pth',
# GPUS
gpus='0',
# Detection thresh
det_thresh=0.5,
# Detection config
det_cfg='mmDetection/gfl_x101_611.py',
# Detection checkpoint
det_checkpoint='mmDetection/weights.pth',
# Show clothe color
clothe_color=True,
# show bboxes
showbox=True
)
self.pose_cfg = update_config(args.pose_cfg)
# Device configuration
args.gpus = [int(i) for i in args.gpus.split(',')] if torch.cuda.device_count() >= 1 else [-1]
args.device = torch.device("cuda:" + str(args.gpus[0]) if args.gpus[0] >= 0 else "cpu")
args.tracking = False
args.pose_track = False
# Copy args
self.args = args
# Preprocess transformation
pose_dataset = builder.retrieve_dataset(self.pose_cfg.DATASET.TRAIN)
self.transformation = SimpleTransform(
pose_dataset, scale_factor=0,
input_size=self.pose_cfg.DATA_PRESET.IMAGE_SIZE,
output_size=self.pose_cfg.DATA_PRESET.HEATMAP_SIZE,
rot=0, sigma=self.pose_cfg.DATA_PRESET.SIGMA,
train=False, add_dpg=False, gpu_device=args.device)
self.norm_type = self.pose_cfg.LOSS.get('NORM_TYPE', None)
# Post process
self.heatmap_to_coord = get_func_heatmap_to_coord(self.pose_cfg)
# Load Detector Model
self.det_model = init_detector(args.det_cfg, checkpoint=args.det_checkpoint, device=args.device)
# Load pose model
self.pose_model = builder.build_sppe(self.pose_cfg.MODEL, preset_cfg=self.pose_cfg.DATA_PRESET)
print(f'Loading pose model from {args.pose_checkpoint}...')
self.pose_model.load_state_dict(torch.load(args.pose_checkpoint, map_location=args.device))
self.pose_model.to(args.device)
self.pose_model.eval()
