app-Copy1.py

import gradio as gr
import numpy as np
from PIL import Image
import os
import requests
import json
from dotenv import load_dotenv
import openai
import base64
import csv
import tempfile
import datetime

# import libraries
from library.utils_model import *
from library.utils_html import *
from library.utils_prompt import *

OR = OpenRouterAPI()

# Define model pricing information (approximate costs per 100 image API calls)
MODEL_PRICING = {
    "google/gemini-2.0-flash-001": "$0.03",
    "gpt-4.1-mini": "$0.07",
    "gpt-4.1": "$0.35",
    "anthropic/claude-3.7-sonnet": "$0.70",
    "google/gemini-2.5-pro-preview-03-25": "$1.20",
    "google/gemini-2.5-flash-preview:thinking": "$0.35",
    "gpt-4.1-nano": "$0.02",
    "openai/chatgpt-4o-latest": "$0.75",
    "meta-llama/llama-4-maverick": "$0.04"
}

def get_sys_prompt(length="medium", photograph=False):
    extra_prompt = ""

    if photograph:
        object_type = "wildlife photography"
        extra_prompt = " Do not guess the exact species of the animals in the photograph unless you are certain - simply use a broader terms to make less errors e.g. say Swan rather mute Swan or Whooper Swan unless you are certain."
    else:
        object_type = "museum objects"
        
    dev_prompt = f"""You are a museum curator tasked with generating long descriptions (as defined by W3C) of {object_type} for visually impaired and blind users from images. Use British English and follow museum accessibility best practices. Do not start with phrases like 'The image shows', 'This is an image of', 'The photograph'. Be precise, concise and avoid filler and subjective statements."""
    
    if length == "short":
        dev_prompt = f"""You are a museum curator tasked with generating alt-text (as defined by W3C) of {object_type} for visually impaired and blind users from images. Use British English and follow museum accessibility best practices. Do not start with phrases like 'The image shows' or 'This is an image of'. Be precise, concise and avoid filler and subjective statements. Repsonses should be a maximum of 130 characters."""
    elif length == "medium":
        dev_prompt += " Repsonses should be a maximum of 250-300 characters."
    else: # long
        dev_prompt += " Repsonses should be a maximum of 450 characters."
    return dev_prompt + extra_prompt

def create_csv_file_simple(results):
    """Create a CSV file from the results and return the path"""
    try:
        with tempfile.NamedTemporaryFile(mode='w', suffix='.csv', delete=False, newline='', encoding='utf-8') as f:
            path = f.name
            writer = csv.writer(f)
            writer.writerow(['image_id', 'content'])
            for result in results:
                writer.writerow([
                    result.get('image_id', ''),
                    result.get('content', '')
                ])
        return path
    except Exception as e:
        print(f"Error creating CSV: {e}")
        return None

def get_base_filename(filepath):
    if not filepath:
        return ""
    basename = os.path.basename(filepath)
    filename = os.path.splitext(basename)[0]
    return filename

