Spaces:
Sleeping
Sleeping
File size: 3,858 Bytes
70de5aa ae059f3 70de5aa b6d576b 70de5aa ae059f3 70de5aa ae059f3 70de5aa 0a3520e 465440b 0a3520e 465440b 0a3520e 70de5aa ae059f3 70de5aa 77f055e 70de5aa ae059f3 70de5aa ae059f3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 |
import requests
from PIL import Image, UnidentifiedImageError
from transformers import AutoTokenizer, AutoImageProcessor, VisionEncoderDecoderModel
import gradio as gr
import os
# Load the model, tokenizer, and image processor with error handling
def load_model_and_components(model_name):
try:
model = VisionEncoderDecoderModel.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
image_processor = AutoImageProcessor.from_pretrained(model_name)
return model, tokenizer, image_processor
except Exception as e:
raise RuntimeError(f"Error loading model components: {e}")
current_model_name = "laicsiifes/swin-distilbertimbau"
model, tokenizer, image_processor = load_model_and_components(current_model_name)
# Function to process the image and generate a caption
def generate_caption(image):
try:
pixel_values = image_processor(image, return_tensors="pt").pixel_values
generated_ids = model.generate(pixel_values)
caption = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
return caption
except Exception:
return "Please upload a valid image."
# Predefined images for selection
image_folder = "images"
predefined_images_paths = [
os.path.join(image_folder, fname) for fname in os.listdir(image_folder) if fname.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.gif'))
]
# Gradio app
def app(image=None):
try:
if image is None:
return "Please upload a valid image."
return generate_caption(image)
except Exception:
return "Please upload a valid image."
# Define UI
with gr.Blocks() as interface:
gr.Markdown("""
# Welcome to the LAICSI-IFES space for Vision Encoder-Decoder (VED) demonstration
---
### Be patient with the Swin-GPorTuguese-2 as it is heavier than the Swin-DistilBERTimbau.
""")
with gr.Row():
with gr.Column():
model_selector = gr.Dropdown(choices=["laicsiifes/swin-distilbertimbau", "laicsiifes/swin-gportuguese-2"],
value="laicsiifes/swin-distilbertimbau",
label="Select Model")
loading_message = gr.Textbox(label="Status Message")
image_display = gr.Image(type="pil", label="Image Preview", interactive=False)
upload_button = gr.File(label="Upload an Image", file_types=["image"], type="filepath")
examples = gr.Examples(predefined_images_paths, inputs=[upload_button], label="Examples")
with gr.Column():
output_text = gr.Textbox(label="Generated Caption")
# Define logic
def handle_uploaded_image(image):
try:
if image is None:
return None, "Please upload a valid image."
pil_image = Image.open(image).convert("RGB")
return pil_image, generate_caption(pil_image)
except Exception:
return None, "Please upload a valid image."
def switch_model(selected_model):
gr.Info("Loading model... Please wait.")
return "Loading model... Please wait.", None, None, None
def load_new_model(selected_model):
global model, tokenizer, image_processor
model, tokenizer, image_processor = load_model_and_components(selected_model)
return "Model loaded successfully.", None, None, None
model_selector.change(fn=switch_model, inputs=model_selector, outputs=[loading_message, upload_button, image_display, output_text])
model_selector.change(fn=load_new_model, inputs=model_selector, outputs=[loading_message, image_display, output_text])
upload_button.change(fn=handle_uploaded_image, inputs=upload_button, outputs=[image_display, output_text])
interface.launch(share=False)
|