app.py

# -*- coding: utf-8 -*-
"""
🦐 새우 검출 통합 시스템
3개의 앱을 하나로 통합: 자동 검출, 라벨링 도구, 데모
RT-DETR 또는 VIDraft/Shrimp 클라우드 모델 선택 가능
"""
import sys
sys.stdout.reconfigure(encoding='utf-8')

import gradio as gr
from PIL import Image, ImageDraw, ImageFont
import numpy as np
import json
import os
import glob
from datetime import datetime
import torch
from transformers import RTDetrForObjectDetection, RTDetrImageProcessor
import requests
import base64
from io import BytesIO
from inference_sdk import InferenceHTTPClient
import tempfile

# ============================================================
# RT-DETR 및 필터링 함수들 (이전 test_visual_validation에서 통합)
# ============================================================
import cv2

def load_rtdetr_model():
    """RT-DETR 모델 로드"""
    print("🔄 RT-DETR 모델 로딩 중...")
    processor = RTDetrImageProcessor.from_pretrained("PekingU/rtdetr_r50vd_coco_o365")
    model = RTDetrForObjectDetection.from_pretrained("PekingU/rtdetr_r50vd_coco_o365")
    model.eval()
    print("✅ RT-DETR 로딩 완료")
    return processor, model

def detect_with_rtdetr(image, processor, model, confidence=0.3):
    """RT-DETR로 객체 검출"""
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)

    target_sizes = torch.tensor([image.size[::-1]])
    results = processor.post_process_object_detection(
        outputs,
        target_sizes=target_sizes,
        threshold=confidence
    )[0]

    detections = []
    for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
        x1, y1, x2, y2 = box.tolist()
        detections.append({
            'bbox': [x1, y1, x2, y2],
            'confidence': score.item(),
            'label': label.item()
        })

    return detections

def calculate_morphological_features(bbox, image_size):
    """형태학적 특징 계산"""
    x1, y1, x2, y2 = bbox
    width = x2 - x1
    height = y2 - y1

    # Aspect ratio (긴 쪽 / 짧은 쪽)
    aspect_ratio = max(width, height) / max(min(width, height), 1)

    # Area ratio (이미지 대비 면적)
    img_w, img_h = image_size
    area_ratio = (width * height) / (img_w * img_h)

    # Compactness (4π * Area / Perimeter²)
    perimeter = 2 * (width + height)
    compactness = (4 * np.pi * width * height) / max(perimeter ** 2, 1)

    return {
        'aspect_ratio': aspect_ratio,
        'area_ratio': area_ratio,
        'compactness': compactness,
        'width': width,
        'height': height
    }

def calculate_visual_features(image_pil, bbox):
    """시각적 특징 계산 (색상, 텍스처)"""
    # PIL → OpenCV
    image_cv = cv2.cvtColor(np.array(image_pil), cv2.COLOR_RGB2BGR)
    x1, y1, x2, y2 = [int(v) for v in bbox]

    # 바운딩 박스 영역 추출
    roi = image_cv[y1:y2, x1:x2]
    if roi.size == 0:
        return {'hue': 100, 'saturation': 255, 'color_std': 255}

    # HSV 변환
    hsv = cv2.cvtColor(roi, cv2.COLOR_BGR2HSV)

    # 색상 (Hue)
    hue_mean = np.mean(hsv[:, :, 0])

    # 채도 (Saturation)
    saturation = np.mean(hsv[:, :, 1])

    # 색상 일관성 (표준편차)
    color_std = np.std(hsv[:, :, 0])

    return {
        'hue': hue_mean,
        'saturation': saturation,
        'color_std': color_std
    }

