Update app.py

app.py

@@ -6,15 +6,17 @@ from PIL import Image
 from collections import deque
 import numpy as np
 
-# Load BLIP model
+# Load main BLIP model for English captioning
 processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
 model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
 
-# Load YOLOv5 model
+# Load YOLOv5 small model for object detection using ultralytics package
 detect_model = YOLO('yolov5s.pt')
 
+# Session memory for last 15 images and captions
 MEMORY_SIZE = 15
-last_texts = deque([], maxlen=MEMORY_SIZE)
+last_images = deque([], maxlen=MEMORY_SIZE)
+last_captions = deque([], maxlen=MEMORY_SIZE)
 
 def preprocess_image(image):
     if image.mode != "RGB":
@@ -35,62 +37,41 @@ def detect_objects(image):
 def generate_caption(image):
     image = preprocess_image(image)
     inputs = processor(image, return_tensors="pt")
-
     out = model.generate(**inputs, max_length=30, num_beams=5, early_stopping=True)
     caption = processor.decode(out[0], skip_special_tokens=True)
     detected_objs = detect_objects(image)
-    tags = ", ".join(detected_objs) if detected_objs else "None"
-
-    combined_text = f"Detected objects: {tags}\nCaption: {caption}"
 
-    last_texts.append(combined_text)
+    # Update session memory
+    last_images.append(image)
+    last_captions.append(caption)
 
-    return combined_text
+    tags = ", ".join(detected_objs) if detected_objs else "None"
+    gallery = [(img, cap) for img, cap in zip(list(last_images), list(last_captions))]
 
-def build_history_table():
-    headers = ["Past Outputs (Click text to copy)"]
-    data = [[text] for text in reversed(last_texts)]
-    return headers, data
+    result_text = f"Detected objects: {tags}\nCaption: {caption}"
+    return result_text, gallery
 
 with gr.Blocks() as iface:
     gr.Markdown("# Image Captioning with Object Detection")
 
-    with gr.Column(scale=2):
-        image_input = gr.Image(type="pil", label="Upload Image")
-        generate_btn = gr.Button("Generate Caption")
-        caption_output = gr.Textbox(
-            label="Caption and Detected Objects",
-            lines=6,
-            interactive=True
-        )
-        copy_btn = gr.Button("Copy Output")
-
-    history_table = gr.Dataframe(
-        headers=["Session History"],
-        datatype=["str"],
-        interactive=True,
-        row_count=(0, MEMORY_SIZE),
-        col_count=1,
-        wrap=True
-    )
+    image_input = gr.Image(type="pil", label="Upload Image")
+
+    caption_output = gr.Textbox(label="Caption and Detected Objects", lines=3, interactive=False)
+
+    gallery = gr.Gallery(label="Last 15 Images and Captions", scale=3)
+
+    generate_btn = gr.Button("Generate Caption")
 
     def on_generate(image):
         if image is None:
-            return "Please upload an image.", (["Session History"], [])
-        combined = generate_caption(image)
-        headers, data = build_history_table()
-        return combined, (headers, data)
-
-    def copy_output(text):
-        return gr.Textbox.update(value=text, interactive=True)
+            return "Please upload an image.", []
+        return generate_caption(image)
 
     generate_btn.click(
         fn=on_generate,
         inputs=image_input,
-        outputs=[caption_output, history_table]
+        outputs=[caption_output, gallery]
     )
 
-    copy_btn.click(fn=copy_output, inputs=caption_output, outputs=caption_output)
-
 if __name__ == "__main__":
     iface.launch()