import gradio as gr

import torch
from PIL import Image
from transformers import CLIPModel, CLIPProcessor

LIST_LABELS = ['agricultural land', 'airplane', 'baseball diamond', 'beach', 'buildings', 'chaparral', 'dense residential area', 'forest', 'freeway', 'golf course', 'harbor', 'intersection', 'medium residential area', 'mobilehome park', 'overpass', 'parking lot', 'river', 'runway', 'sparse residential area', 'storage tanks', 'tennis court']

CLIP_LABELS = [f"A satellite image of {label}" for label in LIST_LABELS]

MODEL_NAME = "NemesisAlm/clip-fine-tuned-satellite"

device = "cuda" if torch.cuda.is_available() else "cpu"
clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device)
clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")

fine_tuned_model = CLIPModel.from_pretrained(MODEL_NAME).to(device)
fine_tuned_processor = CLIPProcessor.from_pretrained(MODEL_NAME)


def classify(image_path, model_number):
    if model_number == "CLIP":
        processor = clip_processor
        model = clip_model
    else:
        processor = fine_tuned_processor
        model = fine_tuned_model
    image = Image.open(image_path).convert('RGB')
    inputs = processor(text=CLIP_LABELS, images=image, return_tensors="pt", padding=True).to(device)
    with torch.no_grad():
        outputs = model(**inputs)
    logits_per_image = outputs.logits_per_image
    prediction = logits_per_image.softmax(dim=1)
    confidences = {LIST_LABELS[i]: float(prediction[0][i].item()) for i in range(len(LIST_LABELS))}
    return confidences

DESCRIPTION="""
<div style="font-family: Arial, sans-serif; line-height: 1.6;  margin: auto; text-align: center;">
    <h2 style="color: #333;">CLIP Fine-Tuned Satellite Model Demo</h2>
    <p>
        This space demonstrates the capabilities of a <strong>fine-tuned CLIP-based model</strong> 
        in classifying satellite images. The model has been specifically trained on the 
        <em>UC Merced</em> satellite image dataset.
    </p>
    <p>
        After just <strong>2 epochs of training</strong>, adjusting only 30% of the model parameters, 
        the model's accuracy in classifying satellite images has significantly improved, from an 
        initial accuracy of <strong>58.8%</strong> to <strong>96.9%</strong> on the test set.
    </p>
    <p>
        Explore this space to see its performance and compare it with the initial CLIP model.
    </p>
</div>
"""

FOOTER = """
<div style="margin-top:50px">
    Link to model: <a href='https://huggingface.co/NemesisAlm/clip-fine-tuned-satellite'>https://huggingface.co/NemesisAlm/clip-fine-tuned-satellite</a><br>
    Link to dataset: <a href='https://huggingface.co/datasets/blanchon/UC_Merced'>https://huggingface.co/datasets/blanchon/UC_Merced</a>
</div>
"""

with gr.Blocks(title="Satellite image classification", css="") as demo:
    logo = gr.HTML("<img src='file/logo_gradio.png' style='margin:auto'/>")
    description = gr.HTML(DESCRIPTION)
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(type='filepath', label='Input image')
            submit_btn = gr.Button("Submit", variant="primary")
        with gr.Column():
            title_1 = gr.HTML("<h1 style='text-align:center'>Original CLIP Model</h1>")
            model_1 = gr.Textbox("CLIP", visible=False)
            output_labels_clip = gr.Label(num_top_classes=10, label="Top 10 classes")
        with gr.Column():
            title_2 = gr.HTML("<h1 style='text-align:center'>Fine-tuned Model</h1>")
            model_2 = gr.Textbox("Fine-tuned", visible=False)
            output_labels_finetuned = gr.Label(num_top_classes=10, label="Top 10 classes")
    examples = gr.Examples([["0.jpg"], ["1.jpg"], ["2.jpg"], ["3.jpg"]   ], input_image)
    footer = gr.HTML(FOOTER)
    submit_btn.click(fn=classify, inputs=[input_image, model_1], outputs=output_labels_clip).then( classify, inputs=[input_image, model_2], outputs=[output_labels_finetuned]  )


demo.queue()
demo.launch(server_name="0.0.0.0",favicon_path='favicon.ico', allowed_paths=["logo_gradio.png", "0.jpg", "1.jpg", "2.jpg", "3.jpg"])