Update app.py

app.py

@@ -1,24 +1,50 @@
-from transformers import BlipProcessor, BlipForConditionalGeneration
-import torch
+from transformers import BlipProcessor, BlipForConditionalGeneration, pipeline
 import gradio as gr
 from PIL import Image
 
-# Load model and processor
+# Load the main image captioning model
 processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
 model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
 
-def generate_caption(image):
+# Load a small language model pipeline for polishing (can change model to a better LLM as needed)
+text_generator = pipeline("text-generation", model="gpt2")
+
+def preprocess_image(image):
+    # Example: convert to RGB if needed, resize if you want consistent input size
+    if image.mode != "RGB":
+        image = image.convert("RGB")
+    return image
+
+def postprocess_caption(raw_caption):
+    # Use language model to polish and expand the caption slightly
+    # Limit max new tokens to keep the output concise
+    polished = text_generator(f"Describe this image in more detail: {raw_caption}", max_length=50, num_return_sequences=1)
+    # Extract generated text, remove prompt part
+    polished_caption = polished[0]['generated_text'].replace(f"Describe this image in more detail: ", "").strip()
+    return polished_caption
+
+def generate_caption(image, max_length=30, num_beams=5):
+    image = preprocess_image(image)
     inputs = processor(image, return_tensors="pt")
-    out = model.generate(**inputs)
-    caption = processor.decode(out[0], skip_special_tokens=True)
-    return caption
+    # Generate caption with adjustable parameters for length and quality
+    out = model.generate(**inputs, max_length=max_length, num_beams=num_beams, early_stopping=True)
+    raw_caption = processor.decode(out[0], skip_special_tokens=True)
+    polished_caption = postprocess_caption(raw_caption)
+    return polished_caption
 
-# Create Gradio interface
+# Gradio interface with sliders for max_length and num_beams parameters and descriptions
 iface = gr.Interface(
     fn=generate_caption,
-    inputs=gr.Image(type="pil"),
+    inputs=[
+        gr.Image(type="pil", label="Upload Image"),
+        gr.Slider(10, 50, value=30, step=5, label="Caption Max Length", 
+                  info="Controls the length of the caption (higher means longer captions)"),
+        gr.Slider(1, 10, value=5, step=1, label="Beam Search Width", 
+                  info="Controls how many caption options the model considers before picking the best")
+    ],
     outputs="text",
-    title="Image Captioning App"
+    title="Enhanced Image Captioning App",
+    description="Upload an image and get a polished, detailed description. Adjust sliders to control caption length and quality."
 )
 
 if __name__ == "__main__":