pip install torch torchvision
pip install opencv-python pillow
pip install ultralytics
pip install numpy tqdm pathlib argparse

class FastSAMAutoLabeler:
    def __init__(self, model_path='FastSAM-x.pt', device='cuda' if torch.cuda.is_available() else 'cpu', classification_model='yolov8n-cls.pt'):
        self.device = device
        self.model = FastSAM(model_path)
        self.classification_model = YOLO(classification_model)

def process_image(self, image_path, output_dir, conf=0.4, iou=0.9, min_area_ratio=0.001, max_area_ratio=0.95):
    # 讀取圖片
    image = cv2.imread(image_path)
    height, width = image.shape[:2]
    
    # FastSAM推理
    everything_results = self.model(image_path, device=self.device, retina_masks=True, imgsz=1024, conf=conf, iou=iou)
    
    # 處理掩碼數(shù)據(jù)
    detections = self._process_masks_manually(ann, everything_results, width, height)
    
    # 過(guò)濾檢測(cè)結(jié)果
    filtered_detections = self._filter_detections(detections, image_area, min_area_ratio, max_area_ratio)
    
    # 對(duì)象分類
    classified_detections = self._classify_objects(image, filtered_detections)
    
    # 生成標(biāo)注
    return self._generate_annotations(image, classified_detections, output_dir, Path(image_path).stem)

def _process_masks_manually(self, ann, everything_results, img_width, img_height):
    masks_np = ann.cpu().numpy()
    boxes = []
    
    for i in range(num_masks):
        mask = masks_np[i]
        y_indices, x_indices = np.where(mask > 0.5)  # 閾值處理
        
        # 計(jì)算邊界框
        x1 = np.min(x_indices)
        y1 = np.min(y_indices)
        x2 = np.max(x_indices)
        y2 = np.max(y_indices)
        
        boxes.append([x1, y1, x2, y2])

def _classify_objects(self, image, detections):
    for i, bbox in enumerate(detections['boxes']):
        x1, y1, x2, y2 = map(int, bbox)
        object_image = image[y1:y2, x1:x2]
        object_image_resized = cv2.resize(object_image, (224, 224))
        
        # 使用YOLO分類模型
        results = self.classification_model(object_image_resized)
        top1 = results[0].probs.top1
        top1conf = results[0].probs.top1conf.item()

class ManualAnnotationVisualizer:
    def draw_annotations(self, image_path, label_path=None, detections=None, output_path=None):
        # 繪制邊界框和標(biāo)簽
        for i, bbox in enumerate(detections['boxes']):
            color = self.colors[i % len(self.colors)]
            cv2.rectangle(annotated_image, (x1, y1), (x2, y2), color, 2)

class AnnotationValidator:
    def validate_annotations(self, image_path, label_path):
        # 檢查數(shù)值范圍、邊界框有效性、重疊等
        issues = []
        for i, line in enumerate(lines):
            # 驗(yàn)證每個(gè)標(biāo)注行的格式和數(shù)值
            if not (0 <= x_center <= 1):
                issues.append(f"第{i+1}行x_center超出范圍 [0,1]: {x_center}")

python autolabeler.py \
    --input images/ \
    --output dataset/ \
    --model models/FastSAM-x.pt \
    --classification-model models/yolov8n-cls.pt \
    --conf 0.4 \
    --iou 0.9 \
    --visualize \
    --validate

image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff']
for img_path in tqdm(image_files, desc="處理圖片"):
    result = labeler.process_image(str(img_path), args.output, args.conf, args.iou)

min_area_ratio=0.001  # 最小面積比例（相對(duì)于圖像面積）
max_area_ratio=0.95   # 最大面積比例

import torch
import cv2
import numpy as np
from PIL import Image
from ultralytics import FastSAM, YOLO
import os
import json
from pathlib import Path
import argparse
import glob
from tqdm import tqdm
import math


