_app_.py

from typing import Tuple, Dict, List, Union
import gradio as gr
import supervision as sv
import numpy as np
from PIL import Image, ImageDraw, ImageFont
from ultralytics import YOLO, YOLOE
import zipfile
import os
import tempfile
import cv2
import json
from datetime import datetime
import io
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib
matplotlib.use('Agg')  # Use non-interactive backend

# Define custom models
MODEL_FILES = {
    "Line Detection": "best_line_detection_yoloe (1).pt",  # Use YOLOE for this
    "Border Detection": "border_model_weights.pt",         # Still YOLO
    "Zones Detection": "zones_model_weights.pt"            # Still YOLO
}

# Dictionary to store loaded models
models: Dict[str, Union[YOLO, YOLOE]] = {}

# Model class definitions - Expected/desired classes
EXPECTED_MODEL_CLASSES = {
    "Line Detection": [
        "line"
    ],
    "Border Detection": [
        "border",
        "decorated_initial",
        "historiated_initial", 
        "illustration",
        "page",
        "simple_initial"
    ],
    "Zones Detection": [
        "CustomZone-PageHeight",
        "CustomZone-PageWidth",
        "DamageZone",
        "DigitizationArtefactZone",
        "DropCapitalZone",
        "GraphicZone",
        "MainZone",
        "MarginTextZone",
        "MusicZone",
        "NumberingZone",
        "PageZone",
        "QuireMarksZone",
        "RunningTitleZone",
        "StampZone",
        "TitlePageZone"
    ]
}

# Model class definitions - will be populated dynamically from actual models
MODEL_CLASSES = {}

# Global variables to store results for download
current_results = []
current_images = []

# Load all custom models
# Get the directory where this script is located
script_dir = os.path.dirname(os.path.abspath(__file__))

for name, model_file in MODEL_FILES.items():
    model_path = os.path.join(script_dir, model_file)
    if os.path.exists(model_path):
        try:
            if name == "Line Detection":
                # Load YOLOE for line detection
                models[name] = YOLOE(model_path)
            else:
                # Load YOLO for other tasks
                models[name] = YOLO(model_path)
            
            # Read actual classes from the model
            if models[name] is not None:
                # Read classes from model
                actual_classes = list(models[name].names.values())
                
                # Map "object" to "line" for Line Detection model in MODEL_CLASSES
                if name == "Line Detection" and "object" in actual_classes:
                    actual_classes = ["line" if c == "object" else c for c in actual_classes]
                    print(f"   Mapped class 'object' to 'line' in Line Detection model for UI")
                
                MODEL_CLASSES[name] = actual_classes
                
                # Check for mismatch with expected classes
                if name in EXPECTED_MODEL_CLASSES:
                    expected = set(EXPECTED_MODEL_CLASSES[name])
                    actual = set(actual_classes)
                    if expected != actual:
                        print(f"⚠️  WARNING: {name} model class mismatch!")
                        print(f"   Expected: {sorted(expected)}")
                        print(f"   Actual: {sorted(actual)}")
                        print(f"   Missing in model: {sorted(expected - actual)}")
                        print(f"   Extra in model: {sorted(actual - expected)}")
                        print(f"   ⚠️  Using ACTUAL classes from model: {sorted(actual)}")
            
            print(f"✓ Loaded {name} model from {model_path}")
            print(f"  Classes available: {MODEL_CLASSES.get(name, 'Unknown')}")
        except Exception as e:
            print(f"✗ Error loading {name} model: {e}")
            models[name] = None
            # Fallback to expected classes if model fails to load
            MODEL_CLASSES[name] = EXPECTED_MODEL_CLASSES.get(name, [])
    else:
        print(f"✗ Warning: Model file {model_path} not found")
        models[name] = None
        # Fallback to expected classes if model file not found
        MODEL_CLASSES[name] = EXPECTED_MODEL_CLASSES.get(name, [])


# Create annotators
LABEL_ANNOTATOR = sv.LabelAnnotator(text_color=sv.Color.BLACK)
BOX_ANNOTATOR = sv.BoxAnnotator()
MASK_ANNOTATOR = sv.MaskAnnotator()

