Spaces:
Sleeping
Sleeping
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"]) | |