app.py

import os
import cv2
import tempfile
import spaces
import gradio as gr
import numpy as np
import torch
import matplotlib
import matplotlib.pyplot as plt
import pandas as pd
from PIL import Image, ImageDraw
from transformers import (
    Sam3Model, Sam3Processor,
    Sam3TrackerModel, Sam3TrackerProcessor
)

# ============ THEME SETUP ============
# Theme disabled to avoid API schema issues on HF Spaces

# ============ GLOBAL SETUP ============
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"🖥️ Using compute device: {device}")

# Load models
print("⏳ Loading SAM3 Models permanently into memory...")
try:
    # 오프라인 모드로 캐시에서 로드 시도
    print("   ... Loading from local cache (offline mode)")
    IMG_MODEL = Sam3Model.from_pretrained("DiffusionWave/sam3", local_files_only=True, device_map="cpu", torch_dtype=torch.float32)
    IMG_PROCESSOR = Sam3Processor.from_pretrained("DiffusionWave/sam3", local_files_only=True)

    TRK_MODEL = Sam3TrackerModel.from_pretrained("DiffusionWave/sam3", local_files_only=True, device_map="cpu", torch_dtype=torch.float32)
    TRK_PROCESSOR = Sam3TrackerProcessor.from_pretrained("DiffusionWave/sam3", local_files_only=True)

    print("✅ All Models loaded successfully from local cache!")
except Exception as e:
    print(f"❌ Cache loading failed: {e}")
    print("   Trying online loading...")
    try:
        IMG_MODEL = Sam3Model.from_pretrained("DiffusionWave/sam3", device_map="cpu", torch_dtype=torch.float32)
        IMG_PROCESSOR = Sam3Processor.from_pretrained("DiffusionWave/sam3")

        TRK_MODEL = Sam3TrackerModel.from_pretrained("DiffusionWave/sam3", device_map="cpu", torch_dtype=torch.float32)
        TRK_PROCESSOR = Sam3TrackerProcessor.from_pretrained("DiffusionWave/sam3")

        print("✅ All Models loaded successfully (CPU mode)!")
    except Exception as e2:
        print(f"❌ Online loading also failed: {e2}")
        IMG_MODEL = IMG_PROCESSOR = TRK_MODEL = TRK_PROCESSOR = None

# ============ LAYER MANAGEMENT ============
class LayerManager:
    """레이어 기반 세그멘테이션 관리 클래스"""
    def __init__(self):
        self.layers = {}  # layer_id -> {'name': str, 'color': tuple, 'points': list, 'point_labels': list, 'masks': list, 'area': float}
        self.current_layer_id = None
        self.layer_counter = 0

    def create_layer(self, name, color=None):
        """새 레이어 생성"""
        if color is None:
            # 무작위 색상 생성
            import random
            color = (random.randint(50, 200), random.randint(50, 200), random.randint(50, 200))

        layer_id = f"layer_{self.layer_counter}"
        self.layers[layer_id] = {
            'name': name,
            'color': color,
            'points': [],
            'point_labels': [],  # 1: positive, 0: negative
            'masks': [],
            'area': 0.0
        }
        self.layer_counter += 1
        return layer_id

    def add_point_to_layer(self, layer_id, point, label=1):
        """레이어에 포인트 추가"""
        if layer_id in self.layers:
            self.layers[layer_id]['points'].append(point)
            self.layers[layer_id]['point_labels'].append(label)
            print(f"[add_point_to_layer] Added point to '{self.layers[layer_id]['name']}': {point}, label={label}")
            print(f"[add_point_to_layer] Total points in '{self.layers[layer_id]['name']}': {len(self.layers[layer_id]['points'])}")

    def add_mask_to_layer(self, layer_id, mask):
        """레이어에 마스크 추가"""
        if layer_id in self.layers:
            # 기존 마스크를 교체 (같은 레이어에 재세그멘테이션 시)
            self.layers[layer_id]['masks'] = [mask]
            
            # 면적 계산 - mask를 numpy array로 변환
            if isinstance(mask, torch.Tensor):
                mask_np = mask.cpu().numpy()
            else:
                mask_np = mask
            
            # 면적 계산
            area = np.sum(mask_np > 0)
            self.layers[layer_id]['area'] = area
            
            # 디버깅: 마스크 정보 출력
            print(f"[add_mask_to_layer] Layer: {self.layers[layer_id]['name']}, Mask shape: {mask_np.shape}, Area: {area}")

    def get_current_layer(self):
        """현재 선택된 레이어 반환"""
        if self.current_layer_id and self.current_layer_id in self.layers:
            return self.layers[self.current_layer_id]
        return None

    def set_current_layer(self, layer_id):
        """현재 레이어 설정"""
        self.current_layer_id = layer_id

    def clear_current_layer(self):
        """현재 레이어 초기화"""
        if self.current_layer_id and self.current_layer_id in self.layers:
            self.layers[self.current_layer_id]['points'] = []
            self.layers[self.current_layer_id]['point_labels'] = []
            self.layers[self.current_layer_id]['masks'] = []
            self.layers[self.current_layer_id]['area'] = 0.0