class FastSAMAutoLabeler:
    def __init__(self, model_path='FastSAM-x.pt', device='cuda' if torch.cuda.is_available() else 'cpu', classification_model='yolov8n-cls.pt'):
        """
        FastSAM自動(dòng)標(biāo)注工具（不依賴supervision）
        """
        self.device = device
        self.model = FastSAM(model_path)
        # 加載分類模型
        self.classification_model = YOLO(classification_model)
        print(f"模型加載完成，使用設(shè)備: {device}")

    def process_image(self, image_path, output_dir, conf=0.4, iou=0.9,
                      min_area_ratio=0.001, max_area_ratio=0.95):
        """
        處理單張圖片并生成標(biāo)注
        """
        try:
            # 讀取圖片
            image = cv2.imread(image_path)
            if image is None:
                raise ValueError(f"無(wú)法讀取圖片: {image_path}")

            height, width = image.shape[:2]
            image_area = height * width

            print(f"圖片尺寸: {width}x{height}, 面積: {image_area}")

            # 使用FastSAM進(jìn)行推理
            everything_results = self.model(
                image_path,
                device=self.device,
                retina_masks=True,
                imgsz=1024,
                conf=conf,
                iou=iou
            )

            # 檢查是否有檢測(cè)結(jié)果
            if not everything_results or len(everything_results) == 0:
                print("警告: 未獲得任何檢測(cè)結(jié)果")
                return self._create_empty_result(image, output_dir, Path(image_path).stem)

            # 獲取掩碼數(shù)據(jù)
            masks_data = everything_results[0].masks
            if masks_data is None:
                print("警告: 未檢測(cè)到任何掩碼")
                return self._create_empty_result(image, output_dir, Path(image_path).stem)

            ann = masks_data.data
            print(f"原始掩碼形狀: {ann.shape}")

            # 處理掩碼維度
            if len(ann.shape) == 2:
                ann = ann.unsqueeze(0)
            elif len(ann.shape) != 3:
                raise ValueError(f"不支持的掩碼形狀: {ann.shape}")

            # 手動(dòng)處理檢測(cè)結(jié)果
            detections = self._process_masks_manually(ann, everything_results, width, height)

            # 過(guò)濾檢測(cè)結(jié)果
            filtered_detections = self._filter_detections(
                detections, image_area, min_area_ratio, max_area_ratio
            )

            # 對(duì)每個(gè)檢測(cè)到的對(duì)象進(jìn)行分類
            classified_detections = self._classify_objects(image, filtered_detections)

            return self._generate_annotations(image, classified_detections, output_dir, Path(image_path).stem)

        except Exception as e:
            print(f"處理圖片時(shí)發(fā)生錯(cuò)誤: {e}")
            import traceback
            traceback.print_exc()
            return self._create_error_result(image, output_dir, Path(image_path).stem, str(e))

    def _process_masks_manually(self, ann, everything_results, img_width, img_height):
        """
        手動(dòng)處理掩碼數(shù)據(jù)，替代supervision的功能
        """
        # 將張量轉(zhuǎn)換為numpy數(shù)組
        if isinstance(ann, torch.Tensor):
            masks_np = ann.cpu().numpy()
        else:
            masks_np = ann

        print(f"掩碼numpy數(shù)組形狀: {masks_np.shape}")

        # 確保是三維的 [N, H, W]
        if len(masks_np.shape) == 2:
            masks_np = np.expand_dims(masks_np, axis=0)

        num_masks = masks_np.shape[0]
        print(f"檢測(cè)到 {num_masks} 個(gè)掩碼")

        if num_masks == 0:
            return self._create_empty_detections()

        # 為每個(gè)掩碼計(jì)算邊界框和相關(guān)信息
        boxes = []
        confidences = []
        class_ids = []
        masks = []

        for i in range(num_masks):
            mask = masks_np[i]
            # 找到掩碼中為True的像素位置
            y_indices, x_indices = np.where(mask > 0.5)  # 閾值處理

            if len(x_indices) == 0 or len(y_indices) == 0:
                continue

            # 計(jì)算邊界框
            x1 = np.min(x_indices)
            y1 = np.min(y_indices)
            x2 = np.max(x_indices)
            y2 = np.max(y_indices)

            # 計(jì)算面積和置信度（使用掩碼面積作為置信度參考）
            bbox_area = (x2 - x1) * (y2 - y1)
            mask_area = len(x_indices)
            confidence = min(mask_area / bbox_area, 1.0) if bbox_area > 0 else 0

            boxes.append([x1, y1, x2, y2])
            confidences.append(confidence)
            class_ids.append(0)  # 默認(rèn)類別ID
            masks.append(mask)

        if not boxes:
            return self._create_empty_detections()

        return {
            'boxes': np.array(boxes),
            'confidences': np.array(confidences),
            'class_ids': np.array(class_ids),
            'masks': np.array(masks)
        }

    def _filter_detections(self, detections, image_area, min_area_ratio, max_area_ratio):
        """根據(jù)面積過(guò)濾檢測(cè)結(jié)果"""
        if len(detections['boxes']) == 0:
            return detections

        filtered_boxes = []
        filtered_confidences = []
        filtered_class_ids = []
        filtered_masks = []

        for i, bbox in enumerate(detections['boxes']):
            x1, y1, x2, y2 = bbox
            area = (x2 - x1) * (y2 - y1)
            area_ratio = area / image_area

            if min_area_ratio <= area_ratio <= max_area_ratio:
                filtered_boxes.append(bbox)
                filtered_confidences.append(detections['confidences'][i])
                filtered_class_ids.append(detections['class_ids'][i])
                if i < len(detections['masks']):
                    filtered_masks.append(detections['masks'][i])

        filtered_detections = {
            'boxes': np.array(filtered_boxes) if filtered_boxes else np.empty((0, 4)),
            'confidences': np.array(filtered_confidences) if filtered_confidences else np.empty(0),
            'class_ids': np.array(filtered_class_ids) if filtered_class_ids else np.empty(0),
            'masks': np.array(filtered_masks) if filtered_masks else np.empty(0)
        }

        print(f"過(guò)濾后保留 {len(filtered_boxes)} 個(gè)檢測(cè)結(jié)果")
        return filtered_detections

    def _classify_objects(self, image, detections):
        """
        對(duì)檢測(cè)到的對(duì)象進(jìn)行分類
        """
        if len(detections['boxes']) == 0:
            return detections

        classified_class_ids = []
        classified_confidences = []
        
        for i, bbox in enumerate(detections['boxes']):
            x1, y1, x2, y2 = map(int, bbox)
            
            # 提取對(duì)象區(qū)域
            object_image = image[y1:y2, x1:x2]
            
            if object_image.size == 0:
                classified_class_ids.append(0)
                classified_confidences.append(detections['confidences'][i])
                continue
                
            # 調(diào)整圖像大小以適應(yīng)分類模型
            object_image_resized = cv2.resize(object_image, (224, 224))
            
            # 使用分類模型進(jìn)行預(yù)測(cè)
            try:
                results = self.classification_model(object_image_resized)
                # 獲取最高置信度的類別
                top1 = results[0].probs.top1
                top1conf = results[0].probs.top1conf.item()
                
                classified_class_ids.append(top1)
                classified_confidences.append(top1conf)
            except Exception as e:
                print(f"分類時(shí)出錯(cuò): {e}")
                # 如果分類失敗，保持原始類別
                classified_class_ids.append(detections['class_ids'][i])
                classified_confidences.append(detections['confidences'][i])
        
        # 更新檢測(cè)結(jié)果
        detections['class_ids'] = np.array(classified_class_ids)
        detections['confidences'] = np.array(classified_confidences)
        
        return detections

    def _generate_annotations(self, image, detections, output_dir, image_name):
        """生成YOLO格式標(biāo)注文件"""
        height, width = image.shape[:2]

        # 創(chuàng)建輸出目錄
        os.makedirs(output_dir, exist_ok=True)
        os.makedirs(os.path.join(output_dir, 'images'), exist_ok=True)
        os.makedirs(os.path.join(output_dir, 'labels'), exist_ok=True)

        # 保存圖片
        image_output_path = os.path.join(output_dir, 'images', f'{image_name}.jpg')
        cv2.imwrite(image_output_path, image)

        # 生成YOLO格式標(biāo)注
        yolo_annotations = []
        for i, bbox in enumerate(detections['boxes']):
            x1, y1, x2, y2 = bbox

            # 轉(zhuǎn)換為YOLO格式 (中心點(diǎn)坐標(biāo)和寬高，歸一化)
            x_center = ((x1 + x2) / 2) / width
            y_center = ((y1 + y2) / 2) / height
            w = (x2 - x1) / width
            h = (y2 - y1) / height

            # 邊界檢查
            x_center = max(0, min(1, x_center))
            y_center = max(0, min(1, y_center))
            w = max(0, min(1, w))
            h = max(0, min(1, h))

            # 如果寬高太小則跳過(guò)
            if w < 0.001 or h < 0.001:
                continue

            # 獲取類別ID和置信度
            class_id = int(detections['class_ids'][i]) if i < len(detections['class_ids']) else 0
            confidence = detections['confidences'][i] if i < len(detections['confidences']) else 1.0

            yolo_annotations.append(f"{class_id} {x_center:.6f} {y_center:.6f} {w:.6f} {h:.6f}")

        # 保存YOLO標(biāo)簽
        label_path = os.path.join(output_dir, 'labels', f'{image_name}.txt')
        with open(label_path, 'w') as f:
            f.write('\n'.join(yolo_annotations))

        return {
            'image_path': image_output_path,
            'label_path': label_path,
            'detections_count': len(yolo_annotations),
            'image_name': image_name
        }

    def _create_empty_detections(self):
        """創(chuàng)建空的檢測(cè)結(jié)果"""
        return {
            'boxes': np.empty((0, 4)),
            'confidences': np.empty(0),
            'class_ids': np.empty(0),
            'masks': np.empty(0)
        }

    def _create_empty_result(self, image, output_dir, image_name):
        """創(chuàng)建空結(jié)果"""
        os.makedirs(output_dir, exist_ok=True)
        os.makedirs(os.path.join(output_dir, 'images'), exist_ok=True)
        os.makedirs(os.path.join(output_dir, 'labels'), exist_ok=True)

        image_output_path = os.path.join(output_dir, 'images', f'{image_name}.jpg')
        cv2.imwrite(image_output_path, image)

        label_path = os.path.join(output_dir, 'labels', f'{image_name}.txt')
        with open(label_path, 'w') as f:
            pass

        return {
            'image_path': image_output_path,
            'label_path': label_path,
            'detections_count': 0,
            'image_name': image_name
        }

    def _create_error_result(self, image, output_dir, image_name, error_msg):
        """創(chuàng)建錯(cuò)誤結(jié)果"""
        print(f"為圖片 {image_name} 創(chuàng)建錯(cuò)誤結(jié)果: {error_msg}")
        return self._create_empty_result(image, output_dir, image_name)


