lmd/Pose_to_SMPL_an_230402
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

name

Browse Source
This commit is contained in:
lmd
2025-07-25 15:05:31 +08:00
parent 24e1e31234
commit 65769e5eb6
18 changed files with 3642 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
fit/configs/CMU_Mocap.json
View File

@ -12,7 +12,7 @@
"USE_GPU": 1,
"DATASET": {
"NAME": "UTD-MHAD",
"PATH": "../Action2Motion/CMU Mocap/mocap/mocap_3djoints/",
"PATH": "/home/lmd/Code/Pose_to_SMPL_an_230402/database",
"TARGET_PATH": "",
"DATA_MAP": [
[

62
fit/configs/UTD_MHAD.json
View File

@ -12,72 +12,100 @@
"USE_GPU": 1,
"DATASET": {
"NAME": "UTD-MHAD",
"PATH": "../UTD-MHAD/Skeleton/Skeleton/",
"PATH": "/home/lmd/Code/Pose_to_SMPL_an_230402/database",
"TARGET_PATH": "",
"DATA_MAP": [
[
12,
1,
1
],
[
2,
2
],
[
0,
3,
3
],
[
16,
4,
4
],
[
18,
5,
5
],
[
20,
6,
6
],
[
22,
7,
7
],
[
17,
8,
8
],
[
19,
9,
9
],
[
21,
10,
10
],
[
23,
11,
11
],
[
1,
12,
12
],
[
4,
13,
13
],
[
7,
14,
14
],
[
15,
15
],
[
2,
16,
16
],
[
5,
17,
17
],
[
8,
18,
18
],
[
19,
19
],
[
20,
20
],
[
21,
21
],
[
22,
22
],
[
23,
23
]
]
},

5
fit/tools/label.py
View File

@ -1284,8 +1284,9 @@ def get_label(file_name, dataset_name):
key = file_name[-5:]
return HumanAct12[key]
elif dataset_name == 'UTD_MHAD':
key = file_name.split('_')[0][1:]
return UTD_MHAD[key]
###key = file_name.split('_')[0][1:]
###return UTD_MHAD[key]
return "single_person"
elif dataset_name == 'CMU_Mocap':
key = file_name.split(':')[0]
return CMU_Mocap[key] if key in CMU_Mocap.keys() else ""

41
fit/tools/load.py
View File

@ -1,25 +1,46 @@
import scipy.io
import numpy as np
import json
import json # 引入json模块
def load(name, path):
# 处理UTD-MHAD的JSON文件（你的单帧数据）
if name == 'UTD_MHAD':
arr = scipy.io.loadmat(path)['d_skel']
new_arr = np.zeros([arr.shape[2], arr.shape[0], arr.shape[1]])
for i in range(arr.shape[2]):
for j in range(arr.shape[0]):
for k in range(arr.shape[1]):
new_arr[i][j][k] = arr[j][k][i]
return new_arr
# 判断文件是否为JSON格式（通过后缀）
if path.endswith('.json'):
with open(path, 'r') as f:
data = json.load(f) # 加载JSON列表
# 转换为NumPy数组，确保形状为[关节数, 3]
data_np = np.array(data)
# 校验数据格式（防止错误）
assert data_np.ndim == 2 and data_np.shape[1] == 3, \
f"UTD-MHAD JSON格式错误，应为[关节数, 3]，实际为{data_np.shape}"
# 若需要单帧维度（[1, 关节数, 3]），可扩展维度
return data_np[np.newaxis, ...] # 输出形状：[1, N, 3]（1表示单帧）
# 保留原UTD_MHAD的.mat文件支持（如果还需要处理.mat数据）
elif path.endswith('.mat'):
arr = scipy.io.loadmat(path)['d_skel']
new_arr = np.zeros([arr.shape[2], arr.shape[0], arr.shape[1]])
for i in range(arr.shape[2]):
for j in range(arr.shape[0]):
for k in range(arr.shape[1]):
new_arr[i][j][k] = arr[j][k][i]
return new_arr
else:
raise ValueError(f"UTD-MHAD不支持的文件格式：{path}")
# 其他数据集的原有逻辑保持不变
elif name == 'HumanAct12':
return np.load(path, allow_pickle=True)
elif name == "CMU_Mocap":
return np.load(path, allow_pickle=True)
elif name == "Human3.6M":
return np.load(path, allow_pickle=True)[0::5] # down_sample
return np.load(path, allow_pickle=True)[0::5] # 下采样
elif name == "NTU":
return np.load(path, allow_pickle=True)[0::2]
elif name == "HAA4D":
return np.load(path, allow_pickle=True)
else:
raise ValueError(f"不支持的数据集名称：{name}")

158
fit/tools/main.py
View File

@ -1,53 +1,104 @@
"""
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
from train import train
from transform import transform
from save import save_pic, save_params
from load import load
import torch
import numpy as np
from tensorboardX import SummaryWriter
from easydict import EasyDict as edict
import time
import logging
import argparse
import json
# 导入自定义模块
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:
@ -56,9 +107,21 @@ def set_device(USE_GPU):
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(
@ -66,6 +129,7 @@ def get_logger(cur_path):
handler.setFormatter(formatter)
logger.addHandler(handler)
# 设置控制台输出处理器，将日志同时输出到终端
handler = logging.StreamHandler()
handler.setLevel(logging.INFO)
formatter = logging.Formatter(
@ -73,54 +137,106 @@ def get_logger(cur_path):
handler.setFormatter(formatter)
logger.addHandler(handler)
writer = SummaryWriter(os.path.join(cur_path, 'tb'))
# 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'
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=cfg.MODEL.GENDER,
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):
if not 'baseball_swing' in file:
continue
file_num += 1
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))

6
fit/tools/train.py
View File

@ -78,9 +78,9 @@ def train(smpl_layer, target,
# epoch, float(loss),float(scale)))
print("Epoch {}, lossPerBatch={:.6f}, scale={:.4f}".format(
epoch, float(loss),float(scale)))
writer.add_scalar('loss', float(loss), epoch)
writer.add_scalar('learning_rate', float(
optimizer.state_dict()['param_groups'][0]['lr']), epoch)
###writer.add_scalar('loss', float(loss), epoch)
###writer.add_scalar('learning_rate', float(
###optimizer.state_dict()['param_groups'][0]['lr']), epoch)
# save_single_pic(res,smpl_layer,epoch,logger,args.dataset_name,target)
logger.info('Train ended, min_loss = {:.4f}'.format(

2
fit/tools/transform.py
View File

@ -3,7 +3,7 @@ import numpy as np
rotate = {
'HumanAct12': [1., -1., -1.],
'CMU_Mocap': [0.05, 0.05, 0.05],
'UTD_MHAD': [-1., 1., -1.],
'UTD_MHAD': [1., 1., 1.],
'Human3.6M': [-0.001, -0.001, 0.001],
'NTU': [1., 1., -1.],
'HAA4D': [1., -1., -1.],