import gradio as gr
from transformers import pipeline
import torch
import os
from PIL import Image
import io

# Check if CUDA is available (Hugging Face Spaces supports GPU acceleration)
device = 0 if torch.cuda.is_available() else -1
print(f"Using device: {'CUDA' if device == 0 else 'CPU'}")

# Initialize the image processing pipeline
# You can replace this with any Hugging Face model that processes images
model_name = "google/vit-base-patch16-224"
image_processor = pipeline("image-classification", model=model_name, device=device)

def process_image(input_image):
    """
    Process the uploaded image through the model and return results
    """
    if input_image is None:
        return [{"label": "No image provided", "score": 0.0}]
    
    # Run the image through the model
    results = image_processor(input_image)
    
    # Return top 5 predictions
    return results[:5]

def save_output(results):
    """
    Convert results to a downloadable format
    """
    if not results or len(results) == 0:
        return None
        
    output_text = "Model Predictions:\n\n"
    for result in results:
        output_text += f"Label: {result['label']}, Score: {result['score']:.4f}\n"
    
    # Create a file for download
    with open("results.txt", "w") as f:
        f.write(output_text)
    
    return "results.txt"

# Create the Gradio interface with a more polished design
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# Image Classification Demo")
    gr.Markdown("Upload an image and get classification results from the model.")
    
    with gr.Row():
        with gr.Column(scale=1):
            # Input components
            input_image = gr.Image(type="pil", label="Upload Image")
            with gr.Row():
                submit_btn = gr.Button("Process Image", variant="primary")
                clear_btn = gr.Button("Clear")
        
        with gr.Column(scale=1):
            # Output components
            output_results = gr.JSON(label="Model Predictions")
            download_btn = gr.Button("Download Results")
            download_output = gr.File(label="Download Output")
    
    # Set up the processing flow
    submit_btn.click(
        fn=process_image,
        inputs=[input_image],
        outputs=[output_results]
    )
    
    clear_btn.click(
        fn=lambda: (None, None, None),
        inputs=[],
        outputs=[input_image, output_results, download_output]
    )
    
    download_btn.click(
        fn=save_output,
        inputs=[output_results],
        outputs=[download_output]
    )
    
    # Add example images
    gr.Examples(
        examples=[
            os.path.join(os.path.dirname(__file__), "examples/cat.jpg"),
            os.path.join(os.path.dirname(__file__), "examples/dog.jpg"),
        ],
        inputs=input_image,
        label="Example Images"
    )
    
    # Add information footer
    gr.Markdown("""
    ### About this demo
    - Model: [google/vit-base-patch16-224](https://huggingface.co/google/vit-base-patch16-224)
    - This demo classifies images into 1000 ImageNet categories
    - Created with Gradio and Hugging Face Transformers
    """)

# For Hugging Face Spaces, we use the Gradio app directly
demo.launch()