def calculate_total_area_ratio(layer_manager, total_pixels):
    """전체 이미지 대비 각 레이어의 면적 비율 계산"""
    ratios = []
    for layer_id, layer in layer_manager.layers.items():
        area = layer['area']
        ratio = (area / total_pixels) * 100 if total_pixels > 0 and area > 0 else 0
        has_mask = len(layer['masks']) > 0
        
        # 디버깅: 레이어 정보 출력
        print(f"[calculate_total_area_ratio] Layer: {layer['name']}, Area: {area}, Ratio: {ratio}%, Masks: {len(layer['masks'])}, Has mask: {has_mask}")
        
        ratios.append({
            'layer_name': layer['name'],
            'area_pixels': int(area),
            'ratio_percent': round(ratio, 2)
        })
    return ratios

def create_area_chart_data(ratios):
    """면적 데이터를 테이블 포맷으로 변환"""
    if not ratios:
        return pd.DataFrame(columns=["Layer", "Area (pixels)", "Ratio(%)"])

    data = []
    for ratio in ratios:
        data.append({
            "Layer": ratio['layer_name'],
            "Area (pixels)": f"{ratio['area_pixels']:,}",
            "Ratio(%)": f"{ratio['ratio_percent']}%"
        })

    return pd.DataFrame(data)

# ============ UTILITY FUNCTIONS ============
def compose_all_layers(base_image, layer_manager, opacity=0.5, border_width=2):
    """모든 레이어를 합성하여 최종 이미지 생성"""
    if isinstance(base_image, np.ndarray):
        base_image = Image.fromarray(base_image)
    base_image = base_image.convert("RGBA")

    if not layer_manager.layers:
        return base_image.convert("RGB")

    composite_layer = Image.new("RGBA", base_image.size, (0, 0, 0, 0))

    for layer_id, layer in layer_manager.layers.items():
        if not layer['masks']:
            continue

        layer_color = layer['color']

        for mask in layer['masks']:
            if isinstance(mask, torch.Tensor):
                mask = mask.cpu().numpy()
            mask = mask.astype(np.uint8)

            if mask.ndim == 3: mask = mask[0]
            if mask.ndim == 2 and mask.shape[0] == 1: mask = mask[0]

            # 마스크를 PIL 이미지로 변환
            mask_img = Image.fromarray((mask * 255).astype(np.uint8))

            # 색상 레이어 생성
            color_layer = Image.new("RGBA", base_image.size, layer_color + (0,))
            mask_alpha = mask_img.point(lambda v: int(v * opacity * 255) if v > 0 else 0)
            color_layer.putalpha(mask_alpha)

            # 테두리 추가
            if border_width > 0:
                try:
                    # 마스크의 테두리 찾기
                    mask_np = np.array(mask_img)
                    kernel_size = border_width * 2 + 1
                    dilated = cv2.dilate(mask_np, np.ones((kernel_size, kernel_size), np.uint8))
                    border = dilated - mask_np
                    border_img = Image.fromarray(border)

                    border_layer = Image.new("RGBA", base_image.size, (255, 255, 255, 255))  # 흰색 테두리
                    border_alpha = border_img.point(lambda v: 255 if v > 0 else 0)
                    border_layer.putalpha(border_alpha)

                    # 테두리를 먼저 합성
                    composite_layer = Image.alpha_composite(composite_layer, border_layer)
                except Exception as e:
                    print(f"Border creation error: {e}")

            # 마스크 레이어 합성
            composite_layer = Image.alpha_composite(composite_layer, color_layer)

    # 최종 합성
    final_result = Image.alpha_composite(base_image, composite_layer)
    return final_result.convert("RGB")