def detect_and_annotate_combined(
    image: np.ndarray,
    conf_threshold: float,
    iou_threshold: float,
    return_annotations: bool = False,
    selected_classes: Dict[str, List[str]] = None
) -> Union[np.ndarray, Tuple[np.ndarray, Dict]]:
    """Run all three models and combine their outputs in a single annotated image"""
    print(f"🔍 Starting detection on image shape: {image.shape}")
    
    # Colors for different models - more distinct colors
    colors = {
        "Line Detection": sv.Color.from_hex("#FF0000"),      # Bright Red
        "Border Detection": sv.Color.from_hex("#00FF00"),   # Bright Green  
        "Zones Detection": sv.Color.from_hex("#0080FF")     # Bright Blue
    }
    
    # Model prefixes for clear labeling
    model_prefixes = {
        "Line Detection": "[LINE]",
        "Border Detection": "[BORDER]", 
        "Zones Detection": "[ZONE]"
    }
    
    annotated_image = image.copy()
    total_detections = 0
    detections_data = {}
    
    # Run each model and annotate with different colors
    for model_name, model in models.items():
        if model is None:
            print(f"⏭️  Skipping {model_name} (model not loaded)")
            detections_data[model_name] = []
            continue
            
        # Check if any classes are selected for this model BEFORE running inference
        if selected_classes and model_name in selected_classes:
            selected_class_names = selected_classes[model_name]
            # If no classes selected for this model, skip it entirely (don't run inference)
            if not selected_class_names:
                print(f"⏭️  Skipping {model_name} (no classes selected)")
                detections_data[model_name] = []
                continue
        elif selected_classes is not None:
            # If selected_classes is provided but this model not in it, skip it
            print(f"⏭️  Skipping {model_name} (model not in selected classes)")
            detections_data[model_name] = []
            continue
        
        print(f"🤖 Running {model_name} model...")
        
        # Perform inference (guard against per-model failures)
        try:
            results = model.predict(
                image,
                conf=conf_threshold,
                iou=iou_threshold
            )[0]
        except Exception as e:
            print(f"✗ {model_name} inference failed: {e}")
            detections_data[model_name] = []
            continue

        model_detections = []
        
        if len(results.boxes) > 0:
            # Convert results to supervision Detections
            boxes = results.boxes.xyxy.cpu().numpy()
            confidence = results.boxes.conf.cpu().numpy()
            class_ids = results.boxes.cls.cpu().numpy().astype(int)
            
            # Filter by selected classes - only show selected classes
            if selected_classes and model_name in selected_classes:
                selected_class_names = selected_classes[model_name]
                
                # Get class names for this model
                model_class_names = results.names
                # Find class IDs that match selected class names
                selected_class_ids = []
                for class_id, class_name in model_class_names.items():
                    # For Line Detection: also match "object" when user selects "line"
                    if model_name == "Line Detection" and class_name == "object" and "line" in selected_class_names:
                        selected_class_ids.append(class_id)
                    elif class_name in selected_class_names:
                        selected_class_ids.append(class_id)
                
                # Filter detections to only show selected classes
                mask = np.isin(class_ids, selected_class_ids)
                if not np.any(mask):
                    print(f"   No detections match selected classes for {model_name}")
                    detections_data[model_name] = []
                    continue
                
                boxes = boxes[mask]
                confidence = confidence[mask]
                class_ids = class_ids[mask]
                print(f"   Filtered to {len(boxes)} detections matching selected classes: {selected_class_names}")
            
            total_detections += len(boxes)
            
            # Store detection data for COCO format
            for i, (box, conf, class_id) in enumerate(zip(boxes, confidence, class_ids)):
                x1, y1, x2, y2 = box
                width = x2 - x1
                height = y2 - y1
                
                class_name = results.names[class_id]
                # Map "object" to "line" for Line Detection model
                if model_name == "Line Detection" and class_name == "object":
                    class_name = "line"
                
                model_detections.append({
                    "bbox": [float(x1), float(y1), float(width), float(height)],  # COCO format: [x, y, width, height]
                    "class_name": class_name,
                    "confidence": float(conf)
                })
            
            
            # Create Detections object for visualization
            detections = sv.Detections(
                xyxy=boxes,
                confidence=confidence,
                mask=results.masks.data.cpu().numpy() if results.masks is not None else None,
                class_id=class_ids
            )
            
            # Create labels with clear model prefixes and confidence scores
            model_prefix = model_prefixes[model_name]
            labels = []
            for class_id, conf in zip(class_ids, confidence):
                class_name = results.names[class_id]
                # Map "object" to "line" for Line Detection model
                if model_name == "Line Detection" and class_name == "object":
                    class_name = "line"
                labels.append(f"{model_prefix} {class_name} ({conf:.2f})")

            # Create annotators with specific colors and improved styling
            box_annotator = sv.BoxAnnotator(
                color=colors[model_name],
                thickness=3  # Thicker boxes for better visibility
            )
            label_annotator = sv.LabelAnnotator(
                text_color=sv.Color.WHITE,
                color=colors[model_name],
                text_thickness=2,
                text_scale=0.6,
                text_padding=8
            )
            
        # Replace the "annotate image" block inside detect_and_annotate_combined with this

            # Annotate image depending on model type
            if model_name == "Line Detection" and results.masks is not None:

                original_h, original_w = annotated_image.shape[:2]

                if detections.mask is not None:
                    all_resized_masks = []
                    for i, mask in enumerate(detections.mask):
                        # ensure binary mask
                        mask_np = (mask > 0).astype(np.uint8)
                        resized_mask = cv2.resize(
                            mask_np,
                            (original_w, original_h),
                            interpolation=cv2.INTER_NEAREST
                        )
                        resized_mask = resized_mask.astype(bool)  # <- important
                        all_resized_masks.append(resized_mask)

                    all_resized_masks = np.stack(all_resized_masks, axis=0)  # (N, H, W)
                    detections.mask = all_resized_masks  # overwrite with clean boolean masks
                    print("Resized masks:", detections.mask.shape, detections.mask.dtype)
                else:
                    detections.mask = None


                # Use MaskAnnotator for line detection
                mask_annotator = sv.MaskAnnotator(
                    color=colors[model_name],
                    opacity=0.6
                )
                annotated_image = mask_annotator.annotate(scene=annotated_image, detections=detections)
                
                # Add labels on top of masks
                annotated_image = label_annotator.annotate(
                    scene=annotated_image,
                    detections=detections,
                    labels=labels
                )
            else:
                # Use BoxAnnotator for Border and Zones
                annotated_image = box_annotator.annotate(scene=annotated_image, detections=detections)
                annotated_image = label_annotator.annotate(scene=annotated_image, detections=detections, labels=labels)

        else:
            print(f"   No detections found for {model_name}")
        
        detections_data[model_name] = model_detections
    
    print(f"🎯 Detection completed. Total detections: {total_detections}")
    
    if return_annotations:
        return annotated_image, detections_data
    else:
        return annotated_image

def process_zip_file(zip_file_path: str, conf_threshold: float, iou_threshold: float, selected_classes: Dict[str, List[str]] = None) -> Tuple[List[Tuple[str, np.ndarray]], List[Tuple[str, Dict]], Dict]:
    """Process all images in a zip file and return annotated images, detection data, and image info"""
    print(f"📁 Opening ZIP file: {zip_file_path}")
    results = []
    annotations_data = []
    image_info = {}
    
    try:
        with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
            print(f"📋 ZIP file contents: {zip_ref.namelist()}")
            
            # Create temporary directory to extract files
            with tempfile.TemporaryDirectory() as temp_dir:
                print(f"📂 Extracting to temporary directory: {temp_dir}")
                zip_ref.extractall(temp_dir)
                
                # List all files in temp directory
                all_files = os.listdir(temp_dir)
                print(f"📄 Files extracted: {all_files}")
                
                # Process each image file (recursively search through folders)
                image_count = 0
                
                # Walk through all directories and subdirectories
                for root, dirs, files in os.walk(temp_dir):
                    print(f"📂 Searching in directory: {root}")
                    
                    for filename in files:
                        # Skip macOS hidden files
                        if filename.startswith('._') or filename.startswith('.DS_Store'):
                            print(f"⏭️  Skipping system file: {filename}")
                            continue
                            
                        if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tiff')):
                            image_count += 1
                            image_path = os.path.join(root, filename)
                            print(f"🖼️  Processing image {image_count}: {filename} (from {os.path.relpath(root, temp_dir)})")
                            
                            # Load image
                            image = cv2.imread(image_path)
                            if image is not None:
                                print(f"✅ Image loaded successfully: {image.shape}")
                                image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
                                
                                # Store image info
                                height, width = image.shape[:2]
                                image_info[filename] = (height, width)
                                
                                # Process with all models and get annotation data
                                print(f"🔍 Running detection models on {filename}...")
                                annotated_image, detections_data = detect_and_annotate_combined(
                                    image, conf_threshold, iou_threshold, return_annotations=True, selected_classes=selected_classes
                                )
                                print(f"✅ Detection completed for {filename}")
                                
                                results.append((filename, annotated_image))
                                annotations_data.append((filename, detections_data))
                            else:
                                print(f"❌ Failed to load image: {filename}")
                        else:
                            print(f"⏭️  Skipping non-image file: {filename}")
                
                print(f"📊 Total images processed: {len(results)} out of {image_count} image files found")
                print(f"📁 Searched through all subdirectories recursively")
        
        print(f"🎉 ZIP processing completed successfully! Processed {len(results)} images")
        return results, annotations_data, image_info
        
    except Exception as e:
        print(f"💥 ERROR in process_zip_file: {str(e)}")
        import traceback
        traceback.print_exc()
        return [], [], {}