# Define the Gradio interface
def create_demo():
    custom_css = """
    /* Container for the image component (#current-image-display is the elem_id of gr.Image) */
        #current-image-display {
            height: 600px;           /* Define container height */
            width: 100%;             /* Define container width (takes column width) */
            display: flex;           /* Use flexbox for alignment */
            justify-content: center; /* Center content horizontally */
            align-items: center;     /* Center content vertically */
            overflow: hidden;        /* Hide any potential overflow from container */
        }
    
        /* The actual <img> element inside the container */
        #current-image-display img {
            object-fit: contain !important;    /* Scale keeping aspect ratio, within bounds */
            max-width: 100%;        /* Prevent image exceeding container width */
            max-height: 600px !important;       /* Prevent image exceeding container height */
            width: auto;            /* Use natural width unless constrained by max-width */
            height: auto;           /* Use natural height unless constrained by max-height */
            display: block;         /* Ensure image behaves predictably in flex */
        }
        
        /* Custom style for model info display */
        #model-info-display {
            font-size: 0.85rem;     /* Small font size */
            color: #666;            /* Subtle color */
            margin-top: 0.5rem;     /* Small top margin */
            margin-bottom: 1rem;    /* Bottom margin before next element */
            padding-left: 0.5rem;   /* Slight indentation */
        }
    """
    # --- Pass css to gr.Blocks ---
    with gr.Blocks(theme=gr.themes.Monochrome(), css=custom_css) as demo:
        with gr.Row():
            with gr.Column(scale=3):
                gr.Markdown("# MATCHA: Museum Alt-Text for Cultural Heritage with AI 🍵 🌿")
                gr.Markdown("Upload one or more images to generate accessible alternative text (designed to meet WCAG Guidelines)")
                gr.Markdown("Developed by the Natural History Museum in Partnership with National Museums Liverpool. Funded by the DCMS Pilot Scheme")
            with gr.Column(scale=1):
                with gr.Row():
                    gr.Image("images/nhm_logo.png", show_label=False, height=100,
                             interactive=False, show_download_button=False,
                             show_share_button=False, show_fullscreen_button=False,
                             container=False, elem_id="nhm-logo")
                    gr.Image("images/nml_logo.png", show_label=False, height=100,
                             interactive=False, show_download_button=False,
                             show_share_button=False, show_fullscreen_button=False,
                             container=False, elem_id="nml-logo")

        # Store model choices and state
        show_all_models_state = gr.State(False)

        # Define preferred and additional models directly in the function
        preferred_models = [
            ("Gemini 2.0 Flash (cheap)", "google/gemini-2.0-flash-001"),
            ("GPT-4.1 Mini", "gpt-4.1-mini"),
            ("GPT-4.1 (Recommended)", "gpt-4.1"),
            ("Claude 3.7 Sonnet", "anthropic/claude-3.7-sonnet"),
            ("Gemini 2.5 Pro", "google/gemini-2.5-pro-preview-03-25"),
            ("Gemini 2.5 Flash Thinking (Recommended)", "google/gemini-2.5-flash-preview:thinking")
        ]
        
        additional_models = [
            ("GPT-4.1 Nano", "gpt-4.1-nano"),
            ("ChatGPT Latest", "openai/chatgpt-4o-latest"),
            ("Llama 4 Maverick", "meta-llama/llama-4-maverick")
        ]
        
        # Calculate all models once
        all_models_list = preferred_models + additional_models
        
        # Default model value
        default_model = "google/gemini-2.0-flash-001"
        
        with gr.Row():
            # Left column: Controls and uploads
            with gr.Column(scale=1):
                upload_button = gr.UploadButton(
                    "Click to Upload Images",
                    file_types=["image"],
                    file_count="multiple"
                )
                
                # Model dropdown
                model_choice = gr.Dropdown(
                    choices=preferred_models, 
                    label="Select Model",
                    value=default_model
                )
                
                length_choice = gr.Radio(
                    choices=["short", "medium", "long"], 
                    label="Response Length",
                    value="medium", 
                    info="Short: max 130 chars | Medium: 250-300 chars | Long: max 450 chars"
                )
                
                # Advanced settings accordion
                with gr.Accordion("Advanced Settings", open=False):
                    show_all_models = gr.Checkbox(
                        label="Show Additional Models", 
                        value=False,
                        info="Display additional model options in the dropdown above"
                    )
                    
                    content_type = gr.Radio(
                        choices=["Museum Object", "Photography"], 
                        label="Content Type",
                        value="Museum Object"
                    )
                
                # Find the default model's display name
                default_model_name = "Unknown Model"
                for name, value in preferred_models:
                    if value == default_model:
                        default_model_name = name
                        break
                
                model_info = gr.Markdown(
                    f"""**Current Model**: {default_model_name}  
                    **Estimated cost per 100 Images**: {MODEL_PRICING[default_model]}""",
                    elem_id="model-info-display"
                )
                
                gr.Markdown("### Uploaded Images")
                input_gallery = gr.Gallery(
                    label="Uploaded Image Previews", columns=3, height=150,
                    object_fit="contain", show_label=False
                )
                analyze_button = gr.Button("Generate Alt-Text", variant="primary", size="lg")
                image_state = gr.State([])
                filename_state = gr.State([])
                csv_download = gr.File(label="Download CSV Results")

            # Right column: Display area
            with gr.Column(scale=2):
                current_image = gr.Image(
                    label="Current Image",
                    type="filepath",
                    elem_id="current-image-display",
                    show_fullscreen_button=True,
                    show_download_button=False,
                    show_share_button=False,
                    show_label=False
                 )

                with gr.Row():
                    prev_button = gr.Button("← Previous", size="sm")
                    image_counter = gr.Markdown("0 of 0", elem_id="image-counter")
                    next_button = gr.Button("Next →", size="sm")

                gr.Markdown("### Generated Alt-text")
                analysis_text = gr.Textbox(
                    label="Generated Text",
                    value="Upload images and click 'Generate Alt-Text'.",
                    lines=6, max_lines=10, interactive=True, show_label=False
                )
                current_index = gr.State(0)
                all_images = gr.State([])
                all_results = gr.State([])
        
        # Handle checkbox change to update model dropdown - modern version
        def toggle_models(show_all, current_model):
            # Make a fresh copy of the models lists to avoid any reference issues
            preferred_choices = list(preferred_models)
            all_choices = list(all_models_list)
            
            if show_all:
                # When showing all models, use the fresh copy of all models
                return gr.Dropdown(choices=all_choices, value=current_model)
            else:
                # Check if current model is in preferred models list
                preferred_values = [value for _, value in preferred_choices]
                
                if current_model in preferred_values:
                    # Keep the current model if it's in preferred models
                    return gr.Dropdown(choices=preferred_choices, value=current_model)
                else:
                    # Reset to default model if current model is not in preferred models
                    return gr.Dropdown(choices=preferred_choices, value=default_model)

        # Update model info when model selection changes
        def update_model_info(model_value):
            # Find display name
            model_name = "Unknown Model"
            for name, value in all_models_list:
                if value == model_value:
                    model_name = name
                    break
            
            # Get cost
            cost = MODEL_PRICING.get(model_value, "Unknown")
            
            # Create markdown
            return f"""**Current Model**: {model_name}  
                   **Estimated cost per 100 Images**: {cost}"""
                
        # Connect checkbox to toggle model choices
        show_all_models.change(
            fn=toggle_models,
            inputs=[show_all_models, model_choice],
            outputs=[model_choice]
        )
        
        # Connect model selection to update info
        model_choice.change(
            fn=update_model_info,
            inputs=[model_choice],
            outputs=[model_info]
        )
        
        # Handle file uploads
        def handle_upload(files, current_paths, current_filenames):
            file_paths = []
            file_names = []
            if files:
                for file in files:
                    file_paths.append(file.name)
                    file_names.append(get_base_filename(file.name))
            return file_paths, file_paths, file_names, 0, None, "0 of 0", "Upload images and click 'Generate Alt-Text'."

        upload_button.upload(
            fn=handle_upload,
            inputs=[upload_button, image_state, filename_state],
            outputs=[input_gallery, image_state, filename_state,
                     current_index, current_image, image_counter, analysis_text]
        )

        # Analyze images
        def analyze_images(image_paths, model_choice, length_choice, filenames, content_type_choice):
            if not image_paths:
                return [], [], 0, None, "0 of 0", "No images uploaded to analyze.", None

            is_photography = content_type_choice == "Photography"
            sys_prompt = get_sys_prompt(length_choice, photograph=is_photography)
            image_results = []
            analysis_progress = gr.Progress(track_tqdm=True)

            for i, image_path in enumerate(analysis_progress.tqdm(image_paths, desc="Analyzing Images")):
                image_id = filenames[i] if i < len(filenames) and filenames[i] else f"Image_{i+1}_{os.path.basename(image_path)}"
                try:
                    img = Image.open(image_path)
                    prompt0 = prompt_new()
                    model_name = model_choice
                    client_to_use = OR # Default client

                    result = client_to_use.generate_caption(
                        img, model=model_name, max_image_size=512,
                        prompt=prompt0, prompt_dev=sys_prompt, temperature=1
                    )
                    image_results.append({"image_id": image_id, "content": result.strip()})
                except FileNotFoundError:
                     error_message = f"Error: File not found at path '{image_path}'"
                     print(error_message)
                     image_results.append({"image_id": image_id, "content": error_message})
                except Exception as e:
                    error_message = f"Error processing {image_id}: {str(e)}"
                    print(error_message)
                    image_results.append({"image_id": image_id, "content": error_message})

            csv_path = create_csv_file_simple(image_results)
            initial_image = image_paths[0] if image_paths else None
            initial_counter = f"1 of {len(image_paths)}" if image_paths else "0 of 0"
            initial_text = image_results[0]["content"] if image_results else "Analysis complete, but no results generated."

            return (image_paths, image_results, 0, initial_image, initial_counter,
                    initial_text, csv_path)

        # Navigate previous
        def go_to_prev(current_idx, images, results):
            if not images or not results or len(images) == 0:
                return current_idx, None, "0 of 0", ""
            new_idx = (current_idx - 1 + len(images)) % len(images)
            counter_text = f"{new_idx + 1} of {len(images)}"
            result_content = results[new_idx]["content"] if new_idx < len(results) else "Error: Result not found"
            return (new_idx, images[new_idx], counter_text, result_content)

        # Navigate next
        def go_to_next(current_idx, images, results):
            if not images or not results or len(images) == 0:
                return current_idx, None, "0 of 0", ""
            new_idx = (current_idx + 1) % len(images)
            counter_text = f"{new_idx + 1} of {len(images)}"
            result_content = results[new_idx]["content"] if new_idx < len(results) else "Error: Result not found"
            return (new_idx, images[new_idx], counter_text, result_content)

        # Connect analyze button
        analyze_button.click(
            fn=analyze_images,
            inputs=[image_state, model_choice, length_choice, filename_state, content_type],
            outputs=[all_images, all_results, current_index, current_image, image_counter,
                     analysis_text, csv_download]
        )

        # Connect navigation buttons
        prev_button.click(
            fn=go_to_prev, inputs=[current_index, all_images, all_results],
            outputs=[current_index, current_image, image_counter, analysis_text], queue=False
        )
        next_button.click(
            fn=go_to_next, inputs=[current_index, all_images, all_results],
            outputs=[current_index, current_image, image_counter, analysis_text], queue=False
        )

        # About section
        with gr.Accordion("About", open=False):
             gr.Markdown("""
            ## About MATCHA 🍵:
            
            This demo generates alternative text for images.
            
            - Upload one or more images using the upload button
            - Choose a model and response length for generation
            - Navigate through the images with the Previous and Next buttons
            - Download CSV with all results
            
            Developed by the Natural History Museum in Partnership with National Museums Liverpool.
            
            If you find any bugs/have any problems/have any suggestions please feel free to get in touch:
            chris.addis@nhm.ac.uk
            """)

    return demo

# Launch the app
if __name__ == "__main__":
    app = create_demo()
    app.launch()