app.py

import os
import shutil
import tempfile
import base64
import asyncio
from io import BytesIO

import cv2
import numpy as np
import torch
import onnxruntime as rt
from PIL import Image
import gradio as gr
from transformers import pipeline
from huggingface_hub import hf_hub_download

# Import necessary function from aesthetic_predictor_v2_5
from aesthetic_predictor_v2_5 import convert_v2_5_from_siglip


#####################################
#         Model Definitions         #
#####################################

class MLP(torch.nn.Module):
    """A simple multi-layer perceptron for image feature regression."""
    def __init__(self, input_size: int, batch_norm: bool = True):
        super().__init__()
        self.input_size = input_size
        self.layers = torch.nn.Sequential(
            torch.nn.Linear(self.input_size, 2048),
            torch.nn.ReLU(),
            torch.nn.BatchNorm1d(2048) if batch_norm else torch.nn.Identity(),
            torch.nn.Dropout(0.3),
            torch.nn.Linear(2048, 512),
            torch.nn.ReLU(),
            torch.nn.BatchNorm1d(512) if batch_norm else torch.nn.Identity(),
            torch.nn.Dropout(0.3),
            torch.nn.Linear(512, 256),
            torch.nn.ReLU(),
            torch.nn.BatchNorm1d(256) if batch_norm else torch.nn.Identity(),
            torch.nn.Dropout(0.2),
            torch.nn.Linear(256, 128),
            torch.nn.ReLU(),
            torch.nn.BatchNorm1d(128) if batch_norm else torch.nn.Identity(),
            torch.nn.Dropout(0.1),
            torch.nn.Linear(128, 32),
            torch.nn.ReLU(),
            torch.nn.Linear(32, 1)
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.layers(x)


class WaifuScorer:
    """WaifuScorer model that uses CLIP for feature extraction and a custom MLP for scoring."""
    def __init__(self, model_path: str = None, device: str = 'cuda', cache_dir: str = None, verbose: bool = False):
        self.verbose = verbose
        self.device = device
        self.dtype = torch.float32
        self.available = False

        try:
            import clip  # local import to avoid dependency issues
            # Set default model path if not provided
            if model_path is None:
                model_path = "Eugeoter/waifu-scorer-v3/model.pth"
                if self.verbose:
                    print(f"Model path not provided. Using default: {model_path}")

            # Download model if not found locally
            if not os.path.isfile(model_path):
                username, repo_id, model_name = model_path.split("/")[-3:]
                model_path = hf_hub_download(f"{username}/{repo_id}", model_name, cache_dir=cache_dir)

            if self.verbose:
                print(f"Loading WaifuScorer model from: {model_path}")

            # Initialize MLP model
            self.mlp = MLP(input_size=768)
            # Load state dict
            if model_path.endswith(".safetensors"):
                from safetensors.torch import load_file
                state_dict = load_file(model_path)
            else:
                state_dict = torch.load(model_path, map_location=device)
            self.mlp.load_state_dict(state_dict)
            self.mlp.to(device)
            self.mlp.eval()

            # Load CLIP model for image preprocessing and feature extraction
            self.clip_model, self.preprocess = clip.load("ViT-L/14", device=device)
            self.available = True
        except Exception as e:
            print(f"Unable to initialize WaifuScorer: {e}")

    @torch.no_grad()
    def __call__(self, images):
        if not self.available:
            return [None] * (len(images) if isinstance(images, list) else 1)
        if isinstance(images, Image.Image):
            images = [images]
        n = len(images)
        # Ensure at least two images for CLIP model compatibility
        if n == 1:
            images = images * 2

        image_tensors = [self.preprocess(img).unsqueeze(0) for img in images]
        image_batch = torch.cat(image_tensors).to(self.device)
        image_features = self.clip_model.encode_image(image_batch)
        # Normalize features
        norm = image_features.norm(2, dim=-1, keepdim=True)
        norm[norm == 0] = 1
        im_emb = (image_features / norm).to(device=self.device, dtype=self.dtype)
        predictions = self.mlp(im_emb)
        scores = predictions.clamp(0, 10).cpu().numpy().reshape(-1).tolist()
        return scores[:n]


#####################################
#   Aesthetic Predictor Functions   #
#####################################

def load_aesthetic_predictor_v2_5():
    """Load and return an instance of Aesthetic Predictor V2.5 with batch processing support."""
    class AestheticPredictorV2_5_Impl:
        def __init__(self):
            print("Loading Aesthetic Predictor V2.5...")
            self.model, self.preprocessor = convert_v2_5_from_siglip(
                low_cpu_mem_usage=True,
                trust_remote_code=True,
            )
            if torch.cuda.is_available():
                self.model = self.model.to(torch.bfloat16).cuda()

        def inference(self, image):
            if isinstance(image, list):
                images_rgb = [img.convert("RGB") for img in image]
                pixel_values = self.preprocessor(images=images_rgb, return_tensors="pt").pixel_values
                if torch.cuda.is_available():
                    pixel_values = pixel_values.to(torch.bfloat16).cuda()
                with torch.inference_mode():
                    scores = self.model(pixel_values).logits.squeeze().float().cpu().numpy()
                if scores.ndim == 0:
                    scores = np.array([scores])
                return scores.tolist()
            else:
                pixel_values = self.preprocessor(images=image.convert("RGB"), return_tensors="pt").pixel_values
                if torch.cuda.is_available():
                    pixel_values = pixel_values.to(torch.bfloat16).cuda()
                with torch.inference_mode():
                    score = self.model(pixel_values).logits.squeeze().float().cpu().numpy()
                return score

    return AestheticPredictorV2_5_Impl()


def load_anime_aesthetic_model():
    """Load and return the Anime Aesthetic ONNX model."""
    model_path = hf_hub_download(repo_id="skytnt/anime-aesthetic", filename="model.onnx")
    return rt.InferenceSession(model_path, providers=['CPUExecutionProvider'])


def predict_anime_aesthetic(img, model):
    """Predict Anime Aesthetic score for a single image."""
    img_np = np.array(img).astype(np.float32) / 255.0
    s = 768
    h, w = img_np.shape[:2]
    if h > w:
        new_h, new_w = s, int(s * w / h)
    else:
        new_h, new_w = int(s * h / w), s
    resized = cv2.resize(img_np, (new_w, new_h))
    # Center the resized image in a square canvas
    canvas = np.zeros((s, s, 3), dtype=np.float32)
    pad_h = (s - new_h) // 2
    pad_w = (s - new_w) // 2
    canvas[pad_h:pad_h+new_h, pad_w:pad_w+new_w] = resized
    # Prepare input for model
    input_tensor = np.transpose(canvas, (2, 0, 1))[np.newaxis, :]
    pred = model.run(None, {"img": input_tensor})[0].item()
    return pred


#####################################
#      Image Evaluation Tool        #
#####################################

class ImageEvaluationTool:
    """Evaluation tool to process images through multiple aesthetic models and generate logs and HTML outputs."""
    def __init__(self):
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        print(f"Using device: {self.device}")
        print("Loading models... This may take some time.")

        # Load models with progress logs
        print("Loading Aesthetic Shadow model...")
        self.aesthetic_shadow = pipeline("image-classification", model="NeoChen1024/aesthetic-shadow-v2-backup", device=self.device)
        print("Loading Waifu Scorer model...")
        self.waifu_scorer = WaifuScorer(device=self.device, verbose=True)
        print("Loading Aesthetic Predictor V2.5...")
        self.aesthetic_predictor = load_aesthetic_predictor_v2_5()
        print("Loading Anime Aesthetic model...")
        self.anime_aesthetic = load_anime_aesthetic_model()
        print("All models loaded successfully!")

        self.temp_dir = tempfile.mkdtemp()
        self.results = []  # Store final results for sorting and display
        self.available_models = {
            "aesthetic_shadow": {"name": "Aesthetic Shadow", "process": self._process_aesthetic_shadow},
            "waifu_scorer": {"name": "Waifu Scorer", "process": self._process_waifu_scorer},
            "aesthetic_predictor_v2_5": {"name": "Aesthetic V2.5", "process": self._process_aesthetic_predictor_v2_5},
            "anime_aesthetic": {"name": "Anime Score", "process": self._process_anime_aesthetic},
        }

    def image_to_base64(self, image: Image.Image) -> str:
        """Convert PIL Image to base64 encoded JPEG string."""
        buffered = BytesIO()
        image.save(buffered, format="JPEG")
        return base64.b64encode(buffered.getvalue()).decode('utf-8')

    def auto_tune_batch_size(self, images: list) -> int:
        """Automatically determine the optimal batch size for processing."""
        batch_size = 1
        max_batch = len(images)
        test_image = images[0:1]
        while batch_size <= max_batch:
            try:
                if "aesthetic_shadow" in self.available_models and self.available_models["aesthetic_shadow"]['selected']: # Check if model is available and selected
                    _ = self.aesthetic_shadow(test_image * batch_size)
                if "waifu_scorer" in self.available_models and self.available_models["waifu_scorer"]['selected']: # Check if model is available and selected
                    _ = self.waifu_scorer(test_image * batch_size)
                if "aesthetic_predictor_v2_5" in self.available_models and self.available_models["aesthetic_predictor_v2_5"]['selected']: # Check if model is available and selected
                    _ = self.aesthetic_predictor.inference(test_image * batch_size)
                batch_size *= 2
                if batch_size > max_batch:
                    break
            except Exception:
                break
        optimal = max(1, batch_size // 2)
        if optimal > 64:
            optimal = 64
            print("Capped optimal batch size to 64")
        print(f"Optimal batch size determined: {optimal}")
        return optimal

    async def process_images_evaluation_with_logs(self, file_paths: list, auto_batch: bool, manual_batch_size: int, selected_models):
        """Asynchronously process images and yield updates with logs, HTML table, and progress bar."""
        self.results = []
        log_events = []
        images = []
        file_names = []

        # Update available models based on selection
        for model_key in self.available_models:
            self.available_models[model_key]['selected'] = model_key in selected_models

        total_files = len(file_paths)
        log_events.append(f"Starting to load {total_files} images...")
        for f in file_paths:
            try:
                img = Image.open(f).convert("RGB")
                images.append(img)
                file_names.append(os.path.basename(f))
            except Exception as e:
                log_events.append(f"Error opening {f}: {e}")

        if not images:
            log_events.append("No valid images loaded.")
            progress_percentage = 0 # Define progress_percentage here for no images case
            progress_html = self._generate_progress_html(progress_percentage)
            yield ("<p>No images loaded.</p>", "", self._format_logs(log_events), progress_html, manual_batch_size)
            return

        yield ("<p>Images loaded. Determining batch size...</p>", "", self._format_logs(log_events),
               self._generate_progress_html(0), manual_batch_size)
        await asyncio.sleep(0.1)

        try:
            manual_batch_size = int(manual_batch_size) if manual_batch_size is not None else 1
        except ValueError:
            manual_batch_size = 1
            log_events.append("Invalid manual batch size. Defaulting to 1.")

        optimal_batch = self.auto_tune_batch_size(images) if auto_batch else manual_batch_size
        log_events.append(f"Using batch size: {optimal_batch}")
        yield ("<p>Processing images in batches...</p>", "", self._format_logs(log_events),
               self._generate_progress_html(0), optimal_batch)
        await asyncio.sleep(0.1)

        total_images = len(images)
        for i in range(0, total_images, optimal_batch):
            batch_images = images[i:i+optimal_batch]
            batch_file_names = file_names[i:i+optimal_batch]
            batch_index = i // optimal_batch + 1
            log_events.append(f"Processing batch {batch_index}: images {i+1} to {min(i+optimal_batch, total_images)}")

            batch_results = {}

            # Aesthetic Shadow processing
            if self.available_models['aesthetic_shadow']['selected']:
                batch_results['aesthetic_shadow'] = await self._process_aesthetic_shadow(batch_images, log_events)
            else:
                batch_results['aesthetic_shadow'] = [None] * len(batch_images)

            # Waifu Scorer processing
            if self.available_models['waifu_scorer']['selected']:
                batch_results['waifu_scorer'] = await self._process_waifu_scorer(batch_images, log_events)
            else:
                batch_results['waifu_scorer'] = [None] * len(batch_images)

            # Aesthetic Predictor V2.5 processing
            if self.available_models['aesthetic_predictor_v2_5']['selected']:
                batch_results['aesthetic_predictor_v2_5'] = await self._process_aesthetic_predictor_v2_5(batch_images, log_events)
            else:
                batch_results['aesthetic_predictor_v2_5'] = [None] * len(batch_images)

            # Anime Aesthetic processing (single image)
            if self.available_models['anime_aesthetic']['selected']:
                batch_results['anime_aesthetic'] = await self._process_anime_aesthetic(batch_images, log_events)
            else:
                batch_results['anime_aesthetic'] = [None] * len(batch_images)


            # Combine results
            for j in range(len(batch_images)):
                scores_to_average = []
                for model_key in self.available_models:
                    if self.available_models[model_key]['selected']: # Only consider selected models
                        score = batch_results[model_key][j]
                        if score is not None:
                            scores_to_average.append(score)

                final_score = float(np.clip(np.mean(scores_to_average), 0.0, 10.0)) if scores_to_average else None
                thumbnail = batch_images[j].copy()
                thumbnail.thumbnail((200, 200))
                result = {
                    'file_name': batch_file_names[j],
                    'img_data': self.image_to_base64(thumbnail),  # Keep this for the HTML display
                    'final_score': final_score,
                }
                for model_key in self.available_models: # Add model scores to result
                    if self.available_models[model_key]['selected']:
                        result[model_key] = batch_results[model_key][j]

                self.results.append(result)
                self.sort_results() # Sort results after adding new result
                progress_percentage = min(100, ((i + len(batch_images)) / total_images) * 100) # Define progress_percentage here
                yield (f"<p>Processed batch {batch_index}.</p>", self.generate_html_table(self.results, selected_models), # Update table immediately
                       self._format_logs(log_events[-10:]), self._generate_progress_html(progress_percentage), optimal_batch)
                await asyncio.sleep(0.1)


        log_events.append("All images processed.")
        self.sort_results() # Final sort after all images processed
        html_table = self.generate_html_table(self.results, selected_models) # Pass selected models to final table generation
        final_progress = self._generate_progress_html(100)
        yield ("<p>All images processed.</p>", html_table,
               self._format_logs(log_events[-10:]), final_progress, optimal_batch)

    async def _process_aesthetic_shadow(self, batch_images, log_events):
        try:
            shadow_results = self.aesthetic_shadow(batch_images)
            log_events.append("Aesthetic Shadow processed for batch.")
        except Exception as e:
            log_events.append(f"Error in Aesthetic Shadow: {e}")
            shadow_results = [None] * len(batch_images)
        aesthetic_shadow_scores = []
        for res in shadow_results:
            try:
                hq_score = next(p for p in res if p['label'] == 'hq')['score']
                score = float(np.clip(hq_score * 10.0, 0.0, 10.0))
            except Exception:
                score = None
            aesthetic_shadow_scores.append(score)
        log_events.append("Aesthetic Shadow scores computed for batch.")
        return aesthetic_shadow_scores

    async def _process_waifu_scorer(self, batch_images, log_events):
        try:
            waifu_scores = self.waifu_scorer(batch_images)
            waifu_scores = [float(np.clip(s, 0.0, 10.0)) if s is not None else None for s in waifu_scores]
            log_events.append("Waifu Scorer processed for batch.")
        except Exception as e:
            log_events.append(f"Error in Waifu Scorer: {e}")
            waifu_scores = [None] * len(batch_images)
        return waifu_scores

    async def _process_aesthetic_predictor_v2_5(self, batch_images, log_events):
        try:
            v2_5_scores = self.aesthetic_predictor.inference(batch_images)
            v2_5_scores = [float(np.round(np.clip(s, 0.0, 10.0), 4)) if s is not None else None for s in v2_5_scores]
            log_events.append("Aesthetic Predictor V2.5 processed for batch.")
        except Exception as e:
            log_events.append(f"Error in Aesthetic Predictor V2.5: {e}")
            v2_5_scores = [None] * len(batch_images)
        return v2_5_scores

    async def _process_anime_aesthetic(self, batch_images, log_events):
        anime_scores = []
        for j, img in enumerate(batch_images):
            try:
                score = predict_anime_aesthetic(img, self.anime_aesthetic)
                anime_scores.append(float(np.clip(score * 10.0, 0.0, 10.0)))
                log_events.append(f"Anime Aesthetic processed for image {j + 1}.")
            except Exception as e:
                log_events.append(f"Error in Anime Aesthetic for image {j + 1}: {e}")
                anime_scores.append(None)
        return anime_scores


    def _generate_progress_html(self, percentage: float) -> str:
        """Generate HTML for a progress bar given a percentage."""
        return f"""
            <div style="width:100%;background-color:#ddd; border-radius:5px;">
              <div style="width:{percentage:.1f}%; background-color:#4CAF50; text-align:center; padding:5px 0; border-radius:5px;">
                {percentage:.1f}%
              </div>
            </div>
        """

    def _format_logs(self, logs: list) -> str:
        """Format log events into an HTML string."""
        return "<div style='max-height:300px; overflow-y:auto;'>" + "<br>".join(logs) + "</div>"

    def sort_results(self, sort_by: str = "Final Score") -> list:
        """Sort results based on the specified column."""
        key_map = {
            "Final Score": "final_score",
            "File Name": "file_name",
            "Aesthetic Shadow": "aesthetic_shadow",
            "Waifu Scorer": "waifu_scorer",
            "Aesthetic V2.5": "aesthetic_predictor_v2_5",
            "Anime Score": "anime_aesthetic"
        }
        key = key_map.get(sort_by, "final_score")
        reverse = sort_by != "File Name"
        self.results.sort(key=lambda r: r.get(key) if r.get(key) is not None else (-float('inf') if not reverse else float('inf')), reverse=reverse)
        return self.results

    def generate_html_table(self, results: list, selected_models) -> str:
        """Generate an HTML table to display the evaluation results."""
        table_html = """
        <style>
            .results-table { width: 100%; border-collapse: collapse; margin: 20px 0; font-family: Arial, sans-serif; }
            .results-table th, .results-table td { color: #eee; border: 1px solid #ddd; padding: 8px; text-align: center; }
            .results-table th { font-weight: bold; }
            .results-table tr:nth-child(even) { background-color: transparent; }
            .results-table tr:hover { background-color: rgba(255, 255, 255, 0.1); }
            .image-preview { max-width: 150px; max-height: 150px; display: block; margin: 0 auto; }
            .good-score { color: #0f0; font-weight: bold; }
            .bad-score { color: #f00; font-weight: bold; }
            .medium-score { color: orange; font-weight: bold; }
        </style>
        <table class="results-table">
            <thead>
                <tr>
                    <th>Image</th>
                    <th>File Name</th>
        """
        visible_models = [] # Keep track of visible model columns
        if "aesthetic_shadow" in selected_models:
            table_html += "<th>Aesthetic Shadow</th>"
            visible_models.append("aesthetic_shadow")
        if "waifu_scorer" in selected_models:
            table_html += "<th>Waifu Scorer</th>"
            visible_models.append("waifu_scorer")
        if "aesthetic_predictor_v2_5" in selected_models:
            table_html += "<th>Aesthetic V2.5</th>"
            visible_models.append("aesthetic_predictor_v2_5")
        if "anime_aesthetic" in selected_models:
            table_html += "<th>Anime Score</th>"
            visible_models.append("anime_aesthetic")
        table_html += "<th>Final Score</th>"
        table_html += "</tr></thead><tbody>"

        for result in results:
            table_html += "<tr>"
            table_html += f'<td><img src="data:image/jpeg;base64,{result["img_data"]}" class="image-preview"></td>'
            table_html += f'<td>{result["file_name"]}</td>'
            for model_key in visible_models: # Iterate through visible models only
                score = result.get(model_key)
                table_html += self._format_score_cell(score)

            score = result.get("final_score")
            table_html += self._format_score_cell(score)
            table_html += "</tr>"
        table_html += """</tbody></table>"""
        return table_html

    def _format_score_cell(self, score):
        score_str = f"{score:.4f}" if isinstance(score, (int, float)) else "N/A"
        score_class = ""
        if isinstance(score, (int, float)):
            if score >= 7:
                score_class = "good-score"
            elif score >= 5:
                score_class = "medium-score"
            else:
                score_class = "bad-score"
        return f'<td class="{score_class}">{score_str}</td>'


    def cleanup(self):
        """Clean up temporary directories."""
        if os.path.exists(self.temp_dir):
            shutil.rmtree(self.temp_dir)


#####################################
#             Interface             #
#####################################

def create_interface():
    evaluator = ImageEvaluationTool()
    sort_options = ["Final Score", "File Name", "Aesthetic Shadow", "Waifu Scorer", "Aesthetic V2.5", "Anime Score"]
    model_options = ["aesthetic_shadow", "waifu_scorer", "aesthetic_predictor_v2_5", "anime_aesthetic"]

    with gr.Blocks(theme=gr.themes.Soft()) as demo:
        gr.Markdown("""
        # Comprehensive Image Evaluation Tool

        Upload images to evaluate them using multiple aesthetic and quality prediction models.

        **New features:**
        - **Dynamic Final Score:** Final score recalculates on model selection changes.
        - **Model Selection:** Choose which models to use for evaluation.
        - **Dynamic Table Updates:** Table updates automatically based on model selection.
        - **Automatic Sorting:** Table is automatically sorted by 'Final Score'.
        - **Detailed Logs:** See major processing events (limited to the last 10).
        - **Progress Bar:** Visual indication of processing status.
        - **Asynchronous Updates:** Streaming status and logs during processing.
        - **Batch Size Controls:** Choose manual batch size or let the tool auto-detect it.
        - **Download Results:** Export the evaluation results as CSV.
        """)

        with gr.Row():
            with gr.Column(scale=1):
                input_images = gr.Files(label="Upload Images", file_count="multiple")
                model_checkboxes = gr.CheckboxGroup(model_options, label="Select Models", value=model_options, info="Choose models for evaluation.")
                auto_batch_checkbox = gr.Checkbox(label="Automatic Batch Size Detection", value=False, info="Enable to automatically determine the optimal batch size.")
                batch_size_input = gr.Number(label="Batch Size", value=1, interactive=True, info="Manually specify the batch size if auto-detection is disabled.")
                sort_dropdown = gr.Dropdown(sort_options, value="Final Score", label="Sort by", info="Select the column to sort results by.")
                process_btn = gr.Button("Evaluate Images", variant="primary")
                clear_btn = gr.Button("Clear Results")
                download_csv = gr.Button("Download CSV", variant="secondary")

            with gr.Column(scale=2):
                progress_bar = gr.HTML(label="Progress Bar", value="""
                <div style='width:100%;background-color:#ddd;'>
                  <div style='width:0%;background-color:#4CAF50;padding:5px 0;text-align:center;'>0%</div>
                </div>
                """)
                log_window = gr.HTML(label="Detailed Logs", value="<div style='max-height:300px; overflow-y:auto;'>Logs will appear here...</div>")
                status_html = gr.HTML(label="Status")
                output_html = gr.HTML(label="Evaluation Results")
                download_file_output = gr.File() # Initialize gr.File component without filename

        # Function to convert results to CSV format, excluding 'img_data'.
        def results_to_csv(selected_models):
            import csv
            import io
            if not evaluator.results:
                return None # Return None when no results are available
            output = io.StringIO()
            fieldnames = ['file_name', 'final_score'] # Base fieldnames
            for model_key in selected_models: # Add selected model names as fieldnames
                if model_key in selected_models: # Double check if model_key is indeed in selected_models list
                    fieldnames.append(model_key)

            writer = csv.DictWriter(output, fieldnames=fieldnames)
            writer.writeheader()
            for res in evaluator.results:
                row_dict = {'file_name': res['file_name'], 'final_score': res['final_score']} # Base data
                for model_key in selected_models: # Add selected model scores
                    if model_key in selected_models: # Double check before accessing res[model_key]
                        row_dict[model_key] = res.get(model_key, 'N/A') # Use get with default 'N/A' if model not in result (shouldn't happen but for safety)
                writer.writerow(row_dict)
            return output.getvalue()


        def update_batch_size_interactivity(auto_batch):
            return gr.update(interactive=not auto_batch)

        async def process_images_and_update(files, auto_batch, manual_batch, selected_models):
            file_paths = [f.name for f in files]
            async for status, table, logs, progress, updated_batch in evaluator.process_images_evaluation_with_logs(file_paths, auto_batch, manual_batch, selected_models):
                yield status, table, logs, progress, gr.update(value=updated_batch, interactive=not auto_batch)

        def update_table_sort(sort_by_column, selected_models):
            sorted_results = evaluator.sort_results(sort_by_column)
            return evaluator.generate_html_table(sorted_results, selected_models)

        def update_table_model_selection(selected_models):
            # Recalculate final scores based on selected models
            for result in evaluator.results:
                scores_to_average = []
                for model_key in evaluator.available_models:
                    if model_key in selected_models and evaluator.available_models[model_key]['selected']: # consider only selected models from checkbox group and available_models
                        score = result.get(model_key)
                        if score is not None:
                            scores_to_average.append(score)
                final_score = float(np.clip(np.mean(scores_to_average), 0.0, 10.0)) if scores_to_average else None
                result['final_score'] = final_score

            sorted_results = evaluator.sort_results() # Keep sorting by Final Score when models change
            return evaluator.generate_html_table(sorted_results, selected_models)


        def clear_results():
            evaluator.results = []
            return (gr.update(value=""),
                    gr.update(value=""),
                    gr.update(value=""),
                    gr.update(value="""
                    <div style='width:100%;background-color:#ddd;'>
                      <div style='width:0%;background-color:#4CAF50;padding:5px 0;text-align:center;'>0%</div>
                    </div>
                    """),
                    gr.update(value=1))

        def download_results_csv_trigger(selected_models): # Changed function name to avoid conflict and clarify purpose
            csv_content = results_to_csv(selected_models)
            if csv_content is None:
                return None # Indicate no file to download

            # Create a temporary file to save the CSV data
            with tempfile.NamedTemporaryFile(suffix=".csv", delete=False) as tmp_file:
                tmp_file.write(csv_content.encode())
                temp_file_path = tmp_file.name # Get the path to the temporary file

            return temp_file_path # Return the path to the temporary file


        auto_batch_checkbox.change(
            update_batch_size_interactivity,
            inputs=[auto_batch_checkbox],
            outputs=[batch_size_input]
        )

        process_btn.click(
            process_images_and_update,
            inputs=[input_images, auto_batch_checkbox, batch_size_input, model_checkboxes],
            outputs=[status_html, output_html, log_window, progress_bar, batch_size_input]
        )
        sort_dropdown.change(
            update_table_sort,
            inputs=[sort_dropdown, model_checkboxes],
            outputs=[output_html]
        )
        model_checkboxes.change( # Added change event for model checkboxes
            update_table_model_selection,
            inputs=[model_checkboxes],
            outputs=[output_html]
        )
        clear_btn.click(
            clear_results,
            inputs=[],
            outputs=[status_html, output_html, log_window, progress_bar, batch_size_input]
        )
        download_csv.click(
            download_results_csv_trigger, # Call the trigger function
            inputs=[model_checkboxes],
            outputs=[download_file_output] # Output is now the gr.File component
        )
        demo.load(lambda: update_table_sort("Final Score", model_options), inputs=None, outputs=[output_html]) # Initial sort and table render
        gr.Markdown("""
        ### Notes
        - Select models to use for evaluation using the checkboxes.
        - The 'Final Score' recalculates dynamically when models are selected/deselected.
        - The table updates automatically when models are selected/deselected and is always sorted by 'Final Score'.
        - The log window displays the most recent 10 events.
        - The progress bar shows overall processing status.
        - When 'Automatic Batch Size Detection' is enabled, the batch size field becomes disabled.
        - Use the download button to export your evaluation results as CSV.
        """)

    return demo

if __name__ == "__main__":
    demo = create_interface()
    demo.queue().launch()