Update app.py

app.py

@@ -8,35 +8,59 @@ import gradio as gr
 from huggingface_hub import HfApi, InferenceClient
 from pydantic import BaseModel, Field
 
+# -------------------------------
+# Environment variables / Constants
+# -------------------------------
 HF_DATASET_REPO = os.environ.get("HF_DATASET_REPO", "OppaAI/Robot_MCP")
 HF_VLM_MODEL = os.environ.get("HF_VLM_MODEL", "Qwen/Qwen2.5-VL-7B-Instruct")
 
-# ---------------------------------------------------
-#  Payload Schema (Remains the same as it already expects image_b64)
-# ---------------------------------------------------
+# -------------------------------
+# Pydantic schema for the tool payload
+# -------------------------------
 class RobotWatchPayload(BaseModel):
+    """
+    Defines the expected input structure for the robot VLM analysis tool.
+    
+    Attributes:
+        hf_token (str): Your Hugging Face API token.
+        robot_id (str): Identifier for the robot (default "unknown").
+        image_b64 (str): Base64 encoded image string to analyze.
+    """
     hf_token: str = Field(description="Your Hugging Face API token.")
     robot_id: str = Field(description="Robot identifier.", default="unknown")
     image_b64: str = Field(description="Base64 encoded image data.")
 
 
-# ---------------------------------------------------
-# Upload Helper (Remains the same)
-# ---------------------------------------------------
+# -------------------------------
+# Helper function: Upload image to Hugging Face dataset
+# -------------------------------
 def upload_image(image_b64: str, hf_token: str):
+    """
+    Decodes a base64 image string, saves it locally, and uploads to Hugging Face dataset.
+
+    Args:
+        image_b64 (str): Base64 encoded image data.
+        hf_token (str): Hugging Face API token.
+
+    Returns:
+        tuple: (local_path, hf_url, filename, size_bytes)
+    """
     try:
         image_bytes = base64.b64decode(image_b64)
         os.makedirs("/tmp", exist_ok=True)
 
+        # Generate unique timestamped filename
         timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
         local_path = f"/tmp/robot_img_{timestamp}.jpg"
 
+        # Save locally
         with open(local_path, "wb") as f:
             f.write(image_bytes)
 
         filename = f"robot_{timestamp}.jpg"
-        api = HfApi()
 
+        # Upload to Hugging Face dataset
+        api = HfApi()
         api.upload_file(
             path_or_fileobj=local_path,
             path_in_repo=f"tmp/{filename}",
@@ -53,10 +77,20 @@ def upload_image(image_b64: str, hf_token: str):
         return None, None, None, 0
 
 
-# ---------------------------------------------------
-# JSON Cleaning Helper (Remains the same)
-# ---------------------------------------------------
+# -------------------------------
+# Helper function: Parse JSON safely
+# -------------------------------
 def safe_parse_json_from_text(text: str):
+    """
+    Attempts to parse JSON from text returned by the VLM model.
+    Strips any leading/trailing characters and handles malformed responses.
+
+    Args:
+        text (str): Raw text output from the model.
+
+    Returns:
+        dict or None: Parsed JSON dictionary, or None if parsing fails.
+    """
     if not text:
         return None
     try:
@@ -76,19 +110,29 @@ def safe_parse_json_from_text(text: str):
         return None
 
 
-# ---------------------------------------------------
-# Core VLM Analysis Logic (Remains the same)
-# ---------------------------------------------------
+# -------------------------------
+# Core VLM analysis function
+# -------------------------------
 def run_vlm_analysis(payload: RobotWatchPayload):
-    # ... (function body remains identical to previous version) ...
+    """
+    Main logic for analyzing an image using Hugging Face VLM model.
+
+    Args:
+        payload (RobotWatchPayload): Validated payload containing token, robot_id, and image.
+
+    Returns:
+        dict: Analysis result including description, objects, and raw VLM output.
+    """
     hf_token = payload.hf_token
     image_b64 = payload.image_b64
     robot_id = payload.robot_id
 
+    # Upload the image to Hugging Face dataset
     _, hf_url, _, size_bytes = upload_image(image_b64, hf_token)
     if not hf_url:
         return {"error": "Image upload failed"}
 
+    # System prompt instructs VLM to return strict JSON
     system_prompt = """
 Respond in STRICT JSON ONLY:
 {
@@ -131,45 +175,57 @@ Respond in STRICT JSON ONLY:
     }
 
 
-# ---------------------------------------------------
-# Gradio UI Function (NOW USES BASE64 STRING INPUT)
-# ---------------------------------------------------
+# -------------------------------
+# Gradio interface function
+# -------------------------------
 def robot_watch(
     hf_token_input: str,
     robot_id_input: str,
-    image_b64_input: str # Changed input type to a string (base64)
+    image_b64_input: str
 ):
     """
-    Handles input from individual Gradio components (including base64 string), 
-    converts to Pydantic model, and calls the core logic.
+    Gradio wrapper for run_vlm_analysis.
+    Converts individual fields into Pydantic model and calls core logic.
+
+    Args:
+        hf_token_input (str): Hugging Face API token input from UI.
+        robot_id_input (str): Robot ID input from UI.
+        image_b64_input (str): Base64 image input from UI.
+
+    Returns:
+        dict: Result from run_vlm_analysis.
     """
     if not image_b64_input:
         return {"error": "Base64 image string is empty."}
     
-    # Create the Pydantic model instance manually
+    # Create the payload instance
     payload_instance = RobotWatchPayload(
         hf_token=hf_token_input,
         robot_id=robot_id_input,
         image_b64=image_b64_input
     )
     
-    # Call the core logic
+    # Run core analysis
     result = run_vlm_analysis(payload_instance)
     return result
 
 
+# -------------------------------
+# Gradio App
+# -------------------------------
 app = gr.Interface(
-    fn=robot_watch, # Use the new multi-input function for the UI
+    fn=robot_watch,
     inputs=[
         gr.Textbox(label="Hugging Face Token", lines=1),
         gr.Textbox(label="Robot ID", lines=1, value="unknown"),
-        gr.Textbox(label="Image Base64 String", lines=5) # Changed input component to Textbox
+        gr.Textbox(label="Image Base64 String", lines=5)
     ],
     outputs=gr.Json(label="Tool Output"),
     title="Robot MCP Server (Base64 Inputs)",
-    description="Interface for the robot VLM analysis using individual fields, including base64 image string.",
+    description="Interface for robot VLM analysis using individual fields, including base64 image string.",
     api_name="predict"
 )
 
 if __name__ == "__main__":
+    # Launch Gradio app with MCP server enabled
     app.launch(mcp_server=True)