class ManualAnnotationVisualizer:
    """
    手動(dòng)實(shí)現(xiàn)的標(biāo)注可視化工具（不依賴supervision）
    """

    def __init__(self, class_names=None, colors=None):
        self.class_names = class_names or ['object']
        self.colors = colors or self._generate_default_colors()

    def _generate_default_colors(self):
        """生成默認(rèn)顏色列表"""
        return [
            (255, 0, 0),  # 紅色
            (0, 255, 0),  # 綠色
            (0, 0, 255),  # 藍(lán)色
            (255, 255, 0),  # 青色
            (255, 0, 255),  # 紫色
            (0, 255, 255),  # 黃色
            (255, 165, 0),  # 橙色
            (128, 0, 128),  # 紫色
            (255, 192, 203),  # 粉色
            (165, 42, 42),  # 棕色
        ]

    def draw_annotations(self, image_path, label_path=None, detections=None,
                         output_path=None, show_labels=True, show_confidences=True):
        """
        繪制標(biāo)注結(jié)果
        """
        image = cv2.imread(image_path)
        if image is None:
            raise ValueError(f"無(wú)法讀取圖片: {image_path}")

        height, width = image.shape[:2]

        if detections is None and label_path is not None:
            # 從YOLO標(biāo)簽文件讀取檢測(cè)結(jié)果
            detections = self._read_yolo_labels(label_path, width, height)
        elif detections is None:
            raise ValueError("必須提供label_path或detections參數(shù)")

        # 繪制邊界框和標(biāo)簽
        annotated_image = self._draw_bounding_boxes(image, detections, show_labels, show_confidences)

        if output_path:
            cv2.imwrite(output_path, annotated_image)
            print(f"可視化結(jié)果已保存: {output_path}")

        return annotated_image

    def _draw_bounding_boxes(self, image, detections, show_labels, show_confidences):
        """繪制邊界框和標(biāo)簽"""
        annotated_image = image.copy()

        for i, bbox in enumerate(detections['boxes']):
            x1, y1, x2, y2 = map(int, bbox)

            # 選擇顏色
            color = self.colors[i % len(self.colors)]

            # 繪制邊界框
            cv2.rectangle(annotated_image, (x1, y1), (x2, y2), color, 2)

            if show_labels:
                # 準(zhǔn)備標(biāo)簽文本
                class_id = int(detections['class_ids'][i]) if i < len(detections['class_ids']) else 0
                class_name = self.class_names[class_id] if class_id < len(self.class_names) else f'class_{class_id}'

                label = class_name
                if show_confidences and i < len(detections['confidences']):
                    confidence = detections['confidences'][i]
                    label += f" {confidence:.2f}"

                # 繪制標(biāo)簽背景
                label_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 2)[0]
                cv2.rectangle(annotated_image,
                              (x1, y1 - label_size[1] - 10),
                              (x1 + label_size[0], y1),
                              color, -1)

                # 繪制標(biāo)簽文本
                cv2.putText(annotated_image, label,
                            (x1, y1 - 5),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)

        return annotated_image

    def _read_yolo_labels(self, label_path, img_width, img_height):
        """讀取YOLO格式標(biāo)簽并轉(zhuǎn)換為檢測(cè)結(jié)果格式"""
        boxes = []
        class_ids = []
        confidences = []

        if not os.path.exists(label_path):
            return self._create_empty_detections()

        with open(label_path, 'r') as f:
            for label_line in f:
                parts = label_line.strip().split()
                if len(parts) < 5:
                    continue

                class_id = int(parts[0])
                x_center, y_center, width, height = map(float, parts[1:5])

                # 轉(zhuǎn)換為絕對(duì)坐標(biāo)
                x_center_abs = x_center * img_width
                y_center_abs = y_center * img_height
                width_abs = width * img_width
                height_abs = height * img_height

                x1 = max(0, x_center_abs - width_abs / 2)
                y1 = max(0, y_center_abs - height_abs / 2)
                x2 = min(img_width, x_center_abs + width_abs / 2)
                y2 = min(img_height, y_center_abs + height_abs / 2)

                boxes.append([x1, y1, x2, y2])
                class_ids.append(class_id)
                confidences.append(1.0)  # YOLO格式?jīng)]有置信度，設(shè)為1.0

        return {
            'boxes': np.array(boxes) if boxes else np.empty((0, 4)),
            'confidences': np.array(confidences) if confidences else np.empty(0),
            'class_ids': np.array(class_ids) if class_ids else np.empty(0),
            'masks': np.empty(0)
        }

    def _create_empty_detections(self):
        """創(chuàng)建空的檢測(cè)結(jié)果"""
        return {
            'boxes': np.empty((0, 4)),
            'confidences': np.empty(0),
            'class_ids': np.empty(0),
            'masks': np.empty(0)
        }


