|
import os |
|
from transformers import BlipProcessor, BlipForConditionalGeneration |
|
import gradio as gr |
|
|
|
|
|
HUGGINGFACE_TOKEN = os.getenv("Image_classification") |
|
|
|
|
|
processor = BlipProcessor.from_pretrained( |
|
"quadranttechnologies/qhub-blip-image-captioning-finetuned", |
|
use_auth_token=HUGGINGFACE_TOKEN |
|
) |
|
model = BlipForConditionalGeneration.from_pretrained( |
|
"quadranttechnologies/qhub-blip-image-captioning-finetuned", |
|
use_auth_token=HUGGINGFACE_TOKEN |
|
) |
|
|
|
|
|
def generate_caption(image): |
|
try: |
|
|
|
inputs = processor(image, return_tensors="pt") |
|
|
|
|
|
outputs = model.generate(**inputs) |
|
caption = processor.decode(outputs[0], skip_special_tokens=True) |
|
return caption |
|
except Exception as e: |
|
return f"Error generating caption: {e}" |
|
|
|
|
|
interface = gr.Interface( |
|
fn=generate_caption, |
|
inputs=gr.Image(type="pil"), |
|
outputs="text", |
|
title="Image Captioning Model", |
|
description="Upload an image to generate a caption using the fine-tuned BLIP model." |
|
) |
|
|
|
|
|
if __name__ == "__main__": |
|
interface.launch(share=True) |
|
|
|
|
|
|
|
|