Update app.py

app.py

@@ -1,38 +1,32 @@
-import gradio as gr
 from transformers import BlipProcessor, BlipForConditionalGeneration
 from PIL import Image
-from fpdf import FPDF
-import os
-from datetime import datetime
+import gradio as gr
+import torch
 
+# Load BLIP model and processor
 processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
 model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
+model.eval()
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device)
 
-def analyze_image(image):
-    raw_image = Image.fromarray(image)
-    text = "Describe the construction site"
-    inputs = processor(raw_image, text, return_tensors="pt")
-    out = model.generate(**inputs)
-    caption = processor.decode(out[0], skip_special_tokens=True)
-
-    date_str = datetime.now().strftime("%Y-%m-%d")
-    pdf = FPDF()
-    pdf.add_page()
-    pdf.set_font("Arial", size=12)
-    pdf.multi_cell(0, 10, f"Daily Progress Report - {date_str}\n\nCaption: {caption}")
+# Inference function
+def generate_caption(image):
+    if image.mode != "RGB":
+        image = image.convert("RGB")
     
-    os.makedirs("reports", exist_ok=True)
-    file_path = f"reports/DPR_{date_str}.pdf"
-    pdf.output(file_path)
-    return caption, file_path
+    inputs = processor(image, return_tensors="pt").to(device, torch.float16)
+    output = model.generate(**inputs, max_new_tokens=50)
+    caption = processor.decode(output[0], skip_special_tokens=True)
+    return caption
 
-demo = gr.Interface(
-    fn=analyze_image,
-    inputs=gr.Image(type="numpy", label="Upload Site Photo"),
-    outputs=[gr.Textbox(label="Generated Caption"), gr.File(label="Download DPR PDF")],
-    title="Auto DPR Generator",
-    description="Upload a construction site image to generate a Daily Progress Report."
+# Gradio interface
+iface = gr.Interface(
+    fn=generate_caption,
+    inputs=gr.Image(type="pil"),
+    outputs="text",
+    title="Construction Site Image-to-Text Generator",
+    description="Upload a site photo. The model will detect and describe construction activities."
 )
 
-if __name__ == "__main__":
-    demo.launch()
+iface.launch()