def apply_universal_filter(detections, image, threshold=90):
    """범용 새우 필터 적용"""
    img_size = image.size
    filtered = []

    for det in detections:
        bbox = det['bbox']

        # 1. 형태학적 특징
        morph = calculate_morphological_features(bbox, img_size)

        # 2. 시각적 특징
        visual = calculate_visual_features(image, bbox)

        # 3. 점수 계산
        score = 0
        reasons = []

        # Aspect ratio (4:1 ~ 9:1)
        if 4.0 <= morph['aspect_ratio'] <= 9.0:
            score += 25
            reasons.append(f"✓ 종횡비 {morph['aspect_ratio']:.1f}")
        elif 3.0 <= morph['aspect_ratio'] < 4.0 or 9.0 < morph['aspect_ratio'] <= 10.0:
            score += 12
            reasons.append(f"△ 종횡비 {morph['aspect_ratio']:.1f}")
        else:
            score -= 5
            reasons.append(f"✗ 종횡비 {morph['aspect_ratio']:.1f}")

        # Compactness (< 0.50, 긴 형태)
        if morph['compactness'] < 0.40:
            score += 30
            reasons.append(f"✓ 세장도 {morph['compactness']:.2f}")
        elif 0.40 <= morph['compactness'] < 0.50:
            score += 15
            reasons.append(f"△ 세장도 {morph['compactness']:.2f}")
        else:
            reasons.append(f"✗ 세장도 {morph['compactness']:.2f}")
            score -= 20

        # Area
        abs_area = morph['width'] * morph['height']
        if 50000 <= abs_area <= 500000:
            score += 35
            reasons.append(f"✓ 면적 {abs_area/1000:.0f}K")
        elif 500000 < abs_area <= 800000:
            score -= 10
            reasons.append(f"△ 면적 {abs_area/1000:.0f}K")
        elif abs_area > 800000:
            score -= 30
            reasons.append(f"✗ 면적 {abs_area/1000:.0f}K (너무큼)")
        else:
            score -= 10
            reasons.append(f"✗ 면적 {abs_area/1000:.0f}K (너무작음)")

        # Hue (색상)
        hue = visual['hue']
        if hue < 40 or hue > 130:
            score += 10
            reasons.append(f"✓ 색상 {hue:.0f}")
        elif 90 <= hue <= 130:
            score -= 5
            reasons.append(f"✗ 색상 {hue:.0f} (배경)")
        else:
            reasons.append(f"△ 색상 {hue:.0f}")

        # Saturation
        if visual['saturation'] < 85:
            score += 20
            reasons.append(f"✓ 채도 {visual['saturation']:.0f}")
        elif 85 <= visual['saturation'] < 120:
            score += 5
            reasons.append(f"△ 채도 {visual['saturation']:.0f}")
        else:
            score -= 15
            reasons.append(f"✗ 채도 {visual['saturation']:.0f} (높음)")

        # Color consistency
        if visual['color_std'] < 50:
            score += 15
            reasons.append(f"✓ 색상일관성 {visual['color_std']:.1f}")
        elif 50 <= visual['color_std'] < 80:
            score += 5
            reasons.append(f"△ 색상일관성 {visual['color_std']:.1f}")
        else:
            score -= 10
            reasons.append(f"✗ 색상일관성 {visual['color_std']:.1f} (불일치)")

        # RT-DETR confidence
        if 'confidence' in det:
            if det['confidence'] >= 0.3:
                score += 15
                reasons.append(f"✓ 신뢰도 {det['confidence']:.0%}")
            elif det['confidence'] >= 0.1:
                score += 8
                reasons.append(f"△ 신뢰도 {det['confidence']:.0%}")
            else:
                reasons.append(f"✗ 신뢰도 {det['confidence']:.0%}")

        det['filter_score'] = score
        det['filter_reasons'] = reasons
        filtered.append(det)

    return filtered

# YOLOv8 import
# ============================================================
# YOLOv8 모델 설정 (선택적)
# ============================================================
YOLO_MODEL_PATH = "runs/train/yolov8m_shrimp2/weights/best.pt"
yolo_model = None
YOLO_AVAILABLE = False

try:
    from ultralytics import YOLO
    import os
    if os.path.exists(YOLO_MODEL_PATH):
        YOLO_AVAILABLE = True
        print(f"✅ YOLOv8 사용 가능: {YOLO_MODEL_PATH}")
    else:
        print(f"⚠️ YOLOv8 모델 파일 없음: {YOLO_MODEL_PATH}")
except ImportError:
    print("⚠️ ultralytics 패키지 없음 - YOLOv8 비활성화")

def load_yolo_model():
    """YOLOv8 모델 로딩"""
    global yolo_model
    if not YOLO_AVAILABLE:
        raise Exception("YOLOv8 모델을 사용할 수 없습니다")
    if yolo_model is None:
        print(f"🔄 YOLOv8 모델 로딩 중: {YOLO_MODEL_PATH}")
        yolo_model = YOLO(YOLO_MODEL_PATH)
        print("✅ YOLOv8 모델 로딩 완료")
    return yolo_model

def detect_with_yolo(image, confidence=0.1):
    """YOLOv8 모델로 검출"""
    try:
        model = load_yolo_model()

        # 추론 실행
        results = model.predict(
            source=image,
            conf=confidence,
            verbose=False
        )

        detections = []
        for result in results:
            boxes = result.boxes
            for box in boxes:
                x1, y1, x2, y2 = box.xyxy[0].tolist()
                conf = box.conf[0].item()

                detections.append({
                    'bbox': [x1, y1, x2, y2],
                    'confidence': conf
                })

        print(f"✅ YOLOv8 검출 완료: {len(detections)}개")
        return detections

    except Exception as e:
        print(f"❌ YOLOv8 검출 오류: {str(e)}")
        import traceback
        traceback.print_exc()
        return []

# ============================================================
# Roboflow SDK 설정 (최적화된 방식)
# ============================================================
# 환경변수에서 API 키 읽기 (Hugging Face Secrets 사용)
ROBOFLOW_API_KEY = os.getenv("ROBOFLOW_API_KEY", "")

# Roboflow Inference SDK 클라이언트 (connection pooling 지원)
roboflow_client = InferenceHTTPClient(
    api_url="https://serverless.roboflow.com",
    api_key=ROBOFLOW_API_KEY
)