def draw_points_on_image(image, layer_manager):
    """이미지에 모든 레이어의 포인트들을 표시"""
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)

    draw_img = image.copy()
    draw = ImageDraw.Draw(draw_img)

    for layer_id, layer in layer_manager.layers.items():
        is_current = (layer_id == layer_manager.current_layer_id)
        
        for i, point in enumerate(layer['points']):
            x, y = point
            label = layer['point_labels'][i]
            
            # 포지티브: 빨간색 원, 네거티브: 파란색 X표시
            if label == 1:  # Positive
                # 큰 빨간색 원
                r = 15 if is_current else 10
                draw.ellipse((x-r, y-r, x+r, y+r), fill="red", outline="white", width=3)
                # 작은 흰색 원 (중앙)
                draw.ellipse((x-3, y-3, x+3, y+3), fill="white")
            else:  # Negative (0)
                # 큰 파란색 원
                r = 15 if is_current else 10
                draw.ellipse((x-r, y-r, x+r, y+r), fill="blue", outline="white", width=3)
                # X 표시
                line_length = 8
                draw.line([(x-line_length, y-line_length), (x+line_length, y+line_length)], fill="white", width=3)
                draw.line([(x-line_length, y+line_length), (x+line_length, y-line_length)], fill="white", width=3)

    return draw_img

# ============ UI FUNCTIONS ============
def create_new_layer(name, current_manager):
    """새 레이어 생성"""
    if not name.strip():
        return current_manager, create_layer_status_html(current_manager), gr.Dropdown(choices=[]), "Please enter a layer name"

    # 중복 이름 체크
    for layer_id, layer in current_manager.layers.items():
        if layer['name'] == name.strip():
            return current_manager, create_layer_status_html(current_manager), gr.Dropdown(choices=[(layer['name'], lid) for lid, layer in current_manager.layers.items()]), f"Layer name '{name}' already exists"

    layer_id = current_manager.create_layer(name.strip())
    current_manager.set_current_layer(layer_id)
    
    # 드롭다운 선택지 업데이트
    choices = [(layer['name'], lid) for lid, layer in current_manager.layers.items()]

    return current_manager, create_layer_status_html(current_manager), gr.Dropdown(choices=choices, value=layer_id), f"Layer '{name}' created"

def create_layer_status_html(current_manager):
    """레이어 상태 표시 HTML 생성 (시각적 표시만)"""
    if not current_manager.layers:
        return "<div style='padding: 10px; text-align: center; color: #888;'>No layers created</div>"
    
    html = "<div style='display: flex; flex-wrap: wrap; gap: 8px; padding: 10px;'>"
    
    for layer_id, layer in current_manager.layers.items():
        is_active = (current_manager.current_layer_id == layer_id)
        
        # 색상 추출
        r, g, b = layer['color']
        color_hex = f"#{r:02x}{g:02x}{b:02x}"
        
        # 활성화 상태에 따른 스타일
        if is_active:
            style = f"""
                background: linear-gradient(135deg, {color_hex}, {color_hex}dd);
                color: white;
                border: 3px solid #4682B4;
                box-shadow: 0 4px 12px rgba(70, 130, 180, 0.4);
            """
        else:
            style = f"""
                background: linear-gradient(135deg, {color_hex}aa, {color_hex}77);
                color: white;
                border: 2px solid {color_hex};
                opacity: 0.7;
            """
        
        # 포인트 개수 계산 (포지티브/네거티브 구분)
        positive_points = sum(1 for label in layer['point_labels'] if label == 1)
        negative_points = sum(1 for label in layer['point_labels'] if label == 0)
        masks_count = len(layer['masks'])
        has_mask = masks_count > 0
        
        # 상태 아이콘
        status_icon = "[OK]" if has_mask else "[ ]"
        
        html += f"""
        <div style="{style}
                    padding: 12px 20px;
                    border-radius: 8px;
                    font-weight: 600;
                    font-size: 14px;
                    min-width: 150px;">
            {status_icon} {layer['name']}<br>
            <small style='font-size: 11px; opacity: 0.9;'>
                <span style='color: #ffcccc;'>+{positive_points}</span> 
                <span style='color: #ccccff;'>-{negative_points}</span> 
                {masks_count}mask
            </small>
        </div>
        """
    
    html += "</div>"
    return html

