整体代码如下:
# 加载图像并转换为数组的形式,调整大小
def figure_trans():
# 加载图像并转换为数组的形式,调整大小
img_path = "E:/IRTC/ceshi/img.png"
image = Image.open(img_path)
transform = transforms.ToTensor()
tensor_image = transform(image)
transform_size = transforms.Resize((512, 512))
out = transform_size(tensor_image)
tensor_image_batch = out.unsqueeze(0)
return tensor_image_batch
# 加载模型并绘制模型
def yml_plot():
# yaml_file_path = 'F:/code/DEKR-main/experiments/coco/w32/w32_4x_reg03_bs10_512_adam_lr1e-3_coco_x140.yaml'
# yaml_file = Path(yaml_file_path)
tensor_image_batch = figure_trans()
# if yaml_file.exists():
# with yaml_file.open('r') as file:
# cfg_dict = yaml.safe_load(file)
# cfg = OmegaConf.create(cfg_dict)
out_img3 = pose(pretrained=True)
# out_img3 = PoseHigherResolutionNet(cfg)
# input = torch.randn(1, 3, 512, 512)
out_img3 = out_img3.forward(tensor_image_batch)
# 选择要可视化的通道索引
selected_channels = [1, 15, 30, 45, 63] # 例如,选择第0、15、30、45和63个通道
# 创建一个包含多个子图的图形
fig, axes = plt.subplots(1, len(selected_channels), figsize=(15, 5))
# 遍历选定的通道并绘制它们
for i, ax in enumerate(axes):
channel_image = out_img3[selected_channels[i]] # 提取选定通道的图像
print(channel_image.shape)
out_img1 = transforms.ToPILImage()(channel_image)
ax.imshow(out_img1, cmap='viridis') # 使用 viridis 颜色图显示图像
ax.set_title(f'Channel {selected_channels[i]}') # 设置子图的标题
ax.axis('off') # 关闭坐标轴
# 显示图形
plt.tight_layout() # 调整子图参数, 使之填充整个图像区域
plt.show()
# print(pose.forward(input))
# else:
# print(f"Error: The YAML file {yaml_file_path} does not exist.")
model_urls = {
'pose': 'F:/code/DEKR-main/model/pose_coco/pose_dekr_hrnetw32_coco.pth',
}
# 加载本地预训练模型并返回模型
def pose(pretrained=False):
yaml_file_path = 'F:/code/DEKR-main/experiments/coco/w32/w32_4x_reg03_bs10_512_adam_lr1e-3_coco_x140.yaml'
yaml_file = Path(yaml_file_path)
if yaml_file.exists():
with yaml_file.open('r') as file:
cfg_dict = yaml.safe_load(file)
cfg = OmegaConf.create(cfg_dict)
model = PoseHigherResolutionNet(cfg)
if pretrained:
# 加载本地预训练模型
model_path = model_urls['pose']
model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu'))) # 如果你在GPU上运行,可能不需要map_location
return model
plt_figure = yml_plot()
这里新加的就是加载本地训练好的模型文件pth
model_urls = {
'pose': 'F:/code/DEKR-main/model/pose_coco/pose_dekr_hrnetw32_coco.pth',
}
并加载进来
model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu'))) # 如果你在GPU上运行,可能不需要map_location