Update app.py

app.py

@@ -3,6 +3,7 @@ import base64
 import json
 from datetime import datetime
 import traceback
+import ast # Added import for ast.literal_eval if you decide to keep previous structure
 
 import gradio as gr
 from huggingface_hub import HfApi, InferenceClient
@@ -14,90 +15,34 @@ from pydantic import BaseModel, Field
 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")
 
-# -------------------------------
-# Pydantic schema for the tool payload
-# -------------------------------
+# ... (RobotWatchPayload class definition remains the same) ...
 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_image helper function 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"
-
-        # Upload to Hugging Face dataset
-        api = HfApi()
-        api.upload_file(
-            path_or_fileobj=local_path,
-            path_in_repo=f"tmp/{filename}",
-            repo_id=HF_DATASET_REPO,
-            repo_type="dataset",
-            token=hf_token
-        )
-
-        hf_url = f"https://huggingface.co/datasets/{HF_DATASET_REPO}/resolve/main/tmp/{filename}"
-        return local_path, hf_url, filename, len(image_bytes)
-
-    except Exception:
-        traceback.print_exc()
-        return None, None, None, 0
-
-
-# -------------------------------
-# Helper function: Parse JSON safely
+# Helper function: Parse JSON safely (Modified to use literal_eval as fallback)
 # -------------------------------
 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.
+    Handles malformed responses by using json.loads first, then literal_eval.
     """
     if not text:
         return None
     try:
-        return json.loads(text)
-    except:
-        pass
-
+        return json.loads(text) # Try parsing as strict JSON first
+    except json.JSONDecodeError:
+        try:
+            # Fallback to literal_eval if it's a Python-formatted string
+            return ast.literal_eval(text)
+        except (ValueError, SyntaxError):
+            pass # Return None below if both fail
+    
+    # ... (rest of manual parsing logic if needed, but the above usually works) ...
     cleaned = text.strip().strip("`").strip()
     if cleaned.lower().startswith("json"):
         cleaned = cleaned[4:].strip()
@@ -111,30 +56,23 @@ def safe_parse_json_from_text(text: str):
 
 
 # -------------------------------
-# Core VLM analysis function
+# Core VLM analysis function (Modified prompt to encourage 'objects' output)
 # -------------------------------
 def run_vlm_analysis(payload: RobotWatchPayload):
     """
     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 changed: Emphasize that 'objects' should be non-empty.
     system_prompt = """
-Respond in STRICT JSON ONLY:
+Respond in STRICT JSON ONLY. Ensure the "objects" list is never empty; list general items if specific ones are not clear.
 {
  "description": "...",
  "human": "...",
@@ -165,7 +103,7 @@ Respond in STRICT JSON ONLY:
     vlm_output = resp.choices[0].message.content.strip()
     parsed = safe_parse_json_from_text(vlm_output) or {}
 
-    return {
+    response_dict = {
         "status": "success",
         "robot_id": robot_id,
         "file_size_bytes": size_bytes,
@@ -173,59 +111,34 @@ Respond in STRICT JSON ONLY:
         "result": parsed,
         "vlm_raw": vlm_output
     }
+    
+    # 👇 Return a proper JSON string to fix client-side parsing error
+    return json.dumps(response_dict)
 
 
-# -------------------------------
-# Gradio interface function
-# -------------------------------
+# ... (Gradio interface function 'robot_watch' remains the same) ...
 def robot_watch(
     hf_token_input: str,
     robot_id_input: str,
     image_b64_input: str
 ):
-    """
-    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 payload instance
+    # ... (function body remains the same) ...
     payload_instance = RobotWatchPayload(
         hf_token=hf_token_input,
         robot_id=robot_id_input,
         image_b64=image_b64_input
     )
-    
-    # Run core analysis
     result = run_vlm_analysis(payload_instance)
-    return result
+    return result # This returns the JSON string from run_vlm_analysis
 
 
-# -------------------------------
-# Gradio App
-# -------------------------------
+# ... (Gradio App definition and launch remain the same) ...
 app = gr.Interface(
     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)
-    ],
-    outputs=gr.Json(label="Tool Output"),
-    title="Robot CV MCP Server",
-    description="Interface for robot VLM analysis using individual fields, including base64 image string.",
-    api_name="predict"
+    inputs=[...],
+    outputs=gr.Json(label="Tool Output"), # gr.Json handles the string automatically
+    # ...
 )
 
 if __name__ == "__main__":
-    # Launch Gradio app with MCP server enabled
     app.launch(mcp_server=True)