class AnnotationValidator:
    """
    標(biāo)注驗(yàn)證器，用于檢查標(biāo)注質(zhì)量和提供統(tǒng)計(jì)信息
    """
    
    def __init__(self, class_names=None):
        self.class_names = class_names or ['object']
    
    def validate_annotations(self, image_path, label_path):
        """
        驗(yàn)證標(biāo)注文件的質(zhì)量
        """
        # 讀取圖像
        image = cv2.imread(image_path)
        if image is None:
            return {"error": f"無(wú)法讀取圖像: {image_path}"}
            
        height, width = image.shape[:2]
        
        # 讀取標(biāo)注文件
        if not os.path.exists(label_path):
            return {"error": f"標(biāo)注文件不存在: {label_path}"}
            
        with open(label_path, 'r') as f:
            lines = f.readlines()
            
        if not lines:
            return {"warning": "標(biāo)注文件為空"}
            
        issues = []
        class_counts = {}
        boxes = []
        
        for i, line in enumerate(lines):
            parts = line.strip().split()
            if len(parts) < 5:
                issues.append(f"第{i+1}行格式錯(cuò)誤: 需要至少5個(gè)值，實(shí)際得到{len(parts)}個(gè)")
                continue
                
            try:
                class_id = int(parts[0])
                x_center = float(parts[1])
                y_center = float(parts[2])
                w = float(parts[3])
                h = float(parts[4])
                
                # 檢查數(shù)值范圍
                if not (0 <= x_center <= 1):
                    issues.append(f"第{i+1}行x_center超出范圍 [0,1]: {x_center}")
                if not (0 <= y_center <= 1):
                    issues.append(f"第{i+1}行y_center超出范圍 [0,1]: {y_center}")
                if not (0 <= w <= 1):
                    issues.append(f"第{i+1}行width超出范圍 [0,1]: {w}")
                if not (0 <= h <= 1):
                    issues.append(f"第{i+1}行height超出范圍 [0,1]: {h}")
                    
                # 檢查邊界框是否有效
                if w <= 0 or h <= 0:
                    issues.append(f"第{i+1}行邊界框尺寸無(wú)效: width={w}, height={h}")
                    
                # 統(tǒng)計(jì)類別
                class_counts[class_id] = class_counts.get(class_id, 0) + 1
                
                # 轉(zhuǎn)換為像素坐標(biāo)用于重疊檢查
                x1 = max(0, (x_center - w/2) * width)
                y1 = max(0, (y_center - h/2) * height)
                x2 = min(width, (x_center + w/2) * width)
                y2 = min(height, (y_center + h/2) * height)
                boxes.append((x1, y1, x2, y2))
                
            except ValueError as e:
                issues.append(f"第{i+1}行數(shù)值轉(zhuǎn)換錯(cuò)誤: {str(e)}")
                
        # 檢查重疊的邊界框
        overlapping_boxes = self._check_overlapping_boxes(boxes)
        if overlapping_boxes:
            issues.append(f"發(fā)現(xiàn){len(overlapping_boxes)}對(duì)重疊的邊界框")
            
        # 生成報(bào)告
        report = {
            "total_objects": len(lines),
            "class_distribution": class_counts,
            "issues": issues,
            "image_size": (width, height)
        }
        
        if class_counts:
            # 添加類別名稱映射
            class_names_mapping = {}
            for class_id in class_counts:
                if class_id < len(self.class_names):
                    class_names_mapping[class_id] = self.class_names[class_id]
                else:
                    class_names_mapping[class_id] = f"未知類別_{class_id}"
            report["class_names"] = class_names_mapping
            
        return report
    
    def _check_overlapping_boxes(self, boxes, overlap_threshold=0.5):
        """
        檢查重疊的邊界框
        """
        overlapping = []
        for i in range(len(boxes)):
            for j in range(i+1, len(boxes)):
                x1_a, y1_a, x2_a, y2_a = boxes[i]
                x1_b, y1_b, x2_b, y2_b = boxes[j]
                
                # 計(jì)算交集
                x_left = max(x1_a, x1_b)
                y_top = max(y1_a, y1_b)
                x_right = min(x2_a, x2_b)
                y_bottom = min(y2_a, y2_b)
                
                if x_right > x_left and y_bottom > y_top:
                    # 計(jì)算交集面積
                    intersection_area = (x_right - x_left) * (y_bottom - y_top)
                    
                    # 計(jì)算兩個(gè)框的面積
                    area_a = (x2_a - x1_a) * (y2_a - y1_a)
                    area_b = (x2_b - x1_b) * (y2_b - y1_b)
                    
                    # 計(jì)算重疊率
                    overlap = intersection_area / min(area_a, area_b)
                    
                    if overlap > overlap_threshold:
                        overlapping.append((i, j, overlap))
                        
        return overlapping