def create_coco_annotations(results_data: List, image_info: Dict) -> Dict:
    """Convert detection results to COCO JSON format"""
    coco_data = {
        "info": {
            "description": "Medieval Manuscript Detection Results",
            "version": "1.0",
            "year": datetime.now().year,
            "contributor": "Medieval YOLO Models",
            "date_created": datetime.now().isoformat()
        },
        "licenses": [
            {
                "id": 1,
                "name": "Custom License",
                "url": ""
            }
        ],
        "images": [],
        "annotations": [],
        "categories": []
    }
    
    # Create categories from all models
    category_id = 1
    category_map = {}
    
    # Add categories for each model type
    for model_name in ["Line Detection", "Border Detection", "Zones Detection"]:
        if model_name in models and models[model_name] is not None:
            model = models[model_name]
            for class_id, class_name in model.names.items():
                full_name = f"{model_name}_{class_name}"
                if full_name not in category_map:
                    category_map[full_name] = category_id
                    coco_data["categories"].append({
                        "id": category_id,
                        "name": full_name,
                        "supercategory": model_name
                    })
                    category_id += 1
    
    annotation_id = 1
    
    for image_idx, (filename, detections_by_model) in enumerate(results_data):
        # Add image info
        image_id = image_idx + 1
        img_height, img_width = image_info.get(filename, (0, 0))
        
        coco_data["images"].append({
            "id": image_id,
            "file_name": filename,
            "width": img_width,
            "height": img_height,
            "license": 1
        })
        
        # Add annotations for each model
        for model_name, detections in detections_by_model.items():
            if detections:
                for detection in detections:
                    bbox = detection["bbox"]  # [x, y, width, height]
                    class_name = detection["class_name"]
                    confidence = detection["confidence"]
                    
                    full_category_name = f"{model_name}_{class_name}"
                    category_id = category_map.get(full_category_name, 1)
                    
                    coco_data["annotations"].append({
                        "id": annotation_id,
                        "image_id": image_id,
                        "category_id": category_id,
                        "bbox": bbox,
                        "area": bbox[2] * bbox[3],
                        "iscrowd": 0,
                        "score": confidence
                    })
                    annotation_id += 1
    
    return coco_data

def create_download_zip(images: List[Tuple[str, np.ndarray]], annotations: Dict) -> str:
    """Create a ZIP file with images and annotations"""
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    zip_filename = f"medieval_detection_results_{timestamp}.zip"
    zip_path = os.path.join(tempfile.gettempdir(), zip_filename)
    
    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
        # Add images
        for filename, image_array in images:
            # Convert numpy array to PIL Image and save as bytes
            pil_image = Image.fromarray(image_array.astype('uint8'))
            img_bytes = io.BytesIO()
            
            # Determine format from filename
            if filename.lower().endswith('.png'):
                pil_image.save(img_bytes, format='PNG')
            else:
                pil_image.save(img_bytes, format='JPEG')
            
            # Add to ZIP
            zipf.writestr(f"images/{filename}", img_bytes.getvalue())
        
        # Add annotations
        annotations_json = json.dumps(annotations, indent=2)
        zipf.writestr("annotations.json", annotations_json)
        
        # Add README
        readme_content = f"""Medieval Manuscript Detection Results
=============================================

Generated on: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}

Contents:
- images/: Annotated images with detection results
- annotations.json: COCO format annotations

Models and Color Coding:
- Line Detection (Red boxes with [LINE] prefix)
- Border Detection (Green boxes with [BORDER] prefix) 
- Zones Detection (Blue boxes with [ZONE] prefix)

Label format: [MODEL] class_name (confidence_score)
Annotation format: COCO JSON
For more info: https://cocodataset.org/#format-data
"""
        zipf.writestr("README.txt", readme_content)
    
    return zip_path

def calculate_statistics(detections_data: Dict, selected_classes: Dict[str, List[str]] = None) -> Dict[str, int]:
    """Calculate statistics (count per class) from detections_data"""
    stats = {}
    
    for model_name, detections in detections_data.items():
        if not detections:
            continue
            
        # Filter by selected classes if provided
        for detection in detections:
            class_name = detection["class_name"]
            
            # Only count if class is in selected classes (if selected_classes is provided)
            if selected_classes:
                if model_name not in selected_classes:
                    continue
                if class_name not in selected_classes[model_name]:
                    continue
            
            # Create full class identifier (model_name + class_name)
            full_class_name = f"{model_name} - {class_name}"
            
            if full_class_name not in stats:
                stats[full_class_name] = 0
            stats[full_class_name] += 1
    
    return stats

def create_statistics_table(stats: Dict[str, int], image_name: str = None) -> pd.DataFrame:
    """Create a pandas DataFrame table from statistics"""
    if not stats:
        return pd.DataFrame(columns=["Class", "Count"])
    
    data = []
    for class_name, count in sorted(stats.items()):
        data.append({"Class": class_name, "Count": count})
    
    df = pd.DataFrame(data)
    if image_name:
        df.insert(0, "Image", image_name)
    
    return df

