Spaces:
Sleeping
Sleeping
import gradio as gr | |
import torch | |
import time | |
from PIL import Image | |
from transformers import BlipProcessor, BlipForConditionalGeneration | |
from utils import create_pdf | |
# Load model and processor | |
processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base") | |
model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base") | |
device = "cuda" if torch.cuda.is_available() else "cpu" | |
model.to(device) | |
def generate_caption(image): | |
start_time = time.time() | |
if image.mode != "RGB": | |
image = image.convert("RGB") | |
inputs = processor(images=image, return_tensors="pt").to(device) | |
output = model.generate(**inputs, max_new_tokens=50) | |
caption = processor.decode(output[0], skip_special_tokens=True) | |
duration = time.time() - start_time | |
if duration > 5: | |
caption = f"⚠️ Took {round(duration, 2)}s: {caption}" | |
return caption | |
def process_images(images): | |
results = [] | |
for i, img in enumerate(images[:10]): # Limit to 10 images | |
caption = generate_caption(img) | |
results.append(f"Image {i+1}: {caption}") | |
pdf_file = create_pdf(results) | |
return "\n\n".join(results), pdf_file | |
iface = gr.Interface( | |
fn=process_images, | |
inputs=gr.File(label="Upload up to 10 Site Images", type="file", file_types=[".jpg", ".png"], multiple=True), | |
outputs=["text", "file"], | |
title="Auto-DPR Generator from Site Images", | |
description="Upload construction site images. AI will auto-generate a progress summary and downloadable PDF.", | |
allow_flagging="never" | |
) | |
if __name__ == "__main__": | |
iface.launch() | |