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