def detect_with_roboflow(image, confidence=0.065):
    """Roboflow API를 사용한 최적화된 검출 (로컬 테스트와 동일)"""
    try:
        # 원본 이미지 보존
        image_original = image
        original_size = image_original.size

        # 리사이즈 (API 전송용)
        image_resized = image_original.copy()
        image_resized.thumbnail((640, 640), Image.Resampling.LANCZOS)
        print(f"📐 이미지 리사이즈: {original_size} → {image_resized.size}")

        # Base64 인코딩
        buffered = BytesIO()
        image_resized.save(buffered, format="JPEG", quality=80)
        img_base64 = base64.b64encode(buffered.getvalue()).decode()
        print(f"📦 Base64 크기: {len(img_base64)} bytes")

        print(f"🔄 Roboflow API 추론 시작...")

        # 🚀 최적화 3: requests로 API 호출 (SDK 대신 사용 - 더 안정적)
        response = requests.post(
            'https://serverless.roboflow.com/vidraft/workflows/find-shrimp-6',
            headers={'Content-Type': 'application/json'},
            json={
                'api_key': ROBOFLOW_API_KEY,
                'inputs': {
                    'image': {'type': 'base64', 'value': img_base64}
                }
            },
            timeout=30
        )

        if response.status_code != 200:
            print(f"❌ Roboflow API 오류: {response.status_code}")
            print(f"응답: {response.text}")
            return []

        result = response.json()
        print(f"🔍 Roboflow 응답: {json.dumps(result, indent=2, ensure_ascii=False)[:500]}...")

        # Workflow 응답 구조 파싱
        detections = []
        predictions = []

        # 방법 1: outputs[0].predictions.predictions (workflow 형태)
        if isinstance(result, dict) and 'outputs' in result and len(result['outputs']) > 0:
            output = result['outputs'][0]
            if isinstance(output, dict) and 'predictions' in output:
                pred_data = output['predictions']
                # predictions가 dict이고 그 안에 predictions 배열이 있는 경우
                if isinstance(pred_data, dict) and 'predictions' in pred_data:
                    predictions = pred_data['predictions']
                # predictions가 바로 배열인 경우
                elif isinstance(pred_data, list):
                    predictions = pred_data
                else:
                    predictions = [pred_data]

        # 방법 2: 직접 predictions
        elif isinstance(result, dict) and 'predictions' in result:
            predictions = result['predictions']

        # 방법 3: 다른 구조
        elif isinstance(result, list):
            predictions = result

        print(f"📦 찾은 predictions: {len(predictions)}개")

        # 스케일 계산 (리사이즈된 좌표 → 원본 좌표)
        scale_x = original_size[0] / image_resized.size[0]
        scale_y = original_size[1] / image_resized.size[1]
        print(f"📏 스케일: x={scale_x:.2f}, y={scale_y:.2f}")

        for pred in predictions:
            # 클래스 필터링 (shrimp만 검출)
            pred_class = pred.get('class', '')
            if pred_class != 'shrimp':
                continue

            # 신뢰도 필터링
            pred_confidence = pred.get('confidence', 0)
            if pred_confidence < confidence:
                continue

            # 바운딩 박스 추출 (리사이즈된 좌표)
            x = pred.get('x', 0)
            y = pred.get('y', 0)
            width = pred.get('width', 0)
            height = pred.get('height', 0)

            # 원본 크기로 스케일 변환
            x_scaled = x * scale_x
            y_scaled = y * scale_y
            width_scaled = width * scale_x
            height_scaled = height * scale_y

            # 중심점 좌표를 좌상단/우하단 좌표로 변환
            x1 = x_scaled - width_scaled / 2
            y1 = y_scaled - height_scaled / 2
            x2 = x_scaled + width_scaled / 2
            y2 = y_scaled + height_scaled / 2

            detections.append({
                'bbox': [x1, y1, x2, y2],
                'confidence': pred_confidence
            })
            print(f"  ✓ 검출 (shrimp): conf={pred_confidence:.2%}, bbox=[{x1:.0f},{y1:.0f},{x2:.0f},{y2:.0f}]")

        print(f"✅ Roboflow 검출 완료: {len(detections)}개")
        return detections

    except Exception as e:
        print(f"❌ Roboflow SDK 오류: {str(e)}")
        import traceback
        traceback.print_exc()
        return []

# ============================================================
# 전역 모델 변수 (지연 로딩)
# ============================================================
processor = None
model = None

def load_rtdetr_on_demand():
    """RT-DETR 모델을 필요시에만 로딩"""
    global processor, model
    if processor is None or model is None:
        processor, model = load_rtdetr_model()
        return "✅ RT-DETR 모델 로딩 완료"
    else:
        return "ℹ️ RT-DETR 모델이 이미 로딩되어 있습니다"

print("✅ VIDraft/Shrimp 클라우드 모델 사용 가능\n")

# ============================================================
# 라벨링 도구 전역 변수
# ============================================================
current_data = {
    'folder': None,
    'images': [],
    'current_idx': 0,
    'detections': {},
    'selections': {},
    'confidence_threshold': 0.2,
    'image_cache': {},
    'model_type': 'RT-DETR'  # 현재 선택된 모델
}

GROUND_TRUTH_FILE = "ground_truth.json"
DATA_BASE = "data/흰다리새우 실측 데이터_익투스에이아이(주)"

# ============================================================
# 모델별 검출 함수
# ============================================================

