import cv2
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
from ultralytics import YOLO
import torch
import os
import time
plt.rcParams['font.sans-serif'] = ['SimHei']  # 用黑体显示中文
plt.rcParams['axes.unicode_minus'] = False    # 正常显示负号
class YOLOv8CPUDetector:
    def __init__(self, model_path=None, model_size='n'):
        """
        初始化YOLOv8检测器
        
        Args:
            model_path: 自定义模型路径，如果为None则自动下载
            model_size: 模型大小 'n'(nano), 's'(small), 'm'(medium)
        """
        self.device = 'cpu'
        print(f"使用设备: {self.device}")
        
        if model_path and os.path.exists(model_path):
            # 加载自定义模型
            self.model = YOLO(model_path)
            print(f"加载自定义模型: {model_path}")
        else:
            # 自动下载预训练模型
            model_name = f"yolov8{model_size}.pt"
            print(f"下载模型: {model_name}")
            self.model = YOLO(model_name)
        
        # 设置模型为评估模式
        self.model.to(self.device)
        self.model.eval()
        
        # COCO数据集类别名称
        self.class_names = [
            'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck',
            'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench',
            'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear', 'zebra',
            'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee',
            'skis', 'snowboard', 'sports ball', 'kite', 'baseball bat', 'baseball glove',
            'skateboard', 'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup',
            'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange',
            'broccoli', 'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch',
            'potted plant', 'bed', 'dining table', 'toilet', 'tv', 'laptop', 'mouse',
            'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink',
            'refrigerator', 'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier',
            'toothbrush'
        ]
        
        # 颜色映射
        self.colors = self.generate_colors(len(self.class_names))
    
    def generate_colors(self, n):
        """生成随机颜色"""
        return [tuple(np.random.randint(0, 255, 3).tolist()) for _ in range(n)]
    
    def preprocess_image(self, image_path):
        """图像预处理"""
        if isinstance(image_path, str):
            image = cv2.imread(image_path)
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        else:
            image = image_path
        
        original_size = image.shape[:2]
        return image, original_size
    
    def draw_detections(self, image, boxes, scores, class_ids):
        """绘制检测结果"""
        result_image = image.copy()
        
        for box, score, class_id in zip(boxes, scores, class_ids):
            x1, y1, x2, y2 = box.astype(int)
            class_name = self.class_names[class_id]
            color = self.colors[class_id]
            
            # 绘制边界框
            cv2.rectangle(result_image, (x1, y1), (x2, y2), color, 2)
            
            # 绘制标签背景
            label = f"{class_name}: {score:.2f}"
            label_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2)[0]
            cv2.rectangle(result_image, (x1, y1 - label_size[1] - 10), 
                         (x1 + label_size[0], y1), color, -1)
            
            # 绘制标签文本
            cv2.putText(result_image, label, (x1, y1 - 5), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
        
        return result_image
# 初始化检测器
detector = YOLOv8CPUDetector(model_size='n')

def __init__(self, model_path=None, model_size='n'):
    self.device = 'cpu'
    print(f"使用设备: {self.device}")

if model_path and os.path.exists(model_path):
    # 加载自定义模型
    self.model = YOLO(model_path)
    print(f"加载自定义模型: {model_path}")
else:
    # 自动下载预训练模型
    model_name = f"yolov8{model_size}.pt"
    print(f"下载模型: {model_name}")
    self.model = YOLO(model_name)

# 设置模型为评估模式
self.model.to(self.device)
self.model.eval()

self.class_names = [
    'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck',
    'boat', 'traffic light', 'fire hydrant', 'stop sign', 'parking meter', 'bench',
    # ... 完整列表
]

def generate_colors(self, n):
    """生成随机颜色"""
    return [tuple(np.random.randint(0, 255, 3).tolist()) for _ in range(n)]

image = cv2.imread(image_path)
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

original_size = image.shape[:2]

cv2.rectangle(result_image, (x1, y1), (x2, y2), color, 2)

label_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2)[0]

cv2.rectangle(result_image, (x1, y1 - label_size[1] - 10), 
             (x1 + label_size[0], y1), color, -1)

cv2.putText(result_image, label, (x1, y1 - 5), 
           cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)