def create_statistics_graph(stats: Dict[str, int], image_name: str = None) -> str:
    """Create a bar chart from statistics and return as image path"""
    if not stats:
        # Return empty graph
        fig, ax = plt.subplots(figsize=(10, 6))
        ax.text(0.5, 0.5, "No detections found", ha='center', va='center', fontsize=14)
        ax.set_xticks([])
        ax.set_yticks([])
    else:
        classes = sorted(stats.keys())
        counts = [stats[c] for c in classes]
        
        fig, ax = plt.subplots(figsize=(12, 6))
        bars = ax.bar(range(len(classes)), counts, color='steelblue')
        ax.set_xlabel('Class', fontsize=12)
        ax.set_ylabel('Count', fontsize=12)
        ax.set_title(f'Detection Statistics{(" - " + image_name) if image_name else ""}', fontsize=14, fontweight='bold')
        ax.set_xticks(range(len(classes)))
        ax.set_xticklabels(classes, rotation=45, ha='right')
        
        # Add count labels on bars
        for bar, count in zip(bars, counts):
            height = bar.get_height()
            ax.text(bar.get_x() + bar.get_width()/2., height,
                   f'{count}',
                   ha='center', va='bottom', fontsize=10)
        
        plt.tight_layout()
    
    # Save to temporary file
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    graph_path = os.path.join(tempfile.gettempdir(), f"statistics_graph_{timestamp}.png")
    fig.savefig(graph_path, dpi=150, bbox_inches='tight')
    plt.close(fig)
    
    return graph_path

def create_statistics_csv(stats: Dict[str, int], image_name: str = None) -> str:
    """Create CSV file from statistics"""
    df = create_statistics_table(stats, image_name)
    
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    csv_path = os.path.join(tempfile.gettempdir(), f"statistics_{timestamp}.csv")
    df.to_csv(csv_path, index=False)
    
    return csv_path

def create_statistics_json(stats: Dict[str, int], image_name: str = None) -> str:
    """Create JSON file from statistics"""
    data = {
        "image": image_name,
        "timestamp": datetime.now().isoformat(),
        "statistics": stats
    }
    
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    json_path = os.path.join(tempfile.gettempdir(), f"statistics_{timestamp}.json")
    
    with open(json_path, 'w') as f:
        json.dump(data, f, indent=2)
    
    return json_path

def calculate_batch_statistics(results_data: List[Tuple[str, Dict]], selected_classes: Dict[str, List[str]] = None) -> pd.DataFrame:
    """Calculate statistics for all images in batch processing - per image"""
    all_stats = []
    
    for filename, detections_by_model in results_data:
        stats = calculate_statistics(detections_by_model, selected_classes)
        df = create_statistics_table(stats, filename)
        if not df.empty:
            all_stats.append(df)
    
    if all_stats:
        combined_df = pd.concat(all_stats, ignore_index=True)
        return combined_df
    else:
        return pd.DataFrame(columns=["Image", "Class", "Count"])

def calculate_batch_statistics_summary(results_data: List[Tuple[str, Dict]], selected_classes: Dict[str, List[str]] = None) -> pd.DataFrame:
    """Calculate overall aggregated statistics for all images in batch"""
    # Aggregate statistics across all images
    all_stats = {}
    
    for filename, detections_by_model in results_data:
        stats = calculate_statistics(detections_by_model, selected_classes)
        for class_name, count in stats.items():
            if class_name not in all_stats:
                all_stats[class_name] = 0
            all_stats[class_name] += count
    
    # Create summary table
    if not all_stats:
        return pd.DataFrame(columns=["Class", "Total Count"])
    
    data = []
    for class_name, count in sorted(all_stats.items()):
        data.append({"Class": class_name, "Total Count": count})
    
    return pd.DataFrame(data)

def create_batch_statistics_graph(results_data: List[Tuple[str, Dict]], selected_classes: Dict[str, List[str]] = None) -> str:
    """Create a graph showing statistics across all images in batch"""
    # Aggregate statistics across all images
    all_stats = {}
    
    for filename, detections_by_model in results_data:
        stats = calculate_statistics(detections_by_model, selected_classes)
        for class_name, count in stats.items():
            if class_name not in all_stats:
                all_stats[class_name] = 0
            all_stats[class_name] += count
    
    return create_statistics_graph(all_stats, "Batch Processing")

def create_batch_statistics_csv(results_data: List[Tuple[str, Dict]], selected_classes: Dict[str, List[str]] = None) -> str:
    """Create CSV file from batch statistics - includes both per-image and summary"""
    # Get per-image statistics
    per_image_df = calculate_batch_statistics(results_data, selected_classes)
    # Get summary statistics
    summary_df = calculate_batch_statistics_summary(results_data, selected_classes)
    
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    csv_path = os.path.join(tempfile.gettempdir(), f"batch_statistics_{timestamp}.csv")
    
    # Write both to CSV with separator
    with open(csv_path, 'w') as f:
        # Write per-image statistics
        f.write("=== PER IMAGE STATISTICS ===\n")
        per_image_df.to_csv(f, index=False)
        f.write("\n\n=== OVERALL SUMMARY STATISTICS ===\n")
        summary_df.to_csv(f, index=False)
    
    return csv_path

def create_batch_statistics_json(results_data: List[Tuple[str, Dict]], selected_classes: Dict[str, List[str]] = None) -> str:
    """Create JSON file from batch statistics - includes both per-image and summary"""
    # Calculate summary statistics
    summary_stats = {}
    for filename, detections_by_model in results_data:
        stats = calculate_statistics(detections_by_model, selected_classes)
        for class_name, count in stats.items():
            if class_name not in summary_stats:
                summary_stats[class_name] = 0
            summary_stats[class_name] += count
    
    data = {
        "batch_processing": True,
        "timestamp": datetime.now().isoformat(),
        "total_images": len(results_data),
        "per_image_statistics": [],
        "overall_summary": summary_stats
    }
    
    for filename, detections_by_model in results_data:
        stats = calculate_statistics(detections_by_model, selected_classes)
        data["per_image_statistics"].append({
            "filename": filename,
            "statistics": stats
        })
    
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    json_path = os.path.join(tempfile.gettempdir(), f"batch_statistics_{timestamp}.json")
    
    with open(json_path, 'w') as f:
        json.dump(data, f, indent=2)
    
    return json_path