def detect_with_selected_model(image, confidence, model_type):
    """선택된 모델로 검출"""
    if model_type == "RT-DETR":
        if processor is None or model is None:
            raise ValueError("⚠️ RT-DETR 모델이 로딩되지 않았습니다. '🔄 RT-DETR 로드' 버튼을 먼저 클릭하세요.")
        return detect_with_rtdetr(image, processor, model, confidence)
    elif model_type == "VIDraft/Shrimp":
        return detect_with_roboflow(image, confidence)
    elif model_type == "YOLOv8":
        return detect_with_yolo(image, confidence)
    else:
        return []

# ============================================================
# 탭 1: 자동 검출 (Interactive Validation)
# ============================================================

def interactive_detect(image, confidence, filter_threshold, show_all, model_type, use_filter):
    """대화형 검출"""
    if image is None:
        return None, "⚠️ 이미지를 업로드하세요."

    try:
        # 선택된 모델로 검출
        all_detections = detect_with_selected_model(image, confidence, model_type)

        # 필터 적용 여부에 따라 처리
        if not use_filter:
            # 필터 미사용: 신뢰도만 적용
            filtered_detections = all_detections
            for det in filtered_detections:
                det['filter_score'] = det['confidence'] * 100
                det['filter_reasons'] = [f"신뢰도: {det['confidence']:.0%} (필터 미사용)"]
            all_detections_scored = filtered_detections
        else:
            # 필터 사용
            if model_type in ["VIDraft/Shrimp", "YOLOv8"]:
                # Roboflow & YOLOv8: 신뢰도를 필터 점수로 사용
                for det in all_detections:
                    det['filter_score'] = det['confidence'] * 100
                    det['filter_reasons'] = [f"{model_type} 신뢰도: {det['confidence']:.0%}"]
                all_detections_scored = all_detections
            else:
                # RT-DETR: Universal Filter 사용
                all_detections_scored = apply_universal_filter(all_detections, image, threshold=0)

            # 필터 임계값 적용
            filtered_detections = [det for det in all_detections_scored if det['filter_score'] >= filter_threshold]

        # 시각화
        img = image.copy()
        draw = ImageDraw.Draw(img)

        try:
            font = ImageFont.truetype("arial.ttf", 14)
            font_large = ImageFont.truetype("arial.ttf", 18)
            font_small = ImageFont.truetype("arial.ttf", 10)
        except:
            font = ImageFont.load_default()
            font_large = ImageFont.load_default()
            font_small = ImageFont.load_default()

        # 제거된 객체 먼저 표시 (빨간색)
        rejected_detections = [det for det in all_detections_scored if det['filter_score'] < filter_threshold]
        for idx, det in enumerate(rejected_detections, 1):
            x1, y1, x2, y2 = det['bbox']
            score = det['filter_score']

            # 빨간색 박스 (제거됨)
            draw.rectangle([x1, y1, x2, y2], outline="red", width=8)

            # 라벨 (작게)
            label = f"✗{idx} {score:.0f}점"
            bbox = draw.textbbox((x1, y1 - 20), label, font=font_small)
            draw.rectangle(bbox, fill="red")
            draw.text((x1, y1 - 20), label, fill="white", font=font_small)

        # 전체 검출 표시 (옵션) - 회색
        if show_all:
            for det in all_detections_scored:
                if det not in filtered_detections and det not in rejected_detections:
                    x1, y1, x2, y2 = det['bbox']
                    draw.rectangle([x1, y1, x2, y2], outline="gray", width=4)

        # 필터링된 결과 (통과) - 녹색/노란색/주황색
        for idx, det in enumerate(filtered_detections, 1):
            x1, y1, x2, y2 = det['bbox']
            score = det['filter_score']

            # 점수에 따라 색상
            if score >= 75:
                color = "lime"
            elif score >= 50:
                color = "yellow"
            else:
                color = "orange"

            # 박스 (두껍게)
            draw.rectangle([x1, y1, x2, y2], outline=color, width=10)

            # 라벨
            label = f"✓#{idx} {score:.0f}점"
            bbox = draw.textbbox((x1, y1 - 25), label, font=font)
            draw.rectangle(bbox, fill=color)
            draw.text((x1, y1 - 25), label, fill="black", font=font)

            # 세부 정보 (작게)
            details = f"{model_type}:{det['confidence']:.0%}"
            draw.text((x1, y2 + 5), details, fill=color, font=font_small)

        # 헤더
        header = f"[{model_type}] ✓ {len(filtered_detections)}개 / ✗ {len(rejected_detections)}개 (전체: {len(all_detections_scored)}개)"
        header_bbox = draw.textbbox((10, 10), header, font=font_large)
        draw.rectangle([5, 5, header_bbox[2]+10, header_bbox[3]+10],
                      fill="black", outline="lime", width=2)
        draw.text((10, 10), header, fill="lime", font=font_large)

        # 정보 생성
        info = f"""
### 📊 검출 결과 (모델: {model_type})

- **전체 검출**: {len(all_detections_scored)}개
- **필터링 후**: {len(filtered_detections)}개
- **제거됨**: {len(rejected_detections)}개

---

### 🎯 검출된 객체 상세 (✅ 통과)

"""

        for idx, det in enumerate(filtered_detections, 1):
            info += f"""
**#{idx} - 점수: {det['filter_score']:.0f}점** ({model_type} 신뢰도: {det['confidence']:.0%})

"""
            # 주요 특징만 5개
            for reason in det['filter_reasons'][:5]:
                info += f"- {reason}\n"

        if not filtered_detections:
            info += """
⚠️ **검출된 객체가 없습니다.**

"""

        # 제거된 객체 정보 추가
        if rejected_detections:
            info += f"""

---

### ❌ 제거된 객체 ({len(rejected_detections)}개)

"""
            for idx, det in enumerate(rejected_detections[:3], 1):  # 최대 3개만 표시
                info += f"""
**제거 #{idx} - 점수: {det['filter_score']:.0f}점** (임계값 미달)
- {model_type} 신뢰도: {det['confidence']:.0%}

"""
                # 실패 이유 표시
                for reason in det['filter_reasons'][:3]:
                    info += f"- {reason}\n"

        return img, info

    except Exception as e:
        import traceback
        error_detail = traceback.format_exc()
        return None, f"❌ 오류 발생:\n\n```\n{error_detail}\n```"


