import os
import cv2
from ultralytics import YOLO
from pathlib import Path
from PIL import Image, ImageDraw, ImageFont

# ==================== 【配置区域】 ====================
class Config:
    # 你的Plate检测模型路径
    MODEL_PATH = r"best.pt"

    CONF_THRESHOLD = 0.5
    IOU_THRESHOLD = 0.45

    BOX_THICKNESS = 2
    FONT_SIZE = 12

    # 输入输出路径
    INPUT_PATH = r"D:\Comproject\image_classify\trian"
    OUTPUT_DIR = r"D:\Comproject\image_classify\Copy"

    BATCH_MODE = True
    GENERATE_VISUALIZATION = True

    # 只有 Plate 一个类别
    LABEL_NAMES = {0: "Plate"}
    COLORS = {'Plate': (0, 255, 0)}

    IMAGE_EXTENSIONS = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff']

# ========================================================

class YOLOv11PlateExtractor:
    def __init__(self, config):
        self.config = config
        self.model = YOLO(config.MODEL_PATH)
        print("✅ Plate 模型加载成功")

    def predict(self, image_path):
        return self.model.predict(
            source=image_path,
            conf=self.config.CONF_THRESHOLD,
            iou=self.config.IOU_THRESHOLD,
            verbose=False
        )[0]

    def crop_plate(self, image, results):
        cropped_regions = []
        if results.boxes is not None:
            boxes = results.boxes.xyxy.cpu().numpy()
            confs = results.boxes.conf.cpu().numpy()

            h, w = image.shape[:2]
            for box, conf in zip(boxes, confs):
                x1, y1, x2, y2 = map(int, box)
                x1 = max(0, x1)
                y1 = max(0, y1)
                x2 = min(w, x2)
                y2 = min(h, y2)
                if x2 > x1 and y2 > y1:
                    crop = image[y1:y2, x1:x2]
                    if crop.size > 0:
                        cropped_regions.append({"image": crop, "conf": conf})
        return cropped_regions

    def visualize_results(self, image, results, save_path):
        img = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
        draw = ImageDraw.Draw(img)
        try:
            font = ImageFont.truetype("simsun.ttc", self.config.FONT_SIZE)
        except:
            font = ImageFont.load_default()

        if results.boxes is not None:
            for box, conf in zip(results.boxes.xyxy, results.boxes.conf):
                x1, y1, x2, y2 = map(int, box)
                draw.rectangle([(x1, y1), (x2, y2)], outline=(0,255,0), width=self.config.BOX_THICKNESS)
                draw.text((x1, y1-15), f"Plate {conf:.2f}", fill=(0,255,0), font=font)
        img.save(save_path)

    def process_single_image(self, image_path):
        image = cv2.imread(image_path)
        if image is None:
            print(f"❌ 读取失败: {image_path}")
            return

        original_stem = Path(image_path).stem
        results = self.predict(image_path)
        cropped_regions = self.crop_plate(image, results)

        # 所有裁剪图直接保存到 OUTPUT_DIR 根目录
        os.makedirs(self.config.OUTPUT_DIR, exist_ok=True)

        for i, region in enumerate(cropped_regions):
            img_crop = region["image"]
            # 命名：原图名_Plate_序号.png
            filename = f"{original_stem}_Plate_{i+1}.png"
            save_path = os.path.join(self.config.OUTPUT_DIR, filename)
            cv2.imwrite(save_path, img_crop)

        # 可视化
        if self.config.GENERATE_VISUALIZATION:
            vis_dir = os.path.join(self.config.OUTPUT_DIR, "visualization")
            os.makedirs(vis_dir, exist_ok=True)
            self.visualize_results(image, results, os.path.join(vis_dir, f"{original_stem}.jpg"))

        print(f"✅ {original_stem} → 裁剪盘子数量：{len(cropped_regions)}")

    def process_batch(self):
        files = []
        for ext in self.config.IMAGE_EXTENSIONS:
            files.extend(Path(self.config.INPUT_PATH).glob(f"*{ext}"))
        files = sorted(list({str(f) for f in files}))

        print(f"\n📂 共找到 {len(files)} 张图片")
        for f in files:
            self.process_single_image(f)

    def run(self):
        print("\n========== Plate 检测 & 批量裁剪（平铺输出） ==========")
        self.process_batch()
        print("\n🎉 全部处理完成！")

def main():
    cfg = Config()
    app = YOLOv11PlateExtractor(cfg)
    app.run()

if __name__ == "__main__":
    main()