基于YOLOv8的关键点估计：实现图像上传与结果可视化

# 安装ultralytics库 (如果尚未安装)
!pip install ultralytics

# 导入YOLO类
from ultralytics import YOLO

# 加载预训练的关键点估计模型 (请替换为你的模型路径)
# 假设你已经有一个名为 'your_pose_model.pt' 的模型文件
model = YOLO('yolov8n-pose.pt') # 例如，加载YOLOv8n姿态估计模型
# 如果是自定义训练的模型，路径可能类似 model = YOLO('/path/to/your/custom_pose_model.pt')

from google.colab import files
import io
from PIL import Image

uploaded = files.upload()

# 获取上传文件的名称
# 假设只上传一张图片
for filename in uploaded.keys():
    print(f'用户已上传文件: {filename}')
    input_image_path = filename
    break # 只处理第一个上传的文件

# 对上传的图像执行推理
# 关键点：设置 save=True 以保存带标注的结果图像
results = model.predict(source=input_image_path, save=True, conf=0.25)

# 打印推理结果摘要 (可选)
for r in results:
    # 打印检测到的关键点信息等
    print(f"检测到 {len(r.keypoints)} 关键点实例。")
    # r.save_dir 包含了结果保存的目录
    output_dir = r.save_dir
    print(f"带关键点标注的图像已保存到: {output_dir}")

# 假设结果会保存在 runs/pose/predict/ 目录下，且文件名为原文件名
# 需要找到实际保存的图像路径
# YOLOv8通常会将结果保存在 runs/pose/predict/ 或 runs/detect/predict/ 目录下
# 并且会保留原始文件名
import os
# 构建输出图像的预期路径
# 注意：如果多次运行，predict目录名可能会变为 predict2, predict3等
# r.save_dir 是最准确的输出目录
processed_image_name = os.path.basename(input_image_path)
output_image_path = os.path.join(output_dir, processed_image_name)

print(f"尝试加载的输出图像路径: {output_image_path}")

import matplotlib.pyplot as plt
import matplotlib.image as mpimg

# 检查文件是否存在
if os.path.exists(output_image_path):
    # 加载并显示处理后的图像
    img = mpimg.imread(output_image_path)
    plt.figure(figsize=(10, 8))
    plt.imshow(img)
    plt.axis('off') # 不显示坐标轴
    plt.title('YOLOv8 关键点估计结果')
    plt.show()
else:
    print(f"错误：未找到处理后的图像文件，路径为: {output_image_path}")
    print("请检查YOLOv8的输出目录结构或推理参数。")