# ============================================================
# 탭 2: 라벨링 도구 (Labeling Tool)
# ============================================================

def detect_with_rtdetr_fast(image, confidence=0.3):
    """RT-DETR 빠른 검출"""
    inputs = processor(images=image, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)

    target_sizes = torch.tensor([image.size[::-1]])
    results = processor.post_process_object_detection(
        outputs,
        target_sizes=target_sizes,
        threshold=confidence
    )[0]

    detections = []
    for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
        x1, y1, x2, y2 = box.tolist()
        detections.append({
            'bbox': [x1, y1, x2, y2],
            'confidence': score.item()
        })

    return detections


def load_existing_ground_truth():
    """기존 ground_truth.json 로드"""
    if os.path.exists(GROUND_TRUTH_FILE):
        with open(GROUND_TRUTH_FILE, 'r', encoding='utf-8') as f:
            return json.load(f)
    return {}


def save_ground_truth(data):
    """ground_truth.json 저장"""
    backup_dir = "backups"
    if not os.path.exists(backup_dir):
        os.makedirs(backup_dir)

    if os.path.exists(GROUND_TRUTH_FILE):
        backup_name = f"ground_truth_backup_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
        backup_path = os.path.join(backup_dir, backup_name)
        import shutil
        shutil.copy2(GROUND_TRUTH_FILE, backup_path)

    with open(GROUND_TRUTH_FILE, 'w', encoding='utf-8') as f:
        json.dump(data, f, ensure_ascii=False, indent=2)

    print(f"✅ Ground Truth 저장 완료: {len(data)}개 이미지")


def get_folders():
    """사용 가능한 폴더 목록"""
    folders = sorted(glob.glob(os.path.join(DATA_BASE, "2*")))
    return [os.path.basename(f) for f in folders if os.path.isdir(f)]


def start_labeling(folder, conf_threshold, model_type):
    """라벨링 시작"""
    if not folder:
        return None, "❌ 폴더를 선택하세요.", ""

    current_data['folder'] = folder
    current_data['confidence_threshold'] = conf_threshold
    current_data['model_type'] = model_type

    folder_path = os.path.join(DATA_BASE, folder)
    all_images = sorted(glob.glob(os.path.join(folder_path, "*.jpg")))

    # -1, -2 등이 붙은 파일 제외 (예: 251017_01-1.jpg 제외, 251017_01.jpg만 포함)
    import re
    images = [img for img in all_images if not re.search(r'-\d+\.jpg$', os.path.basename(img))]

    if not images:
        return None, "❌ 이미지 없음", ""

    print(f"📁 폴더: {folder}")
    print(f"   전체 이미지: {len(all_images)}개")
    print(f"   라벨링 대상: {len(images)}개 (-숫자 파일 제외)")

    current_data['images'] = images
    current_data['current_idx'] = 0
    current_data['detections'] = {}
    current_data['selections'] = {}

    # 기존 GT 로드
    gt = load_existing_ground_truth()

    # 이미 라벨링된 이미지 건너뛰기
    for i, img_path in enumerate(images):
        filename = os.path.basename(img_path)
        if filename in gt:
            current_data['selections'][filename] = [j for j in range(len(gt[filename]))]
            print(f"⏭️  건너뛰기: {filename} (이미 라벨링됨)")

    # 첫 미라벨링 이미지 찾기
    while current_data['current_idx'] < len(images):
        filename = os.path.basename(images[current_data['current_idx']])
        if filename not in current_data['selections']:
            break
        current_data['current_idx'] += 1

    if current_data['current_idx'] >= len(images):
        return None, "✅ 모든 이미지 라벨링 완료!", ""

    return show_current_image()