def click_on_image(current_manager, image, point_mode, evt: gr.SelectData):
    """이미지 클릭 처리 - Include/Exclude 모드에 따라 포인트 추가"""
    if image is None or current_manager.current_layer_id is None:
        return image, current_manager, create_layer_status_html(current_manager), "Please select image and layer"

    x, y = evt.index
    
    # 포인트 모드에 따라 레이블 결정 (positive=1, negative=0)
    label = 1 if point_mode == "positive" else 0
    
    layer_name = current_manager.layers[current_manager.current_layer_id]['name']
    print(f"\n[click_on_image] ================")
    print(f"[click_on_image] Layer: {layer_name}")
    print(f"[click_on_image] Point mode: {point_mode}, Label: {label}, Position: ({x}, {y})")
    
    current_manager.add_point_to_layer(current_manager.current_layer_id, [x, y], label)

    # 포인트 표시된 이미지 생성 (원본 이미지에 포인트 표시)
    result_image = draw_points_on_image(image, current_manager)
    
    mode_text = "Include" if label == 1 else "Exclude"
    
    return result_image, current_manager, create_layer_status_html(current_manager), f"{mode_text} point added to '{layer_name}' at ({x}, {y})"

def segment_all_layers(current_manager, image, opacity, border_width):
    """모든 레이어를 순서대로 세그멘테이션 실행"""
    if image is None:
        return None, current_manager, create_layer_status_html(current_manager), "Please upload an image", pd.DataFrame()
    
    if not current_manager.layers:
        return None, current_manager, create_layer_status_html(current_manager), "Please create layers first", pd.DataFrame()

    try:
        print(f"\n[segment_all_layers] Starting segmentation for all layers...")
        segmented_count = 0
        skipped_count = 0
        
        # 모든 레이어를 순회하며 세그멘테이션
        for layer_id, layer in current_manager.layers.items():
            layer_name = layer['name']
            
            # 포인트가 없는 레이어는 건너뛰기
            if not layer['points']:
                print(f"[segment_all_layers] Skipping '{layer_name}' - no points")
                skipped_count += 1
                continue
            
            print(f"\n[segment_all_layers] Processing layer: {layer_name}")
            print(f"[segment_all_layers] Points: {len(layer['points'])}, Labels: {layer['point_labels']}")
            
            # SAM3 Tracker로 세그멘테이션
            points_list = layer['points']
            labels_list = layer['point_labels']

            input_points = [[points_list]]
            input_labels = [[labels_list]]

            inputs = TRK_PROCESSOR(images=image, input_points=input_points, input_labels=input_labels, return_tensors="pt").to(device)

            with torch.no_grad():
                outputs = TRK_MODEL(**inputs, multimask_output=False)

            masks = TRK_PROCESSOR.post_process_masks(outputs.pred_masks.cpu(), inputs["original_sizes"], binarize=True)[0]

            # 레이어에 마스크 추가
            current_manager.add_mask_to_layer(layer_id, masks[0])
            segmented_count += 1
            print(f"[segment_all_layers] Completed '{layer_name}'")

        # 결과 이미지 생성 (포인트 포함)
        result_image = compose_all_layers(image, current_manager, opacity, border_width)
        result_image = draw_points_on_image(result_image, current_manager)

        # 면적 분석
        total_pixels = image.size[0] * image.size[1]
        ratios = calculate_total_area_ratio(current_manager, total_pixels)
        chart_data = create_area_chart_data(ratios)
        
        status_msg = f"Segmentation completed! Processed: {segmented_count} layers, Skipped: {skipped_count} layers"
        print(f"\n[segment_all_layers] {status_msg}")

        return result_image, current_manager, create_layer_status_html(current_manager), status_msg, chart_data

    except Exception as e:
        import traceback
        print(f"[segment_all_layers] Error: {str(e)}")
        traceback.print_exc()
        return None, current_manager, create_layer_status_html(current_manager), f"Error: {str(e)}", pd.DataFrame()