def detect_single_image(image_path, conf_threshold=0.25, iou_threshold=0.7):
    """
    单张图像目标检测
    
    Args:
        image_path: 图像路径
        conf_threshold: 置信度阈值
        iou_threshold: IoU阈值
    """
    
    # 预处理图像
    image, original_size = detector.preprocess_image(image_path)
    
    print("开始目标检测...")
    
    # 使用YOLO进行推理
    results = detector.model(
        image, 
        conf=conf_threshold, 
        iou=iou_threshold, 
        device=detector.device,
        verbose=False
    )
    
    # 解析结果
    result = results[0]
    boxes = result.boxes.xyxy.cpu().numpy() if result.boxes is not None else np.array([])
    scores = result.boxes.conf.cpu().numpy() if result.boxes is not None else np.array([])
    class_ids = result.boxes.cls.cpu().numpy().astype(int) if result.boxes is not None else np.array([])
    
    print(f"检测到 {len(boxes)} 个目标")
    
    # 绘制结果
    if len(boxes) > 0:
        result_image = detector.draw_detections(image, boxes, scores, class_ids)
    else:
        result_image = image
    
    # 显示结果
    plt.figure(figsize=(15, 7))
    
    plt.subplot(1, 2, 1)
    plt.imshow(image)
    plt.title("原始图像")
    plt.axis('off')
    
    plt.subplot(1, 2, 2)
    plt.imshow(result_image)
    plt.title(f"检测结果 (检测到 {len(boxes)} 个目标)")
    plt.axis('off')
    
    plt.tight_layout()
    plt.show()
    
    # 保存结果
    output_path = "detection_result.jpg"
    result_image_bgr = cv2.cvtColor(result_image, cv2.COLOR_RGB2BGR)
    cv2.imwrite(output_path, result_image_bgr)
    print(f"结果已保存至: {output_path}")
    
    return boxes, scores, class_ids
# 使用示例
boxes, scores, class_ids = detect_single_image("test_image.jpg", conf_threshold=0.3)

def realtime_camera_detection(camera_id=0, conf_threshold=0.25, show_fps=True):
    """
    实时摄像头目标检测
    
    Args:
        camera_id: 摄像头ID (0=默认摄像头)
        conf_threshold: 置信度阈值
        show_fps: 是否显示FPS
    """
    
    # 打开摄像头
    cap = cv2.VideoCapture(camera_id)
    if not cap.isOpened():
        print("无法打开摄像头")
        return
    
    print("开始实时检测，按 'q' 退出，按 's' 保存当前帧")
    
    fps_counter = 0
    fps_time = time.time()
    
    while True:
        ret, frame = cap.read()
        if not ret:
            print("无法读取视频帧")
            break
        
        # 转换颜色空间
        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        
        # 进行检测
        results = detector.model(
            frame_rgb, 
            conf=conf_threshold, 
            device=detector.device,
            verbose=False
        )
        
        # 解析结果
        result = results[0]
        boxes = result.boxes.xyxy.cpu().numpy() if result.boxes is not None else np.array([])
        scores = result.boxes.conf.cpu().numpy() if result.boxes is not None else np.array([])
        class_ids = result.boxes.cls.cpu().numpy().astype(int) if result.boxes is not None else np.array([])
        
        # 绘制检测结果
        frame_with_detections = detector.draw_detections(frame_rgb, boxes, scores, class_ids)
        frame_bgr = cv2.cvtColor(frame_with_detections, cv2.COLOR_RGB2BGR)
        
        # 计算并显示FPS
        if show_fps:
            fps_counter += 1
            if time.time() - fps_time >= 1.0:
                fps = fps_counter / (time.time() - fps_time)
                fps_text = f"FPS: {fps:.1f}"
                cv2.putText(frame_bgr, fps_text, (10, 30), 
                           cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
                fps_counter = 0
                fps_time = time.time()
        
        # 显示检测数量
        count_text = f"Objects: {len(boxes)}"
        cv2.putText(frame_bgr, count_text, (10, 70), 
                   cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2)
        
        # 显示结果
        cv2.imshow('YOLOv8 Real-time Detection', frame_bgr)
        
        # 按键处理
        key = cv2.waitKey(1) & 0xFF
        if key == ord('q'):
            break
        elif key == ord('s'):
            # 保存当前帧
            timestamp = int(time.time())
            cv2.imwrite(f"capture_{timestamp}.jpg", frame_bgr)
            print(f"截图已保存: capture_{timestamp}.jpg")
    
    # 释放资源
    cap.release()
    cv2.destroyAllWindows()
# 使用示例
realtime_camera_detection(camera_id=0, conf_threshold=0.3)

# 解决方案：减小批处理大小和图像尺寸
def low_memory_detection(image_path):
    image = cv2.imread(image_path)
    # 缩小图像
    image = cv2.resize(image, (0,0), fx=0.5, fy=0.5)
    # 使用较低的置信度阈值
    results = detector.model(image, conf=0.15, device='cpu')
    return results

# 解决方案：优化配置
def speed_optimization():
    # 设置线程数
    torch.set_num_threads(1)  # 单线程可能更快
    # 使用较小的模型
    detector = YOLOv8CPUDetector(model_size='n')
    # 缩小输入尺寸
    detector.model.imgsz = 320

# 解决方案：优化配置
def speed_optimization():
    # 设置线程数
    torch.set_num_threads(1)  # 单线程可能更快
    # 使用较小的模型
    detector = YOLOv8CPUDetector(model_size='n')
    # 缩小输入尺寸
    detector.model.imgsz = 320