243 lines
8.4 KiB
Python
243 lines
8.4 KiB
Python
"""
|
||
SMPL模型拟合主程序
|
||
该脚本用于将人体姿态数据拟合到SMPL（Skinned Multi-Person Linear）模型中
|
||
主要功能：
|
||
1. 加载人体姿态数据
|
||
2. 使用SMPL模型进行拟合优化
|
||
3. 保存拟合结果和可视化图像
|
||
"""
|
||
|
||
import os
|
||
import sys
|
||
sys.path.append(os.getcwd())
|
||
|
||
# 导入自定义模块
|
||
from meters import Meters # 用于跟踪训练指标的工具类
|
||
from smplpytorch.pytorch.smpl_layer import SMPL_Layer # SMPL模型层
|
||
from train import train # 训练函数
|
||
from transform import transform # 数据变换函数
|
||
from save import save_pic, save_params # 保存结果的函数
|
||
from load import load # 数据加载函数
|
||
|
||
# 导入标准库
|
||
import torch # PyTorch深度学习框架
|
||
import numpy as np # 数值计算库
|
||
from easydict import EasyDict as edict # 用于创建字典对象的便捷工具
|
||
import time # 时间处理
|
||
import logging # 日志记录
|
||
|
||
import argparse # 命令行参数解析
|
||
import json # JSON文件处理
|
||
|
||
# 启用CUDNN加速，提高训练效率
|
||
torch.backends.cudnn.enabled = True
|
||
torch.backends.cudnn.benchmark = True
|
||
|
||
|
||
def parse_args():
|
||
"""
|
||
解析命令行参数
|
||
|
||
Returns:
|
||
args: 包含解析后参数的命名空间对象
|
||
"""
|
||
parser = argparse.ArgumentParser(description='Fit SMPL')
|
||
|
||
# 实验名称，默认使用当前时间戳
|
||
parser.add_argument('--exp', dest='exp',
|
||
help='Define exp name',
|
||
default=time.strftime('%Y-%m-%d %H-%M-%S', time.localtime(time.time())), type=str)
|
||
|
||
# 数据集名称，用于选择对应的配置文件
|
||
parser.add_argument('--dataset_name', '-n', dest='dataset_name',
|
||
help='select dataset',
|
||
default='', type=str)
|
||
|
||
# 数据集路径，可以覆盖配置文件中的默认路径
|
||
parser.add_argument('--dataset_path', dest='dataset_path',
|
||
help='path of dataset',
|
||
default=None, type=str)
|
||
|
||
args = parser.parse_args()
|
||
return args
|
||
|
||
|
||
def get_config(args):
|
||
"""
|
||
根据数据集名称加载对应的配置文件
|
||
|
||
Args:
|
||
args: 命令行参数对象
|
||
|
||
Returns:
|
||
cfg: 配置对象，包含所有训练和模型参数
|
||
"""
|
||
# 根据数据集名称构建配置文件路径
|
||
config_path = 'fit/configs/{}.json'.format(args.dataset_name)
|
||
|
||
# 读取JSON配置文件
|
||
with open(config_path, 'r') as f:
|
||
data = json.load(f)
|
||
|
||
# 将字典转换为edict对象，支持点号访问属性
|
||
cfg = edict(data.copy())
|
||
|
||
# 如果命令行指定了数据集路径，则覆盖配置文件中的设置
|
||
if not args.dataset_path == None:
|
||
cfg.DATASET.PATH = args.dataset_path
|
||
|
||
return cfg
|
||
|
||
|
||
def set_device(USE_GPU):
|
||
"""
|
||
根据配置和硬件可用性设置计算设备
|
||
|
||
Args:
|
||
USE_GPU: 是否使用GPU的布尔值
|
||
|
||
Returns:
|
||
device: PyTorch设备对象（'cuda' 或 'cpu'）
|
||
"""
|
||
if USE_GPU and torch.cuda.is_available():
|
||
device = torch.device('cuda')
|
||
else:
|
||
device = torch.device('cpu')
|
||
return device
|
||
|
||
|
||
def get_logger(cur_path):
|
||
"""
|
||
设置日志记录器，同时输出到文件和控制台
|
||
|
||
Args:
|
||
cur_path: 当前实验路径，用于保存日志文件
|
||
|
||
Returns:
|
||
logger: 日志记录器对象
|
||
writer: TensorBoard写入器（当前设置为None）
|
||
"""
|
||
# 创建日志记录器
|
||
logger = logging.getLogger(__name__)
|
||
logger.setLevel(level=logging.INFO)
|
||
|
||
# 设置文件输出处理器，将日志保存到log.txt文件
|
||
handler = logging.FileHandler(os.path.join(cur_path, "log.txt"))
|
||
handler.setLevel(logging.INFO)
|
||
formatter = logging.Formatter(
|
||
'%(asctime)s - %(name)s - %(levelname)s - %(message)s')
|
||
handler.setFormatter(formatter)
|
||
logger.addHandler(handler)
|
||
|
||
# 设置控制台输出处理器，将日志同时输出到终端
|
||
handler = logging.StreamHandler()
|
||
handler.setLevel(logging.INFO)
|
||
formatter = logging.Formatter(
|
||
'%(asctime)s - %(name)s - %(levelname)s - %(message)s')
|
||
handler.setFormatter(formatter)
|
||
logger.addHandler(handler)
|
||
|
||
# TensorBoard写入器（目前被注释掉，设置为None）
|
||
# from tensorboardX import SummaryWriter
|
||
# writer = SummaryWriter(os.path.join(cur_path, 'tb'))
|
||
writer = None
|
||
|
||
return logger, writer
|
||
|
||
|
||
if __name__ == "__main__":
|
||
"""
|
||
主函数：执行SMPL模型拟合流程
|
||
|
||
主要步骤：
|
||
1. 解析命令行参数
|
||
2. 创建实验目录
|
||
3. 加载配置文件
|
||
4. 设置日志记录
|
||
5. 初始化SMPL模型
|
||
6. 遍历数据集进行拟合
|
||
7. 保存结果
|
||
"""
|
||
# 解析命令行参数
|
||
args = parse_args()
|
||
|
||
# 创建实验目录，使用时间戳或用户指定的实验名称
|
||
cur_path = os.path.join(os.getcwd(), 'exp', args.exp)
|
||
assert not os.path.exists(cur_path), 'Duplicate exp name' # 确保实验名称不重复
|
||
os.mkdir(cur_path)
|
||
|
||
# 加载配置文件
|
||
cfg = get_config(args)
|
||
# 将配置保存到实验目录中，便于后续追踪实验设置
|
||
json.dump(dict(cfg), open(os.path.join(cur_path, 'config.json'), 'w'))
|
||
|
||
# 设置日志记录器
|
||
logger, writer = get_logger(cur_path)
|
||
logger.info("Start print log")
|
||
|
||
# 设置计算设备（GPU或CPU）
|
||
device = set_device(USE_GPU=cfg.USE_GPU)
|
||
logger.info('using device: {}'.format(device))
|
||
|
||
# 初始化SMPL模型层
|
||
# center_idx=0: 设置中心关节点索引
|
||
# gender='male': 设置性别为男性（注释掉的cfg.MODEL.GENDER可能用于从配置文件读取）
|
||
# model_root: SMPL模型文件的路径
|
||
smpl_layer = SMPL_Layer(
|
||
center_idx=0,
|
||
gender='male', #cfg.MODEL.GENDER,
|
||
model_root='smplpytorch/native/models')
|
||
|
||
# 初始化指标记录器，用于跟踪训练损失等指标
|
||
meters = Meters()
|
||
file_num = 0
|
||
# 遍历数据集目录中的所有文件
|
||
for root, dirs, files in os.walk(cfg.DATASET.PATH):
|
||
for file in sorted(files): # 按文件名排序处理
|
||
# 可选的文件过滤器（当前被注释掉）
|
||
# 可以用于只处理特定的文件，如包含'baseball_swing'的文件
|
||
###if not 'baseball_swing' in file:
|
||
###continue
|
||
|
||
file_num += 1 # 文件计数器
|
||
logger.info(
|
||
'Processing file: {} [{} / {}]'.format(file, file_num, len(files)))
|
||
|
||
# 加载并变换目标数据
|
||
# 1. load(): 根据数据集类型加载原始数据
|
||
# 2. transform(): 将数据转换为模型所需的格式
|
||
# 3. torch.from_numpy(): 将numpy数组转换为PyTorch张量
|
||
# 4. .float(): 确保数据类型为float32
|
||
target = torch.from_numpy(transform(args.dataset_name, load(args.dataset_name,
|
||
os.path.join(root, file)))).float()
|
||
logger.info("target shape:{}".format(target.shape))
|
||
|
||
# 执行SMPL模型拟合训练
|
||
# 传入SMPL层、目标数据、日志记录器、设备信息等
|
||
res = train(smpl_layer, target,
|
||
logger, writer, device,
|
||
args, cfg, meters)
|
||
|
||
# 更新平均损失指标
|
||
# k=target.shape[0] 表示批次大小，用于加权平均
|
||
meters.update_avg(meters.min_loss, k=target.shape[0])
|
||
|
||
# 重置早停计数器，为下一个文件的训练做准备
|
||
meters.reset_early_stop()
|
||
|
||
# 记录当前的平均损失
|
||
logger.info("avg_loss:{:.4f}".format(meters.avg))
|
||
|
||
# 保存拟合结果
|
||
# 1. save_params(): 保存拟合得到的SMPL参数
|
||
# 2. save_pic(): 保存可视化图像，包括拟合结果和原始目标的对比
|
||
save_params(res, file, logger, args.dataset_name)
|
||
save_pic(res, smpl_layer, file, logger, args.dataset_name, target)
|
||
|
||
# 清空GPU缓存，防止内存溢出
|
||
torch.cuda.empty_cache()
|
||
|
||
# 所有文件处理完成，记录最终的平均损失
|
||
logger.info(
|
||
"Fitting finished! Average loss: {:.9f}".format(meters.avg))
|