def show_current_image():
    """현재 이미지 표시"""
    if current_data['current_idx'] >= len(current_data['images']):
        return None, "✅ 완료!", ""

    img_path = current_data['images'][current_data['current_idx']]
    filename = os.path.basename(img_path)

    # 캐시 확인
    if filename in current_data['image_cache']:
        image = current_data['image_cache'][filename]
    else:
        image = Image.open(img_path)
        current_data['image_cache'][filename] = image

    # 선택된 모델로 검출
    if filename not in current_data['detections']:
        if current_data['model_type'] == 'RT-DETR':
            detections = detect_with_rtdetr_fast(image, current_data['confidence_threshold'])
        elif current_data['model_type'] == 'YOLOv8':
            detections = detect_with_yolo(image, current_data['confidence_threshold'])
        else:  # VIDraft/Shrimp
            detections = detect_with_roboflow(image, current_data['confidence_threshold'])
        current_data['detections'][filename] = detections
    else:
        detections = current_data['detections'][filename]

    # 선택된 박스
    selected_indices = current_data['selections'].get(filename, [])

    # 시각화
    vis_image = draw_detections(image, detections, selected_indices)

    info = f"""
### 📁 {current_data['folder']} - 이미지 {current_data['current_idx']+1}/{len(current_data['images'])}

**파일**: {filename}
**모델**: {current_data['model_type']}

**검출**: {len(detections)}개
**선택**: {len(selected_indices)}개

---

### 🖱️ 사용 방법:
1. 이미지를 클릭하여 박스 선택/해제
2. "다음" 버튼으로 저장 후 이동
3. "건너뛰기"로 선택 없이 이동
"""

    return vis_image, info, filename


def draw_detections(image, detections, selected_indices):
    """검출 결과 그리기"""
    img = image.copy()
    draw = ImageDraw.Draw(img)

    try:
        font_tiny = ImageFont.truetype("arial.ttf", 10)
        font_large = ImageFont.truetype("arial.ttf", 40)
    except:
        font_tiny = ImageFont.load_default()
        font_large = ImageFont.load_default()

    # 선택되지 않은 박스 먼저 (뒤쪽 레이어)
    for idx, det in enumerate(detections):
        if idx not in selected_indices:
            x1, y1, x2, y2 = det['bbox']
            draw.rectangle([x1, y1, x2, y2], outline="lime", width=20)
            corner_label = f"#{idx+1}"
            draw.rectangle([x1-2, y1-24, x1+30, y1-2], fill="lime")
            draw.text((x1, y1 - 22), corner_label, fill="white", font=font_tiny)

    # 선택된 박스 나중에 (앞쪽 레이어)
    for idx, det in enumerate(detections):
        if idx in selected_indices:
            x1, y1, x2, y2 = det['bbox']
            draw.rectangle([x1, y1, x2, y2], outline="blue", width=28)
            corner_label = f"✓#{idx+1}"
            draw.rectangle([x1-2, y1-24, x1+40, y1-2], fill="blue")
            draw.text((x1, y1 - 22), corner_label, fill="white", font=font_tiny)

    # 원형 버튼
    for idx, det in enumerate(detections):
        x1, y1, x2, y2 = det['bbox']
        center_x = (x1 + x2) / 2
        center_y = (y1 + y2) / 2

        selected = idx in selected_indices
        btn_color = "blue" if selected else "lime"
        btn_text = f"✓{idx+1}" if selected else f"{idx+1}"

        box_width = x2 - x1
        box_height = y2 - y1
        radius = min(55, box_width * 0.18, box_height * 0.35)

        # 원형 버튼
        draw.ellipse(
            [center_x - radius, center_y - radius,
             center_x + radius, center_y + radius],
            fill=btn_color, outline="white", width=4
        )
        draw.text((center_x - radius*0.5, center_y - radius*0.6),
                 btn_text, fill="white", font=font_large)

    return img


def labeling_click(image, filename, evt: gr.SelectData):
    """이미지 클릭 이벤트"""
    if not filename or filename not in current_data['detections']:
        return image, "⚠️ 이미지를 먼저 로드하세요."

    click_x, click_y = evt.index[0], evt.index[1]
    detections = current_data['detections'][filename]
    selected_indices = set(current_data['selections'].get(filename, []))

    # 클릭한 박스 찾기
    clicked_idx = None
    button_candidates = []

    # 버튼 영역 확인
    for idx, det in enumerate(detections):
        x1, y1, x2, y2 = det['bbox']
        center_x = (x1 + x2) / 2
        center_y = (y1 + y2) / 2

        box_width = x2 - x1
        box_height = y2 - y1
        radius = min(55, box_width * 0.18, box_height * 0.35)

        distance = ((click_x - center_x) ** 2 + (click_y - center_y) ** 2) ** 0.5

        if distance <= radius:
            button_candidates.append((idx, distance))

    # 버튼 클릭이 있으면 선택
    if button_candidates:
        button_candidates.sort(key=lambda x: x[1])
        clicked_idx = button_candidates[0][0]
    else:
        # 박스 영역 클릭 확인
        for idx, det in enumerate(detections):
            x1, y1, x2, y2 = det['bbox']
            if x1 <= click_x <= x2 and y1 <= click_y <= y2:
                clicked_idx = idx
                break

    # 선택 토글
    if clicked_idx is not None:
        if clicked_idx in selected_indices:
            selected_indices.remove(clicked_idx)
            print(f"❌ 선택 해제: 박스 #{clicked_idx+1}")
        else:
            selected_indices.add(clicked_idx)
            print(f"✅ 선택: 박스 #{clicked_idx+1}")

        current_data['selections'][filename] = list(selected_indices)

        # 이미지 다시 그리기
        img_path = current_data['images'][current_data['current_idx']]
        image = Image.open(img_path)
        vis_image = draw_detections(image, detections, list(selected_indices))

        info = f"✅ 박스 #{clicked_idx+1} {'선택' if clicked_idx in selected_indices else '해제'}"
        return vis_image, info

    return image, "❌ 박스를 찾을 수 없습니다."


