import cv2
import numpy as np
from ultralytics import YOLO

def extract_clothing_area(image_path, output_path="output.jpg"):
    # 加载预训练的人体检测模型（YOLOv8）
    model = YOLO("/data/data/luosy/project/oral/models/yolov8n.pt")  # 使用轻量级模型

    # 读取图像
    image = cv2.imread(image_path)
    if image is None:
        raise FileNotFoundError(f"Image not found at {image_path}")

    # 人体检测
    results = model.predict(image, classes=[0])  # classes=[0] 表示只检测人（COCO类别中0为"person"）
  
    # 提取检测到的人体区域
    for result in results:
        boxes = result.boxes.xyxy.cpu().numpy()  # 获取边界框坐标
        if len(boxes) == 0:
            print("未检测到人体")
            return
      
        # 选择第一个检测到的人体（假设图像中只有一人）
        x1, y1, x2, y2 = map(int, boxes[0])
        human_roi = image[y1:y2, x1:x2]

        # 简单方法：假设服装区域为人体上半部分（可调整）
        height = y2 - y1
        clothing_y1 = int(height * 0.1)  # 顶部留空10%
        clothing_y2 = int(height * 0.6)  # 上半部分（可根据需求调整）
        clothing_roi = human_roi[clothing_y1:clothing_y2, :]

        # 保存服装区域
        cv2.imwrite(output_path, clothing_roi)
        print(f"服装区域已保存至 {output_path}")

if __name__ == "__main__":
    # 使用示例
    image_path1 = '/data/data/luosy/project/oral/data/key_frame/frame_013.jpg'
    image_path2 = '/data/data/luosy/project/oral/data/key_frame/frame_013_detect.jpg'
    extract_clothing_area(image_path1, image_path2)