def main():
    """主函數(shù)示例"""
    parser = argparse.ArgumentParser(description='FastSAM自動(dòng)標(biāo)注工具（無(wú)依賴版）')
    parser.add_argument('--input', type=str,default="images", help='輸入圖片目錄或文件路徑')
    parser.add_argument('--output', type=str, default="dataset", help='輸出目錄')
    parser.add_argument('--model', type=str, default='FastSAM-x.pt', help='FastSAM模型路徑')
    parser.add_argument('--classification-model', type=str, default='yolov8x-cls.pt', help='分類模型路徑')
    parser.add_argument('--conf', type=float, default=0.4, help='置信度閾值')
    parser.add_argument('--iou', type=float, default=0.9, help='IOU閾值')
    parser.add_argument('--min-area', type=float, default=0.001, help='最小面積比例')
    parser.add_argument('--max-area', type=float, default=0.95, help='最大面積比例')
    parser.add_argument('--visualize', action='store_true', help='是否生成可視化結(jié)果')
    parser.add_argument('--validate', action='store_true', help='是否驗(yàn)證標(biāo)注結(jié)果')

    args = parser.parse_args()

    # 創(chuàng)建輸出目錄
    os.makedirs(args.output, exist_ok=True)

    # 初始化標(biāo)注器
    labeler = FastSAMAutoLabeler(args.model, classification_model=args.classification_model)

    # 處理輸入
    if os.path.isfile(args.input):
        # 單文件處理
        result = labeler.process_image(
            args.input, args.output, args.conf, args.iou, args.min_area, args.max_area
        )
        print(f"處理完成: {result}")

        if args.visualize:
            # 獲取分類模型的類別名稱
            class_names = labeler.classification_model.names if hasattr(labeler.classification_model, 'names') else None
            visualizer = ManualAnnotationVisualizer(class_names=class_names)
            vis_path = os.path.join(args.output, 'visualization', f"{Path(args.input).stem}_annotated.jpg")
            os.makedirs(os.path.dirname(vis_path), exist_ok=True)
            visualizer.draw_annotations(
                result['image_path'], result['label_path'], output_path=vis_path
            )

        # 如果需要驗(yàn)證，執(zhí)行驗(yàn)證
        if args.validate:
            class_names = labeler.classification_model.names if hasattr(labeler.classification_model, 'names') else None
            validator = AnnotationValidator(class_names=class_names)
            validation_report = validator.validate_annotations(result['image_path'], result['label_path'])
            
            print("\n標(biāo)注驗(yàn)證報(bào)告:")
            print("=" * 50)
            if "error" in validation_report:
                print(f"錯(cuò)誤: {validation_report['error']}")
            elif "warning" in validation_report:
                print(f"警告: {validation_report['warning']}")
            else:
                print(f"圖像尺寸: {validation_report['image_size'][0]}x{validation_report['image_size'][1]}")
                print(f"總對(duì)象數(shù): {validation_report['total_objects']}")
                print("\n類別分布:")
                for class_id, count in validation_report['class_distribution'].items():
                    class_name = validation_report['class_names'].get(class_id, f"未知類別_{class_id}")
                    print(f"  {class_name} ({class_id}): {count}個(gè)")
                    
                if validation_report['issues']:
                    print(f"\n發(fā)現(xiàn)問(wèn)題 ({len(validation_report['issues'])}個(gè)):")
                    for issue in validation_report['issues']:
                        print(f"  - {issue}")
                else:
                    print("\n標(biāo)注質(zhì)量良好，未發(fā)現(xiàn)問(wèn)題。")

    else:
        # 目錄處理
        image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff']
        image_files = []

        for ext in image_extensions:
            pattern = f'**/*{ext}'
            image_files.extend(Path(args.input).glob(pattern))
            image_files.extend(Path(args.input).glob(pattern.upper()))

        image_files = list(set(image_files))
        print(f"找到 {len(image_files)} 張圖片")

        results = []
        successful = 0
        failed = 0

        for img_path in tqdm(image_files, desc="處理圖片"):
            try:
                result = labeler.process_image(
                    str(img_path), args.output, args.conf, args.iou, args.min_area, args.max_area
                )
                results.append(result)
                successful += 1
                print(f"? 成功處理: {img_path.name} (檢測(cè)到 {result['detections_count']} 個(gè)對(duì)象)")
            except Exception as e:
                failed += 1
                print(f"? 處理失敗: {img_path.name} - 錯(cuò)誤: {e}")

        print(f"\n處理完成!")
        print(f"成功: {successful}, 失敗: {failed}")

        if args.visualize and results:
            # 為前幾張圖片生成可視化結(jié)果
            # 獲取分類模型的類別名稱
            class_names = labeler.classification_model.names if hasattr(labeler.classification_model, 'names') else None
            visualizer = ManualAnnotationVisualizer(class_names=class_names)
            vis_dir = os.path.join(args.output, 'visualization')
            os.makedirs(vis_dir, exist_ok=True)

            sample_count = min(5, len(results))
            print(f"\n為前 {sample_count} 張圖片生成可視化結(jié)果...")

            for i, result in enumerate(results[:sample_count]):
                vis_path = os.path.join(vis_dir, f"{result['image_name']}_annotated.jpg")
                visualizer.draw_annotations(
                    result['image_path'], result['label_path'], output_path=vis_path
                )

        # 如果需要驗(yàn)證，執(zhí)行驗(yàn)證
        if args.validate and results:
            print("\n開始驗(yàn)證標(biāo)注結(jié)果...")
            class_names = labeler.classification_model.names if hasattr(labeler.classification_model, 'names') else None
            validator = AnnotationValidator(class_names=class_names)
            
            total_issues = 0
            for result in results:
                validation_report = validator.validate_annotations(result['image_path'], result['label_path'])
                if "issues" in validation_report and validation_report["issues"]:
                    total_issues += len(validation_report["issues"])
                    print(f"\n{result['image_name']}發(fā)現(xiàn)問(wèn)題:")
                    for issue in validation_report["issues"]:
                        print(f"  - {issue}")
                        
            if total_issues == 0:
                print("所有標(biāo)注文件驗(yàn)證通過(guò)，未發(fā)現(xiàn)問(wèn)題。")
            else:
                print(f"\n總共發(fā)現(xiàn) {total_issues} 個(gè)問(wèn)題。")


if __name__ == "__main__":
    main()