DPR / app.py
Rammohan0504's picture
Update app.py
6b0a154 verified
raw
history blame
1.63 kB
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()