Update app.py

app.py

@@ -1,11 +1,12 @@
 import gradio as gr
 import torch
-from PIL import Image
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
-# ---- CPU-only config ----
+# -----------------------------
+# Model configuration
+# -----------------------------
 MID = "apple/FastVLM-0.5B"
-IMAGE_TOKEN_INDEX = -200  # special image token id used by FastVLM
+IMAGE_TOKEN_INDEX = -200
 
 tok = None
 model = None
@@ -13,50 +14,43 @@ model = None
 def load_model():
     global tok, model
     if tok is None or model is None:
-        print("Loading model (CPU)…")
+        print("Loading model on CPU...")
         tok = AutoTokenizer.from_pretrained(MID, trust_remote_code=True)
-        # Force CPU + float32 (fp16 is unsafe on CPU)
         model = AutoModelForCausalLM.from_pretrained(
             MID,
             torch_dtype=torch.float32,
             device_map="cpu",
             trust_remote_code=True,
         )
-        print("Model loaded successfully on CPU!")
+        print("Model loaded successfully!")
     return tok, model
 
-def caption_image(image, custom_prompt=None):
-    """
-    Generate a caption for the input image (CPU-only).
-    """
+
+def run_fastvlm(image, prompt):
     if image is None:
         return "Please upload an image first."
 
     try:
         tok, model = load_model()
 
-        # Convert image to RGB if needed
         if image.mode != "RGB":
             image = image.convert("RGB")
 
-        prompt = custom_prompt if custom_prompt else "Describe this image in detail."
-
-        # Single-turn chat with an <image> placeholder
         messages = [{"role": "user", "content": f"<image>\n{prompt}"}]
-        rendered = tok.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)
+        rendered = tok.apply_chat_template(
+            messages,
+            add_generation_prompt=True,
+            tokenize=False
+        )
 
-        # Split around the literal "<image>"
         pre, post = rendered.split("<image>", 1)
 
-        # Tokenize text around the image token
         pre_ids = tok(pre, return_tensors="pt", add_special_tokens=False).input_ids
         post_ids = tok(post, return_tensors="pt", add_special_tokens=False).input_ids
 
-        # Derive device/dtype from the loaded model (CPU here, but future-proof)
         model_device = next(model.parameters()).device
         model_dtype = next(model.parameters()).dtype
 
-        # Insert IMAGE token id at placeholder position
         img_tok = torch.tensor([[IMAGE_TOKEN_INDEX]], dtype=pre_ids.dtype, device=model_device)
         input_ids = torch.cat(
             [pre_ids.to(model_device), img_tok, post_ids.to(model_device)],
@@ -64,23 +58,22 @@ def caption_image(image, custom_prompt=None):
         )
         attention_mask = torch.ones_like(input_ids, device=model_device)
 