def save_and_next():
    """저장 후 다음"""
    if current_data['current_idx'] >= len(current_data['images']):
        return None, "✅ 완료!", ""

    img_path = current_data['images'][current_data['current_idx']]
    filename = os.path.basename(img_path)

    # GT 저장
    gt = load_existing_ground_truth()
    selected_indices = current_data['selections'].get(filename, [])

    if selected_indices:
        detections = current_data['detections'][filename]
        gt[filename] = [
            {
                'bbox': detections[i]['bbox'],
                'folder': current_data['folder']
            }
            for i in selected_indices
        ]
        save_ground_truth(gt)
        print(f"💾 저장: {filename} - {len(selected_indices)}개 박스")
    else:
        print(f"⏭️  건너뛰기: {filename} - 선택 없음")

    # 다음 이미지
    current_data['current_idx'] += 1

    # 다음 미라벨링 이미지 찾기
    while current_data['current_idx'] < len(current_data['images']):
        next_filename = os.path.basename(current_data['images'][current_data['current_idx']])
        if next_filename not in current_data['selections']:
            break
        current_data['current_idx'] += 1

    if current_data['current_idx'] >= len(current_data['images']):
        return None, "✅ 모든 이미지 라벨링 완료!", ""

    return show_current_image()


def skip_image():
    """건너뛰기"""
    current_data['current_idx'] += 1

    if current_data['current_idx'] >= len(current_data['images']):
        return None, "✅ 완료!", ""

    return show_current_image()


# ============================================================
# Gradio 인터페이스 - 2개 탭으로 통합
# ============================================================