def clear_current_layer(current_manager, image, opacity, border_width):
    """현재 레이어 초기화"""
    if current_manager.current_layer_id:
        current_manager.clear_current_layer()

        if image:
            result_image = compose_all_layers(image, current_manager, opacity, border_width)
            result_image = draw_points_on_image(result_image, current_manager)
        else:
            result_image = None

        total_pixels = image.size[0] * image.size[1] if image else 0
        ratios = calculate_total_area_ratio(current_manager, total_pixels)
        chart_data = create_area_chart_data(ratios)

        return result_image, current_manager, create_layer_status_html(current_manager), "Layer cleared", chart_data

    return None, current_manager, create_layer_status_html(current_manager), "Please select a layer", pd.DataFrame()

def refresh_visualization(current_manager, image, opacity, border_width):
    """시각화 새로고침"""
    if image is None:
        return None, "Please upload an image", pd.DataFrame()

    result_image = compose_all_layers(image, current_manager, opacity, border_width)
    result_image = draw_points_on_image(result_image, current_manager)

    total_pixels = image.size[0] * image.size[1]
    ratios = calculate_total_area_ratio(current_manager, total_pixels)
    chart_data = create_area_chart_data(ratios)

    return result_image, "Visualization updated", chart_data


# ============ GRADIO INTERFACE ============
custom_css="""
#col-container { margin: 0 auto; max-width: 1200px; }
#main-title h1 { font-size: 2.1em !important; }
.layer-button { margin: 2px; }
"""

# No custom JavaScript needed anymore
custom_js = ""

# 전역 레이어 매니저
layer_manager = LayerManager()