-        # Preprocess image using model's vision tower
-        px = model.get_vision_tower().image_processor(
-            images=image, return_tensors="pt"
+        pixel_values = model.get_vision_tower().image_processor(
+            images=image,
+            return_tensors="pt"
         )["pixel_values"].to(device=model_device, dtype=model_dtype)
 
-        # Generate caption (deterministic)
         with torch.no_grad():
             out = model.generate(
                 inputs=input_ids,
                 attention_mask=attention_mask,
-                images=px,
-                max_new_tokens=128,
-                do_sample=False,  # temperature is ignored when sampling is off
+                images=pixel_values,
+                max_new_tokens=220,
+                do_sample=False
             )
 
-        # Decode and slice to the assistant part if present
         generated_text = tok.decode(out[0], skip_special_tokens=True)
+
         if "Assistant:" in generated_text:
             response = generated_text.split("Assistant:", 1)[-1].strip()
         elif "assistant" in generated_text:
@@ -91,53 +84,160 @@ def caption_image(image, custom_prompt=None):
         return response
 
     except Exception as e:
-        return f"Error generating caption: {str(e)}"
-
-# ---- Gradio UI (CPU) ----
-with gr.Blocks(title="FastVLM Image Captioning (CPU)") as demo:
-    gr.Markdown(
-        """
-        # 🖼️ FastVLM Image Captioning (CPU)
-        Upload an image to generate a detailed caption using Apple's FastVLM-0.5B.
-        This build runs on **CPU only**. Expect slower generation than GPU.
-        """
-    )
+        return f"Error: {str(e)}"
+
+
+def build_prompt(mode, user_context):
+    context_part = f"\nExtra user context: {user_context}" if user_context.strip() else ""
+
+    prompts = {
+        "Scene Description":
+            f"""
+You are an AI assistant helping a visually impaired person.
+Describe the image in simple, human-friendly language.
+
+Return output in this format:
+1. Quick Summary
+2. Main Objects Seen
+3. Relative Position of Important Objects
+4. Helpful Note
+
+Keep the language simple and practical.{context_part}
+""",
+
+        "Hazard Detection":
+            f"""
+You are an AI safety assistant helping a visually impaired person.
+Analyze the image for possible hazards.
+
+Return output in this format:
+1. Quick Summary
+2. Possible Hazards
+3. Risk Level (Low/Medium/High)
+4. Safety Advice
+
+Be practical and avoid exaggeration.{context_part}
+""",
+
+        "Important Object Summary":
+            f"""
+You are an AI visual assistant.
+Identify the most important objects in the image that a visually impaired person should know about.
+
+Return output in this format:
+1. Key Objects
+2. What Looks Most Important
+3. Why These Objects Matter
+4. Short Spoken Summary
+
+Keep it easy to understand.{context_part}
+""",
+
+        "Safe Action Suggestion":
+            f"""
+You are an AI guidance assistant for a visually impaired person.
+Based on the image, suggest the next safest action.
+
+Return output in this format:
+1. What the Scene Looks Like
+2. What Needs Attention
+3. Recommended Action
+4. One-Line Safety Tip
+
+Do not assume too much. Give cautious guidance.{context_part}
+"""
+    }
+
+    return prompts.get(mode, prompts["Scene Description"])
+
+
+def analyze_image(image, mode, user_context):
+    if image is None:
+        return "Please upload an image."
+
+    prompt = build_prompt(mode, user_context)
+    return run_fastvlm(image, prompt)
+
+
+def exhibition_pitch(mode):
+    pitches = {
+        "Scene Description":
+            "This mode explains the surrounding environment in simple words so a visually impaired person can understand the scene.",
+        "Hazard Detection":
+            "This mode checks whether the image contains obstacles or risky elements such as vehicles, stairs, clutter, or unsafe walking areas.",
+        "Important Object Summary":
+            "This mode highlights the most useful objects in the scene so the user can focus on what matters most.",
+        "Safe Action Suggestion":
+            "This mode provides the next practical action the user should consider, based on the visual situation."
+    }
+    return pitches.get(mode, "")
+
+
+with gr.Blocks(title="VisionMate AI - Smart Visual Assistant") as demo:
+    gr.Markdown("""
+# 👁️ VisionMate AI
+## Smart Visual Assistant for Visually Impaired People
+
+Upload an image and let the AI explain the scene, identify hazards, summarize important objects, or suggest the safest next action.
+
+### Exhibition Theme
+**AI for Social Good**
+""")
 
     with gr.Row():
-        with gr.Column():
-            image_input = gr.Image(type="pil", label="Upload Image", elem_id="image-upload")
-            custom_prompt = gr.Textbox(
-                label="Custom Prompt (Optional)",
-                placeholder="Leave empty for default: 'Describe this image in detail.'",
+        with gr.Column(scale=1):
+            image_input = gr.Image(type="pil", label="Upload Scene Image")
+
+            mode = gr.Radio(
+                choices=[
+                    "Scene Description",
+                    "Hazard Detection",
+                    "Important Object Summary",
+                    "Safe Action Suggestion"
+                ],
+                value="Scene Description",
+                label="Select Assistance Mode"
+            )
+
+            user_context = gr.Textbox(
+                label="Optional Context",
+                placeholder="Example: Person is walking alone on a road / indoor corridor / market area",
                 lines=2
             )
+
             with gr.Row():
-                clear_btn = gr.ClearButton([image_input, custom_prompt])
-                generate_btn = gr.Button("Generate Caption", variant="primary")
+                analyze_btn = gr.Button("Analyze Scene", variant="primary")
+                clear_btn = gr.ClearButton([image_input, user_context])
+
+        with gr.Column(scale=1):
+            mode_explanation = gr.Textbox(
+                label="Mode Purpose",
+                value=exhibition_pitch("Scene Description"),
+                interactive=False,
+                lines=4
+            )
 
-        with gr.Column():
             output = gr.Textbox(
-                label="Generated Caption",
-                lines=8,
-                max_lines=15,
+                label="AI Assistance Output",
+                lines=16,
+                max_lines=25,
                 show_copy_button=True
             )
 
-    generate_btn.click(fn=caption_image, inputs=[image_input, custom_prompt], outputs=output)
-
-    # Also generate on image upload if no custom prompt
-    def _auto_caption(img, prompt):
-        return caption_image(img, prompt) if (img is not None and not prompt) else None
+    mode.change(fn=exhibition_pitch, inputs=mode, outputs=mode_explanation)
+    analyze_btn.click(fn=analyze_image, inputs=[image_input, mode, user_context], outputs=output)
 
-    image_input.change(fn=_auto_caption, inputs=[image_input, custom_prompt], outputs=output)
+    gr.Markdown("""
+---
+### Suggested Demo Images for Exhibition
+- A road with vehicles and pedestrians
+- A classroom or hallway
+- A kitchen or home environment
+- A supermarket shelf or crowded place
 
-    gr.Markdown(
-        """
-        ---
-        **Model:** [apple/FastVLM-0.5B](https://huggingface.co/apple/FastVLM-0.5B)  
-        **Note:** CPU-only run. For speed, switch to a CUDA GPU build or a GPU Space.
-        """
-    )
+### Expected Impact
+This project shows how computer vision and multimodal AI can improve accessibility and independence for visually impaired users.
+""")
 
 if __name__ == "__main__":
     demo.launch(