with gr.Blocks(title="🦐 새우 검출 통합 시스템", theme=gr.themes.Soft()) as demo:

    gr.Markdown("""
    # 🦐 새우 검출 통합 시스템

    **2가지 탭으로 새우를 정확하게 검출하고 관리하세요**

    ---
    """)

    # ==================== 최상단: 모델 선택 ====================
    with gr.Row():
        with gr.Column(scale=3):
            # YOLOv8 사용 가능 여부에 따라 선택지 결정
            model_choices = ["RT-DETR", "VIDraft/Shrimp"]
            default_model = "RT-DETR"
            if YOLO_AVAILABLE:
                model_choices.append("YOLOv8")
                default_model = "YOLOv8"

            model_selector = gr.Radio(
                choices=model_choices,
                value=default_model,
                label="🤖 검출 모델 선택",
                info="모든 탭에 적용됩니다"
            )
        with gr.Column(scale=1):
            load_rtdetr_btn = gr.Button("🔄 RT-DETR 로드", size="sm", variant="secondary")
            rtdetr_status = gr.Textbox(label="모델 상태", value="⏸️ RT-DETR 미로드 (VIDraft/Shrimp 클라우드 모델 사용 가능)", interactive=False, lines=1)

    # RT-DETR 로딩 버튼 이벤트
    load_rtdetr_btn.click(
        load_rtdetr_on_demand,
        inputs=[],
        outputs=[rtdetr_status]
    )

    gr.Markdown("---")

    with gr.Tabs():
        # ==================== 탭 1: 자동 검출 ====================
        with gr.TabItem("🤖 자동 검출 & 검증"):
            gr.Markdown("""
            ### 실시간으로 파라미터를 조정하며 검출 결과를 확인
            최적화된 파라미터로 새우 검출을 테스트하세요.
            """)

            with gr.Row():
                with gr.Column():
                    input_image_detect = gr.Image(label="입력 이미지", type="pil")

                    confidence_slider_detect = gr.Slider(
                        0.01, 1.0, 0.1,
                        step=0.01,
                        label="신뢰도 임계값",
                        info="RT-DETR: 0.065 | VIDraft/Shrimp: 0.3~0.5 | YOLOv8: 0.1~0.3 권장"
                    )

                    use_filter_check = gr.Checkbox(
                        label="🔍 필터 점수 임계값 사용",
                        value=False,
                        info="체크하면 필터 점수 기준으로 추가 필터링"
                    )

                    filter_slider_detect = gr.Slider(
                        0, 100, 90,
                        step=5,
                        label="필터 점수 임계값",
                        info="RT-DETR: Universal Filter | VIDraft/Shrimp: 신뢰도 기반",
                        visible=True
                    )

                    show_all_check = gr.Checkbox(
                        label="전체 검출 결과 표시 (회색)",
                        value=False
                    )

                    detect_btn = gr.Button("🚀 검출 실행", variant="primary", size="lg")

                    # 예제 이미지 (결과 파일 제외)
                    example_images = [
                        "examples/250818_03.jpg",
                        "examples/test_shrimp_tank.png",
                        "examples/250818_05.jpg",
                    ]

                    # 파일이 존재하는 것만 필터링
                    example_images = [img for img in example_images if os.path.exists(img)]

                    if example_images:
                        gr.Examples(
                            examples=[[img] for img in example_images],
                            inputs=[input_image_detect],
                            label="📷 예제 이미지"
                        )

                with gr.Column():
                    output_image_detect = gr.Image(label="검출 결과")
                    output_info_detect = gr.Markdown()

            detect_btn.click(
                interactive_detect,
                [input_image_detect, confidence_slider_detect, filter_slider_detect, show_all_check, model_selector, use_filter_check],
                [output_image_detect, output_info_detect]
            )

            # 필터 사용 체크박스에 따라 필터 슬라이더 활성화/비활성화
            def update_filter_interactivity(use_filter):
                return gr.update(interactive=use_filter)

            use_filter_check.change(
                update_filter_interactivity,
                inputs=[use_filter_check],
                outputs=[filter_slider_detect]
            )

            gr.Markdown("""
            ### 💡 사용 팁
            - 모델을 선택하고 신뢰도를 조정하여 검출 결과를 확인하세요
            - 검출이 적을 때는 신뢰도를 낮추고, 오검출이 많을 때는 높이세요
            - 필터 기능을 사용하여 더 정확한 결과를 얻을 수 있습니다

            **박스 색상:** 🟢 녹색(높은 확률) | 🟡 노란색(중간 확률) | 🟠 주황색(낮은 확률) | 🔴 빨간색(제거됨)
            """)

        # ==================== 탭 2: 라벨링 도구 ====================
        with gr.TabItem("📝 Ground Truth 라벨링"):
            gr.Markdown("""
            ### 선택된 모델의 검출 결과에서 올바른 박스만 선택하여 라벨링
            이미지를 클릭하여 새우 박스를 선택/해제하세요.
            """)

            with gr.Row():
                with gr.Column(scale=1):
                    folder_dropdown = gr.Dropdown(
                        choices=get_folders(),
                        label="📁 폴더 선택",
                        info="라벨링할 폴더를 선택하세요"
                    )

                    conf_slider_label = gr.Slider(
                        0.01, 0.5, 0.2,
                        step=0.05,
                        label="신뢰도",
                        info="검출 민감도 조정"
                    )

                    start_btn = gr.Button("▶️ 라벨링 시작", variant="primary", size="lg")

                    gr.Markdown("---")

                    next_btn = gr.Button("⏭️ 저장 & 다음", variant="secondary", size="lg")
                    skip_btn = gr.Button("⏩ 건너뛰기", size="lg")

                    labeling_info = gr.Markdown("폴더를 선택하고 '라벨링 시작'을 클릭하세요.")

                with gr.Column(scale=2):
                    labeling_image = gr.Image(
                        label="🖱️ 클릭하여 박스 선택/해제",
                        type="pil",
                        interactive=True
                    )

                    labeling_filename = gr.Textbox(visible=False)
                    click_info = gr.Markdown()

            # 이벤트 핸들러
            start_btn.click(
                start_labeling,
                [folder_dropdown, conf_slider_label, model_selector],
                [labeling_image, labeling_info, labeling_filename]
            )

            labeling_image.select(
                labeling_click,
                [labeling_image, labeling_filename],
                [labeling_image, click_info]
            )

            next_btn.click(
                save_and_next,
                [],
                [labeling_image, labeling_info, labeling_filename]
            )

            skip_btn.click(
                skip_image,
                [],
                [labeling_image, labeling_info, labeling_filename]
            )

            gr.Markdown("""
            ### 🖱️ 사용 방법
            1. **모델 선택** (최상단에서 선택)
            2. 폴더 선택 후 "라벨링 시작"
            3. 이미지에서 **원형 버튼 클릭** 또는 **박스 영역 클릭**으로 선택/해제
            4. "저장 & 다음"으로 다음 이미지로 이동 (자동 저장)
            5. "건너뛰기"로 선택 없이 다음 이미지로

            **💾 저장 위치:** `ground_truth.json` (자동 백업: `backups/`)
            """)

    gr.Markdown("""
    ---

    ### 🤖 모델 설명
    - **RT-DETR**: 로컬 모델, 빠른 추론 속도, 오프라인 사용 가능
    - **VIDraft/Shrimp**: 클라우드 모델, 인터넷 연결 필요
    - **YOLOv8**: 로컬 커스텀 학습 모델, 빠른 추론 속도

    ---

    © 2025 VIDraft. All rights reserved.
    """)

if __name__ == "__main__":
    print("\n" + "="*60)
    print("🦐 새우 검출 통합 시스템 v2.2 시작")
    print("="*60)
    print("🤖 사용 가능한 모델:")
    print("   1. RT-DETR (로컬)")
    print("   2. VIDraft/Shrimp (클라우드)")
    print("   3. YOLOv8 (로컬 학습) ⭐ 기본값")
    print(f"\n📦 YOLOv8 모델: {YOLO_MODEL_PATH}")
    print("="*60)

    demo.launch(
        server_name="0.0.0.0",
        server_port=None,  # 자동으로 빈 포트 찾기
        share=False
    )