# Create Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Medieval Manuscript Detection with Custom YOLO Models")
    gr.Markdown("""
    **Models and Color Coding:**
    - 🔵**Line Detection** - Red boxes with [LINE] prefix
    - 🟢 **Border Detection** - Green boxes with [BORDER] prefix  
    - 🟠 **Zones Detection** - Blue boxes with [ZONE] prefix
    
    Each detection shows: **[MODEL] class_name (confidence_score)**
    """)
    
    with gr.Tabs():
        # Single Image Tab
        with gr.TabItem("Single Image"):
            with gr.Row():
                with gr.Column():
                    input_image = gr.Image(
                        label="Input Image",
                        type='numpy'
                    )
                    with gr.Accordion("Detection Settings", open=True):
                        with gr.Row():
                            conf_threshold = gr.Slider(
                                label="Confidence Threshold",
                                minimum=0.0,
                                maximum=1.0,
                                step=0.05,
                                value=0.25,
                            )
                            iou_threshold = gr.Slider(
                                label="IoU Threshold",
                                minimum=0.0,
                                maximum=1.0,
                                step=0.05,
                                value=0.45,
                                info="Decrease for stricter detection, increase for more overlapping boxes"
                            )
                    
                    with gr.Accordion("Class Selection", open=False):
                        gr.Markdown("**Select which classes to detect for each model:**")
                        with gr.Row():
                            with gr.Column():
                                line_classes = gr.CheckboxGroup(
                                    label="Line Detection Classes",
                                    choices=MODEL_CLASSES["Line Detection"],
                                    value=MODEL_CLASSES["Line Detection"],  # All selected by default
                                    info="Select at least one class for detection"
                                )
                                with gr.Row():
                                    line_select_all = gr.Button("Select All", size="sm")
                                    line_unselect_all = gr.Button("Unselect All", size="sm")
                            with gr.Column():
                                border_classes = gr.CheckboxGroup(
                                    label="Border Detection Classes", 
                                    choices=MODEL_CLASSES["Border Detection"],
                                    value=MODEL_CLASSES["Border Detection"],  # All selected by default
                                    info="Select at least one class for detection"
                                )
                                with gr.Row():
                                    border_select_all = gr.Button("Select All", size="sm")
                                    border_unselect_all = gr.Button("Unselect All", size="sm")
                        with gr.Row():
                            with gr.Column():
                                zones_classes = gr.CheckboxGroup(
                                    label="Zones Detection Classes",
                                    choices=MODEL_CLASSES["Zones Detection"],
                                    value=MODEL_CLASSES["Zones Detection"],  # All selected by default
                                    info="Select at least one class for detection"
                                )
                                with gr.Row():
                                    zones_select_all = gr.Button("Select All", size="sm")
                                    zones_unselect_all = gr.Button("Unselect All", size="sm")
                    with gr.Row():
                        clear_btn = gr.Button("Clear")
                        detect_btn = gr.Button("Detect with All Models", variant="primary")
                        
                with gr.Column():
                    output_image = gr.Image(
                        label="Combined Detection Result",
                        type='numpy'
                    )
                    
                    # Single image download buttons
                    with gr.Row():
                        single_download_json_btn = gr.Button(
                            "📄 Download Annotations (JSON)",
                            variant="secondary",
                            size="sm"
                        )
                        single_download_image_btn = gr.Button(
                            "🖼️ Download Image",
                            variant="secondary",
                            size="sm"
                        )
                    
                    # Single image file outputs
                    single_json_output = gr.File(
                        label="📄 JSON Download",
                        visible=True,
                        height=50
                    )
                    single_image_output = gr.File(
                        label="🖼️ Image Download",
                        visible=True,
                        height=50
                    )
                    
                    # Statistics section for single image
                    with gr.Accordion("📊 Statistics", open=False):
                        with gr.Tabs():
                            with gr.TabItem("Table"):
                                single_stats_table = gr.Dataframe(
                                    label="Detection Statistics",
                                    headers=["Class", "Count"],
                                    wrap=True
                                )
                            with gr.TabItem("Graph"):
                                single_stats_graph = gr.Image(
                                    label="Detection Statistics Graph",
                                    type='filepath'
                                )
                        
                        # Statistics download buttons
                        with gr.Row():
                            single_download_stats_csv_btn = gr.Button(
                                "📊 Download Statistics (CSV)",
                                variant="secondary",
                                size="sm"
                            )
                            single_download_stats_json_btn = gr.Button(
                                "📊 Download Statistics (JSON)",
                                variant="secondary",
                                size="sm"
                            )
                        
                        single_stats_csv_output = gr.File(
                            label="📊 Statistics CSV Download",
                            visible=False,
                            height=50
                        )
                        single_stats_json_output = gr.File(
                            label="📊 Statistics JSON Download",
                            visible=False,
                            height=50
                        )
        
        # Batch Processing Tab
        with gr.TabItem("Batch Processing (ZIP)"):
            with gr.Row():
                with gr.Column():
                    zip_file = gr.File(
                        label="Upload ZIP file with images",
                        file_types=[".zip"]
                    )
                    with gr.Accordion("Detection Settings", open=True):
                        with gr.Row():
                            batch_conf_threshold = gr.Slider(
                                label="Confidence Threshold",
                                minimum=0.0,
                                maximum=1.0,
                                step=0.05,
                                value=0.25,
                            )
                            batch_iou_threshold = gr.Slider(
                                label="IoU Threshold",
                                minimum=0.0,
                                maximum=1.0,
                                step=0.05,
                                value=0.45,
                            )
                    
                    with gr.Accordion("Class Selection", open=False):
                        gr.Markdown("**Select which classes to detect for each model:**")
                        with gr.Row():
                            with gr.Column():
                                batch_line_classes = gr.CheckboxGroup(
                                    label="Line Detection Classes",
                                    choices=MODEL_CLASSES["Line Detection"],
                                    value=MODEL_CLASSES["Line Detection"],  # All selected by default
                                    info="Select at least one class for detection"
                                )
                                with gr.Row():
                                    batch_line_select_all = gr.Button("Select All", size="sm")
                                    batch_line_unselect_all = gr.Button("Unselect All", size="sm")
                            with gr.Column():
                                batch_border_classes = gr.CheckboxGroup(
                                    label="Border Detection Classes", 
                                    choices=MODEL_CLASSES["Border Detection"],
                                    value=MODEL_CLASSES["Border Detection"],  # All selected by default
                                    info="Select at least one class for detection"
                                )
                                with gr.Row():
                                    batch_border_select_all = gr.Button("Select All", size="sm")
                                    batch_border_unselect_all = gr.Button("Unselect All", size="sm")
                        with gr.Row():
                            with gr.Column():
                                batch_zones_classes = gr.CheckboxGroup(
                                    label="Zones Detection Classes",
                                    choices=MODEL_CLASSES["Zones Detection"],
                                    value=MODEL_CLASSES["Zones Detection"],  # All selected by default
                                    info="Select at least one class for detection"
                                )
                                with gr.Row():
                                    batch_zones_select_all = gr.Button("Select All", size="sm")
                                    batch_zones_unselect_all = gr.Button("Unselect All", size="sm")
                    
                    # Add status message box
                    batch_status = gr.Textbox(
                        label="Processing Status",
                        value="Ready to process ZIP file...",
                        interactive=False,
                        max_lines=3
                    )
                    
                    with gr.Row():
                        clear_batch_btn = gr.Button("Clear")
                        process_batch_btn = gr.Button("Process ZIP", variant="primary")
                        
                with gr.Column():
                    batch_gallery = gr.Gallery(
                        label="Batch Processing Results",
                        show_label=True,
                        elem_id="gallery",
                        columns=2,
                        rows=2,
                        height="auto",
                        type="numpy"  # Explicitly handle numpy arrays
                    )
                    
                    # Download buttons
                    with gr.Row():
                        download_json_btn = gr.Button(
                            "📄 Download COCO Annotations (JSON)",
                            variant="secondary"
                        )
                        download_zip_btn = gr.Button(
                            "📦 Download Results (ZIP)",
                            variant="secondary"
                        )
                    
                    # File outputs for downloads
                    json_file_output = gr.File(
                        label="📄 JSON Download",
                        visible=True,
                        height=50
                    )
                    zip_file_output = gr.File(
                        label="📦 ZIP Download",
                        visible=True,
                        height=50
                    )
                    
                    # Statistics section for batch processing
                    with gr.Accordion("📊 Statistics", open=False):
                        with gr.Tabs():
                            with gr.TabItem("Per Image"):
                                batch_stats_table = gr.Dataframe(
                                    label="Detection Statistics Per Image",
                                    wrap=True
                                )
                            with gr.TabItem("Overall Summary"):
                                batch_stats_summary_table = gr.Dataframe(
                                    label="Overall Statistics Summary (All Images Combined)",
                                    wrap=True
                                )
                            with gr.TabItem("Graph"):
                                batch_stats_graph = gr.Image(
                                    label="Detection Statistics Graph (Aggregated)",
                                    type='filepath'
                                )
                        
                        # Statistics download buttons
                        with gr.Row():
                            batch_download_stats_csv_btn = gr.Button(
                                "📊 Download Statistics (CSV)",
                                variant="secondary",
                                size="sm"
                            )
                            batch_download_stats_json_btn = gr.Button(
                                "📊 Download Statistics (JSON)",
                                variant="secondary",
                                size="sm"
                            )
                        
                        batch_stats_csv_output = gr.File(
                            label="📊 Statistics CSV Download",
                            visible=False,
                            height=50
                        )
                        batch_stats_json_output = gr.File(
                            label="📊 Statistics JSON Download",
                            visible=False,
                            height=50
                        )

    # Global variables for single image results
    single_image_result = None
    single_image_annotations = None
    single_image_filename = None
    single_image_selected_classes = None
    
    def process_single_image(
        image: np.ndarray,
        conf_threshold: float,
        iou_threshold: float,
        line_classes: List[str],
        border_classes: List[str],
        zones_classes: List[str]
    ) -> Tuple[np.ndarray, np.ndarray, pd.DataFrame, str]:
        global single_image_result, single_image_annotations, single_image_filename, single_image_selected_classes
        
        if image is None:
            single_image_result = None
            single_image_annotations = None
            single_image_filename = None
            single_image_selected_classes = None
            return None, None, pd.DataFrame(columns=["Class", "Count"]), None
            
        # Validate that at least one class is selected
        all_selected = (line_classes or []) + (border_classes or []) + (zones_classes or [])
        if not all_selected:
            raise gr.Error("⚠️ Please select at least one class for detection!")
        
        # Prepare selected classes dictionary
        selected_classes = {
            "Line Detection": line_classes or [],
            "Border Detection": border_classes or [],
            "Zones Detection": zones_classes or []
        }
        
        # Process with annotations
        try:
            annotated_image, detections_data = detect_and_annotate_combined(
                image, conf_threshold, iou_threshold, return_annotations=True, selected_classes=selected_classes
            )
        except Exception as e:
            # Surface a nice error to the UI without crashing the app
            raise gr.Error(f"Detection failed: {str(e)}")
        
        # Calculate statistics
        stats = calculate_statistics(detections_data, selected_classes)
        stats_table = create_statistics_table(stats, single_image_filename)
        stats_graph_path = create_statistics_graph(stats, single_image_filename)
        
        # Store results globally for download
        single_image_result = annotated_image
        single_image_annotations = detections_data
        single_image_selected_classes = selected_classes
        single_image_filename = f"detection_result_{datetime.now().strftime('%Y%m%d_%H%M%S')}.jpg"
        
        return image, annotated_image, stats_table, stats_graph_path

    # Global variables for batch results
    current_batch_results = []
    current_batch_selected_classes = None
    
    def process_batch_images_with_status(
        zip_file,
        conf_threshold: float,
        iou_threshold: float,
        line_classes: List[str],
        border_classes: List[str],
        zones_classes: List[str]
    ):
        global current_batch_results, current_batch_selected_classes
        
        print("🚀 ========== BATCH PROCESSING STARTED ==========")
        
        if zip_file is None:
            print("❌ No ZIP file provided")
            return [], "Please upload a ZIP file first.", pd.DataFrame(columns=["Image", "Class", "Count"]), pd.DataFrame(columns=["Class", "Total Count"]), None
        
        print(f"📁 ZIP file received: {zip_file.name}")
        print(f"⚙️  Settings: conf_threshold={conf_threshold}, iou_threshold={iou_threshold}")
        
        try:
            # Validate that at least one class is selected
            all_selected = (line_classes or []) + (border_classes or []) + (zones_classes or [])
            if not all_selected:
                raise gr.Error("⚠️ Please select at least one class for detection!")
            
            # Prepare selected classes dictionary
            selected_classes = {
                "Line Detection": line_classes or [],
                "Border Detection": border_classes or [],
                "Zones Detection": zones_classes or []
            }
            current_batch_selected_classes = selected_classes
            
            # Process zip file
            print("🔄 Starting ZIP file processing...")
            results, annotations_data, image_info = process_zip_file(zip_file.name, conf_threshold, iou_threshold, selected_classes)
            
            # Store batch results globally
            current_batch_results = results
            
            if not results:
                error_msg = "No valid images found in ZIP file."
                print(f"❌ {error_msg}")
                return [], error_msg
            
            # Store data globally for download
            global current_results, current_images
            current_images = results
            current_results = annotations_data
            
            print(f"📊 ZIP processing returned {len(results)} results")
            
            # Convert results to format expected by Gallery
            print("🔄 Converting results for Gradio Gallery...")
            gallery_images = []
            
            for i, (filename, annotated_image) in enumerate(results):
                print(f"🖼️  Converting image {i+1}/{len(results)}: {filename}")
                print(f"   Image shape: {annotated_image.shape}, dtype: {annotated_image.dtype}")
                
                # Ensure the image is in the right format and range
                if annotated_image.dtype != 'uint8':
                    print(f"   Converting dtype from {annotated_image.dtype} to uint8")
                    # Normalize if needed
                    if annotated_image.max() <= 1.0:
                        annotated_image = (annotated_image * 255).astype('uint8')
                        print(f"   Normalized from [0,1] to [0,255]")
                    else:
                        annotated_image = annotated_image.astype('uint8')
                        print(f"   Cast to uint8")
                
                print(f"   Final image shape: {annotated_image.shape}, dtype: {annotated_image.dtype}")
                
                # For Gradio gallery, we can pass numpy arrays directly
                # Format: (image_data, caption)
                gallery_images.append((annotated_image, filename))
                print(f"   ✅ Added {filename} to gallery")
            
            # Calculate statistics (use annotations_data, not results)
            stats_table = calculate_batch_statistics(annotations_data, selected_classes)
            stats_summary_table = calculate_batch_statistics_summary(annotations_data, selected_classes)
            stats_graph_path = create_batch_statistics_graph(annotations_data, selected_classes)
            
            success_msg = f"✅ Successfully processed {len(gallery_images)} images!"
            print(f"🎉 {success_msg}")
            print(f"📋 Gallery contains {len(gallery_images)} items")
            print("🏁 ========== BATCH PROCESSING COMPLETED ==========\n")
            
            return gallery_images, success_msg, stats_table, stats_summary_table, stats_graph_path
            
        except Exception as e:
            error_msg = f"❌ Error: {str(e)}"
            print(f"💥 EXCEPTION in process_batch_images_with_status: {error_msg}")
            import traceback
            traceback.print_exc()
            print("💀 ========== BATCH PROCESSING FAILED ==========\n")
            return [], error_msg, pd.DataFrame(columns=["Image", "Class", "Count"]), pd.DataFrame(columns=["Class", "Total Count"]), None

    def clear_single():
        global single_image_result, single_image_annotations, single_image_filename, single_image_selected_classes
        single_image_result = None
        single_image_annotations = None
        single_image_filename = None
        single_image_selected_classes = None
        return None, None, pd.DataFrame(columns=["Class", "Count"]), None
    
    def clear_batch():
        global current_results, current_images
        current_results = []
        current_images = []
        return None, [], "Ready to process ZIP file..."
    
    def download_annotations():
        """Create and return COCO JSON annotations file"""
        global current_results, current_images
        
        if not current_results:
            print("❌ No annotation data available for download")
            return None
        
        try:
            # Create image info dictionary
            image_info = {}
            for filename, image_array in current_images:
                height, width = image_array.shape[:2]
                image_info[filename] = (height, width)
            
            # Create COCO annotations
            coco_data = create_coco_annotations(current_results, image_info)
            
            # Save to temporary file with proper name
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            json_filename = f"medieval_annotations_{timestamp}.json"
            json_path = os.path.join(tempfile.gettempdir(), json_filename)
            
            with open(json_path, 'w') as f:
                json.dump(coco_data, f, indent=2)
            
            print(f"💾 Created annotations file: {json_path}")
            print(f"📁 File size: {os.path.getsize(json_path)} bytes")
            
            # Verify file exists and is readable
            if os.path.exists(json_path) and os.path.getsize(json_path) > 0:
                return json_path
            else:
                print(f"❌ File verification failed: {json_path}")
                return None
            
        except Exception as e:
            print(f"❌ Error creating annotations: {e}")
            import traceback
            traceback.print_exc()
            return None
    
    def download_results_zip():
        """Create and return ZIP file with images and annotations"""
        global current_results, current_images
        
        if not current_results or not current_images:
            print("❌ No results data available for ZIP download")
            return None
        
        try:
            # Create image info dictionary
            image_info = {}
            for filename, image_array in current_images:
                height, width = image_array.shape[:2]
                image_info[filename] = (height, width)
            
            # Create COCO annotations
            coco_data = create_coco_annotations(current_results, image_info)
            
            # Create ZIP file
            zip_path = create_download_zip(current_images, coco_data)
            
            print(f"💾 Created results ZIP: {zip_path}")
            print(f"📁 ZIP file size: {os.path.getsize(zip_path)} bytes")
            
            # Verify file exists and is readable
            if os.path.exists(zip_path) and os.path.getsize(zip_path) > 0:
                return zip_path
            else:
                print(f"❌ ZIP file verification failed: {zip_path}")
                return None
            
        except Exception as e:
            print(f"❌ Error creating ZIP file: {e}")
            import traceback
            traceback.print_exc()
            return None
    
    def download_single_annotations():
        """Download COCO annotations for single image"""
        global single_image_annotations, single_image_result, single_image_filename
        
        if single_image_annotations is None or single_image_result is None:
            print("❌ No single image annotation data available")
            return None
        
        try:
            # Create image info
            height, width = single_image_result.shape[:2]
            image_info = {single_image_filename: (height, width)}
            
            # Create annotations data in the expected format
            annotations_data = [(single_image_filename, single_image_annotations)]
            
            # Create COCO annotations
            coco_data = create_coco_annotations(annotations_data, image_info)
            
            # Save to temporary file
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            json_filename = f"single_image_annotations_{timestamp}.json"
            json_path = os.path.join(tempfile.gettempdir(), json_filename)
            
            with open(json_path, 'w') as f:
                json.dump(coco_data, f, indent=2)
            
            print(f"💾 Created single image annotations: {json_path}")
            print(f"📁 File size: {os.path.getsize(json_path)} bytes")
            
            # Verify file exists
            if os.path.exists(json_path) and os.path.getsize(json_path) > 0:
                return json_path
            else:
                print(f"❌ Single image file verification failed: {json_path}")
                return None
            
        except Exception as e:
            print(f"❌ Error creating single image annotations: {e}")
            import traceback
            traceback.print_exc()
            return None
    
    def download_single_image():
        """Download processed single image"""
        global single_image_result, single_image_filename
        
        if single_image_result is None:
            print("❌ No single image result available")
            return None
        
        try:
            # Convert to PIL and save
            pil_image = Image.fromarray(single_image_result.astype('uint8'))
            
            # Save to temporary file
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            img_filename = f"processed_image_{timestamp}.jpg"
            img_path = os.path.join(tempfile.gettempdir(), img_filename)
            
            pil_image.save(img_path, 'JPEG', quality=95)
            
            print(f"💾 Created single image file: {img_path}")
            print(f"📁 Image file size: {os.path.getsize(img_path)} bytes")
            
            # Verify file exists
            if os.path.exists(img_path) and os.path.getsize(img_path) > 0:
                return img_path
            else:
                print(f"❌ Single image file verification failed: {img_path}")
                return None
            
        except Exception as e:
            print(f"❌ Error creating single image file: {e}")
            import traceback
            traceback.print_exc()
            return None

    # Connect buttons to functions for single image
    detect_btn.click(
        process_single_image,
        inputs=[input_image, conf_threshold, iou_threshold, line_classes, border_classes, zones_classes],
        outputs=[input_image, output_image, single_stats_table, single_stats_graph]
    )
    clear_btn.click(
        clear_single,
        inputs=None,
        outputs=[input_image, output_image, single_stats_table, single_stats_graph]
    )
    
    # Select All/Unselect All handlers for single image
    line_select_all.click(
        fn=lambda: MODEL_CLASSES["Line Detection"],
        outputs=[line_classes]
    )
    line_unselect_all.click(
        fn=lambda: [],
        outputs=[line_classes]
    )
    border_select_all.click(
        fn=lambda: MODEL_CLASSES["Border Detection"],
        outputs=[border_classes]
    )
    border_unselect_all.click(
        fn=lambda: [],
        outputs=[border_classes]
    )
    zones_select_all.click(
        fn=lambda: MODEL_CLASSES["Zones Detection"],
        outputs=[zones_classes]
    )
    zones_unselect_all.click(
        fn=lambda: [],
        outputs=[zones_classes]
    )
    
    # Connect buttons to functions for batch processing
    process_batch_btn.click(
        process_batch_images_with_status,
        inputs=[zip_file, batch_conf_threshold, batch_iou_threshold, batch_line_classes, batch_border_classes, batch_zones_classes],
        outputs=[batch_gallery, batch_status, batch_stats_table, batch_stats_summary_table, batch_stats_graph]
    )
    clear_batch_btn.click(
        clear_batch,
        inputs=None,
        outputs=[zip_file, batch_gallery, batch_status]
    )
    
    # Select All/Unselect All handlers for batch processing
    batch_line_select_all.click(
        fn=lambda: MODEL_CLASSES["Line Detection"],
        outputs=[batch_line_classes]
    )
    batch_line_unselect_all.click(
        fn=lambda: [],
        outputs=[batch_line_classes]
    )
    batch_border_select_all.click(
        fn=lambda: MODEL_CLASSES["Border Detection"],
        outputs=[batch_border_classes]
    )
    batch_border_unselect_all.click(
        fn=lambda: [],
        outputs=[batch_border_classes]
    )
    batch_zones_select_all.click(
        fn=lambda: MODEL_CLASSES["Zones Detection"],
        outputs=[batch_zones_classes]
    )
    batch_zones_unselect_all.click(
        fn=lambda: [],
        outputs=[batch_zones_classes]
    )
    
    # Connect download buttons
    download_json_btn.click(
        fn=download_annotations,
        inputs=[],
        outputs=[json_file_output]
    )
    download_zip_btn.click(
        fn=download_results_zip,
        inputs=[],
        outputs=[zip_file_output]
    )
    
    # Connect single image download buttons
    single_download_json_btn.click(
        fn=download_single_annotations,
        inputs=[],
        outputs=[single_json_output]
    )
    single_download_image_btn.click(
        fn=download_single_image,
        inputs=[],
        outputs=[single_image_output]
    )
    
    # Statistics download handlers for single image
    def download_single_stats_csv():
        global single_image_annotations, single_image_filename, single_image_selected_classes
        if single_image_annotations is None:
            return None
        stats = calculate_statistics(single_image_annotations, single_image_selected_classes)
        csv_path = create_statistics_csv(stats, single_image_filename)
        return csv_path
    
    def download_single_stats_json():
        global single_image_annotations, single_image_filename, single_image_selected_classes
        if single_image_annotations is None:
            return None
        stats = calculate_statistics(single_image_annotations, single_image_selected_classes)
        json_path = create_statistics_json(stats, single_image_filename)
        return json_path
    
    single_download_stats_csv_btn.click(
        fn=download_single_stats_csv,
        inputs=[],
        outputs=[single_stats_csv_output]
    )
    single_download_stats_json_btn.click(
        fn=download_single_stats_json,
        inputs=[],
        outputs=[single_stats_json_output]
    )
    
    # Statistics download handlers for batch processing
    def download_batch_stats_csv():
        global current_results, current_batch_selected_classes
        if not current_results:
            return None
        csv_path = create_batch_statistics_csv(current_results, current_batch_selected_classes)
        return csv_path
    
    def download_batch_stats_json():
        global current_results, current_batch_selected_classes
        if not current_results:
            return None
        json_path = create_batch_statistics_json(current_results, current_batch_selected_classes)
        return json_path
    
    batch_download_stats_csv_btn.click(
        fn=download_batch_stats_csv,
        inputs=[],
        outputs=[batch_stats_csv_output]
    )
    batch_download_stats_json_btn.click(
        fn=download_batch_stats_json,
        inputs=[],
        outputs=[batch_stats_json_output]
    )

if __name__ == "__main__":
    # Configure launch settings for better stability
    # Enable Gradio queue for more robust concurrency and error isolation
    demo.queue()
    demo.launch(
        debug=False,  # Disable debug mode for production
        show_error=True,
        server_name="0.0.0.0",
        server_port=8000,
        share=False,
        max_threads=4,  # Limit concurrent requests
        auth=None,
        inbrowser=False,
        favicon_path=None,
        ssl_verify=True,
        quiet=False
    )