with gr.Blocks(css=custom_css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("# **SAM3 Layer Segmentation Tool**", elem_id="main-title")
        gr.Markdown("**Layer-based object separation and area analysis tool** | 1. Create layers 2. Select point mode and click 3. Run segmentation (processes all layers)")

        with gr.Row():
            with gr.Column(scale=1):
                img_input = gr.Image(type="pil", label="Upload Image", interactive=True, height=400)

                # 레이어 생성
                with gr.Row():
                    layer_name_input = gr.Textbox(label="Layer Name", placeholder="e.g. bench, tree, person")
                    create_layer_btn = gr.Button("Create", variant="primary")

                # 레이어 상태 표시
                gr.Markdown("### Layers Status")
                layer_buttons_html = gr.HTML("<div style='padding: 10px; text-align: center; color: #888;'>No layers created</div>")
                
                # 레이어 선택
                layer_selector = gr.Dropdown(label="Select Layer to Add Points", choices=[], interactive=True)
                
                # 포인트 모드 선택
                gr.Markdown("### Point Mode")
                with gr.Row():
                    include_btn = gr.Button("Include Point", variant="primary", size="sm")
                    exclude_btn = gr.Button("Exclude Point", variant="secondary", size="sm")
                
                point_mode_text = gr.Textbox(label="Current Mode", value="Include Point (Red)", interactive=False)
                
                # 포인트 안내
                gr.Markdown("""
                **Instructions:**
                - Select a layer from dropdown
                - Choose point mode (Include/Exclude)
                - Click on image to add point
                - **Red circle (●)**: Include this area
                - **Blue circle with X**: Exclude this area
                """)

                # 컨트롤
                with gr.Row():
                    segment_btn = gr.Button("Run All Segmentation", variant="primary", size="lg")
                    clear_btn = gr.Button("Clear Current Layer", variant="secondary")

                # 상태
                status_text = gr.Textbox(label="Status", interactive=False)
                st_layer_manager = gr.State(layer_manager)
                point_mode_state = gr.State("positive")  # "positive" or "negative"

            with gr.Column(scale=2):
                img_output = gr.Image(type="pil", label="Segmentation Result", height=400, interactive=False)

                # 면적 테이블
                area_table = gr.Dataframe(
                    label="Area Ratio by Layer"
                )

                # 설정
                with gr.Accordion("Visualization Settings", open=False):
                    opacity_slider = gr.Slider(0.1, 1.0, value=0.5, step=0.1, label="Mask Opacity")
                    border_slider = gr.Slider(0, 5, value=2, step=1, label="Border Width")

        # 이벤트 연결
        create_layer_btn.click(
            create_new_layer,
            inputs=[layer_name_input, st_layer_manager],
            outputs=[st_layer_manager, layer_buttons_html, layer_selector, status_text]
        )
        
        # 레이어 선택
        def on_layer_select(layer_id, mgr):
            if layer_id:
                mgr.set_current_layer(layer_id)
                return mgr, create_layer_status_html(mgr), f"Layer '{mgr.layers[layer_id]['name']}' selected"
            return mgr, create_layer_status_html(mgr), "Please select a layer"
        
        layer_selector.change(
            on_layer_select,
            inputs=[layer_selector, st_layer_manager],
            outputs=[st_layer_manager, layer_buttons_html, status_text]
        )
        
        # 포인트 모드 변경
        def set_include_mode():
            return "positive", "Include Point (Red)"
        
        def set_exclude_mode():
            return "negative", "Exclude Point (Blue)"
        
        include_btn.click(
            set_include_mode,
            outputs=[point_mode_state, point_mode_text]
        )
        
        exclude_btn.click(
            set_exclude_mode,
            outputs=[point_mode_state, point_mode_text]
        )

        # 이미지 클릭 이벤트 - img_input과 img_output 모두에서 클릭 받기
        img_input.select(
            click_on_image,
            inputs=[st_layer_manager, img_input, point_mode_state],
            outputs=[img_output, st_layer_manager, layer_buttons_html, status_text]
        )
        
        img_output.select(
            click_on_image,
            inputs=[st_layer_manager, img_input, point_mode_state],
            outputs=[img_output, st_layer_manager, layer_buttons_html, status_text]
        )

        # 모든 레이어 세그멘테이션 실행
        segment_btn.click(
            segment_all_layers,
            inputs=[st_layer_manager, img_input, opacity_slider, border_slider],
            outputs=[img_output, st_layer_manager, layer_buttons_html, status_text, area_table]
        )

        clear_btn.click(
            clear_current_layer,
            inputs=[st_layer_manager, img_input, opacity_slider, border_slider],
            outputs=[img_output, st_layer_manager, layer_buttons_html, status_text, area_table]
        )

        # 투명도 및 테두리 슬라이더 실시간 업데이트
        opacity_slider.change(
            refresh_visualization,
            inputs=[st_layer_manager, img_input, opacity_slider, border_slider],
            outputs=[img_output, status_text, area_table]
        )
        
        border_slider.change(
            refresh_visualization,
            inputs=[st_layer_manager, img_input, opacity_slider, border_slider],
            outputs=[img_output, status_text, area_table]
        )

        # 이미지 업로드 시 초기화
        def on_image_upload(img):
            new_manager = LayerManager()
            empty_html = "<div style='padding: 10px; text-align: center; color: #888;'>No layers created</div>"
            # 업로드한 이미지를 출력에도 표시
            return new_manager, img, pd.DataFrame(), empty_html, gr.Dropdown(choices=[], value=None), "positive", "Include Point (Red)", "New image uploaded"
        
        img_input.change(
            on_image_upload,
            inputs=[img_input],
            outputs=[st_layer_manager, img_output, area_table, layer_buttons_html, layer_selector, point_mode_state, point_mode_text, status_text]
        )

if __name__ == "__main__":
    demo.launch(show_error=True)