Update app.py

app.py

@@ -186,7 +186,7 @@ def get_uitars_prompt(task, image):
 
 # --- Holo2 Prompt ---
 def get_holo2_prompt(task, image):
-    # Holo2 typically uses standard chat formatting
+    # Holo2 often expects a simple user prompt with the image
     return [
         {
             "role": "user",
@@ -211,7 +211,7 @@ def get_image_proc_params(processor) -> Dict[str, int]:
 # -----------------------------------------------------------------------------
 
 def parse_uitars_response(text: str) -> List[Dict]:
-    """Parse UI-TARS specific output formats"""
+    """Parse UI-TARS output"""
     actions = []
     text = text.strip()
     
@@ -221,13 +221,10 @@ def parse_uitars_response(text: str) -> List[Dict]:
     m = re.findall(r"point=\[\s*(\d+)\s*,\s*(\d+)\s*\]", text, re.IGNORECASE)
     for p in m: actions.append({"type": "click", "x": int(p[0]), "y": int(p[1]), "text": ""})
 
-    m = re.search(r"start_box=['\"]?\(\s*(\d+)\s*,\s*(\d+)\s*\)['\"]?", text, re.IGNORECASE)
-    if m: actions.append({"type": "click", "x": int(m[0]), "y": int(m[1]), "text": ""})
-
     return actions
 
 def parse_fara_response(response: str) -> List[Dict]:
-    """Parse Fara <tool_call> JSON format"""
+    """Parse Fara output"""
     actions = []
     matches = re.findall(r"<tool_call>(.*?)</tool_call>", response, re.DOTALL)
     for match in matches:
@@ -245,70 +242,56 @@ def parse_fara_response(response: str) -> List[Dict]:
     return actions
 
 def parse_holo2_response(generated_ids, processor, input_len) -> Tuple[str, str, List[Dict]]:
-    """Parse Holo2 reasoning tokens and JSON content"""
+    """Parse Holo2 reasoning and actions"""
     all_ids = generated_ids[0].tolist()
     
-    # Token IDs for <|thought_start|> and <|thought_end|> (Qwen/Holo specific)
+    # Qwen/Holo specific reasoning tokens
     THOUGHT_START = 151667
     THOUGHT_END = 151668
     
     thinking_content = ""
     content = ""
     
-    try:
-        if THOUGHT_START in all_ids:
-            start_idx = all_ids.index(THOUGHT_START)
-            try:
-                end_idx = all_ids.index(THOUGHT_END)
-            except ValueError:
-                end_idx = len(all_ids)
-                
-            thinking_ids = all_ids[start_idx+1:end_idx]
-            thinking_content = processor.decode(thinking_ids, skip_special_tokens=True).strip()
-            
-            # Content is everything after thought_end
-            content_ids = all_ids[end_idx+1:]
-            content = processor.decode(content_ids, skip_special_tokens=True).strip()
-        else:
-            # Fallback if no reasoning tokens found (just raw output)
-            # Slice off input tokens first
-            output_ids = all_ids[input_len:]
-            content = processor.decode(output_ids, skip_special_tokens=True).strip()
-    except Exception as e:
-        print(f"Holo Parsing Error: {e}")
-        content = processor.decode(all_ids[input_len:], skip_special_tokens=True).strip()
-
-    # Parse JSON Content
+    # 1. Extract Thinking
+    if THOUGHT_START in all_ids:
+        start_idx = all_ids.index(THOUGHT_START)
+        try:
+            end_idx = all_ids.index(THOUGHT_END)
+        except ValueError:
+            end_idx = len(all_ids)
+        thinking_ids = all_ids[start_idx+1:end_idx]
+        thinking_content = processor.decode(thinking_ids, skip_special_tokens=True).strip()
+        # Content is after thought_end
+        output_ids = all_ids[end_idx+1:]
+        content = processor.decode(output_ids, skip_special_tokens=True).strip()
+    else:
+        output_ids = all_ids[input_len:]
+        content = processor.decode(output_ids, skip_special_tokens=True).strip()
+
+    # 2. Extract Coordinates (Robust parsing)
     actions = []
-    try:
-        # Holo2 outputs strictly valid JSON usually
-        # E.g. {"x": 500, "y": 300, "description": "search bar"}
-        # Or {"action": "click", "point": [100, 200]}
-        # Flattening to common format
-        if "{" in content and "}" in content:
-            # Find JSON block if surrounded by text
-            json_str = re.search(r"(\{.*\})", content, re.DOTALL).group(1)
-            data = json.loads(json_str)
-            
-            x, y = 0, 0
-            if "x" in data and "y" in data:
-                x, y = data["x"], data["y"]
-            elif "point" in data:
-                x, y = data["point"][0], data["point"][1]
-            elif "coordinate" in data:
-                x, y = data["coordinate"][0], data["coordinate"][1]
-            
-            if x or y:
-                # Holo2 output is 0-1000 scale
-                actions.append({
-                    "type": "click",
-                    "x": float(x),
-                    "y": float(y), 
-                    "text": data.get("description", "") or data.get("text", ""),
-                    "scale_base": 1000 # Flag to indicate this needs normalization from 1000
-                })
-    except Exception as e:
-        print(f"Holo JSON Parse Failed: {e}")
+    
+    # Pattern A: point=[x, y] (Common in Holo)
+    points = re.findall(r"point=\[\s*(\d+)\s*,\s*(\d+)\s*\]", content)
+    for p in points:
+        actions.append({"type": "click", "x": float(p[0]), "y": float(p[1]), "scale_base": 1000})
+
+    # Pattern B: JSON {"point": [x, y]}
+    json_candidates = re.findall(r"\{.*?\}", content, re.DOTALL)
+    for jc in json_candidates:
+        try:
+            data = json.loads(jc)
+            if "point" in data:
+                actions.append({"type": "click", "x": float(data["point"][0]), "y": float(data["point"][1]), "scale_base": 1000})
+            if "coordinate" in data:
+                actions.append({"type": "click", "x": float(data["coordinate"][0]), "y": float(data["coordinate"][1]), "scale_base": 1000})
+        except: pass
+
+    # Pattern C: Plain [x, y] at end of string
+    if not actions:
+        plain_coords = re.findall(r"\[\s*(\d+)\s*,\s*(\d+)\s*\]", content)
+        for p in plain_coords:
+            actions.append({"type": "click", "x": float(p[0]), "y": float(p[1]), "scale_base": 1000})
 
     return content, thinking_content, actions
 
@@ -325,38 +308,41 @@ def create_localized_image(original_image: Image.Image, actions: list[dict]) ->
         x = act['x']
         y = act['y']
         
-        # Holo2 Special Case (0-1000 scaling)
+        # Scaling Logic
+        pixel_x, pixel_y = 0, 0
+        
+        # Case 1: Holo2 0-1000 scale
         if act.get('scale_base') == 1000:
-            pixel_x = int((x / 1000) * width)
-            pixel_y = int((y / 1000) * height)
-        # Normalized (0-1)
+            pixel_x = int((x / 1000.0) * width)
+            pixel_y = int((y / 1000.0) * height)
+        # Case 2: Normalized 0-1
         elif x <= 1.0 and y <= 1.0 and x > 0:
             pixel_x = int(x * width)
             pixel_y = int(y * height)
-        # Absolute Pixels
+        # Case 3: Absolute Pixels
         else:
             pixel_x = int(x)
             pixel_y = int(y)
             
-        color = 'red' if 'click' in act['type'].lower() else 'blue'
+        color = 'red' 
         
-        # Draw Visuals
+        # Draw Markers (Thicker for visibility)
         r = 15
-        draw.ellipse([pixel_x - r, pixel_y - r, pixel_x + r, pixel_y + r], outline=color, width=4)
-        draw.ellipse([pixel_x - 3, pixel_y - 3, pixel_x + 3, pixel_y + 3], fill=color)
+        draw.ellipse([pixel_x - r, pixel_y - r, pixel_x + r, pixel_y + r], outline=color, width=5)
+        draw.ellipse([pixel_x - 4, pixel_y - 4, pixel_x + 4, pixel_y + 4], fill=color)
         
-        # Draw Cross
-        draw.line([pixel_x - 10, pixel_y, pixel_x + 10, pixel_y], fill=color, width=2)
-        draw.line([pixel_x, pixel_y - 10, pixel_x, pixel_y + 10], fill=color, width=2)
+        # Crosshair
+        draw.line([pixel_x - 20, pixel_y, pixel_x + 20, pixel_y], fill=color, width=3)
+        draw.line([pixel_x, pixel_y - 20, pixel_x, pixel_y + 20], fill=color, width=3)
         
-        # Label
-        label = f"{act['type']}"
-        if act['text']: label += f": {act['text']}"
+        # Text Label
+        label = f"{act.get('type','Action')}"
+        text_pos = (pixel_x + 20, pixel_y - 15)
         
-        text_pos = (pixel_x + 20, pixel_y - 10)
-        bbox = draw.textbbox(text_pos, label, font=font)
-        draw.rectangle((bbox[0]-4, bbox[1]-2, bbox[2]+4, bbox[3]+2), fill="black")
-        draw.text(text_pos, label, fill="white", font=font)
+        if font:
+            bbox = draw.textbbox(text_pos, label, font=font)
+            draw.rectangle((bbox[0]-4, bbox[1]-2, bbox[2]+4, bbox[3]+2), fill="black")
+            draw.text(text_pos, label, fill="white", font=font)
 
     return img_copy
 
@@ -366,18 +352,17 @@ def create_localized_image(original_image: Image.Image, actions: list[dict]) ->
 
 @spaces.GPU(duration=120)
 def process_screenshot(input_numpy_image: np.ndarray, task: str, model_choice: str):
-    if input_numpy_image is None: return "⚠️ Please upload an image.", None, None
+    if input_numpy_image is None: return "⚠️ Please upload an image.", None
 
     input_pil_image = array_to_image(input_numpy_image)
     orig_w, orig_h = input_pil_image.size
     actions = []
     raw_response = ""
-    reasoning_text = None
+    reasoning_text = ""
 
     # --- UI-TARS Logic ---
     if model_choice == "UI-TARS-1.5-7B":
-        if model_x is None: return "Error: UI-TARS model failed to load.", None, None
-        print("Running UI-TARS...")
+        if model_x is None: return "Error: UI-TARS model failed to load.", None
         
         ip_params = get_image_proc_params(processor_x)
         resized_h, resized_w = smart_resize(
@@ -400,7 +385,7 @@ def process_screenshot(input_numpy_image: np.ndarray, task: str, model_choice: s
         
         actions = parse_uitars_response(raw_response)
         
-        # Rescale
+        # Rescale UI-TARS coords
         scale_x = orig_w / resized_w
         scale_y = orig_h / resized_h
         for a in actions:
@@ -409,8 +394,7 @@ def process_screenshot(input_numpy_image: np.ndarray, task: str, model_choice: s
 
     # --- Holo2 Logic ---
     elif model_choice == "Holo2-8B":
-        if model_h is None: return "Error: Holo2 model failed to load.", None, None
-        print("Running Holo2...")
+        if model_h is None: return "Error: Holo2 model failed to load.", None
         
         messages = get_holo2_prompt(task, input_pil_image)
         text_prompt = processor_h.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
@@ -421,18 +405,12 @@ def process_screenshot(input_numpy_image: np.ndarray, task: str, model_choice: s
         with torch.no_grad():
             generated_ids = model_h.generate(**inputs, max_new_tokens=512)
         
-        # Parse Reasoning + Content
         input_len = len(inputs.input_ids[0])
-        content, thinking, parsed_actions = parse_holo2_response(generated_ids, processor_h, input_len)
-        
-        raw_response = content
-        reasoning_text = thinking
-        actions = parsed_actions
+        raw_response, reasoning_text, actions = parse_holo2_response(generated_ids, processor_h, input_len)
 
     # --- Fara Logic ---
     else: 
-        if model_v is None: return "Error: Fara model failed to load.", None, None
-        print("Running Fara...")
+        if model_v is None: return "Error: Fara model failed to load.", None
         messages = get_fara_prompt(task, input_pil_image)
         text_prompt = processor_v.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
         image_inputs, video_inputs = process_vision_info(messages)
@@ -446,20 +424,17 @@ def process_screenshot(input_numpy_image: np.ndarray, task: str, model_choice: s
         raw_response = processor_v.batch_decode(generated_ids, skip_special_tokens=True)[0]
         actions = parse_fara_response(raw_response)
 
-    print(f"Raw: {raw_response}")
-    if reasoning_text: print(f"Thinking: {reasoning_text}")
-
     # Visualize
     output_image = input_pil_image
     if actions:
         vis = create_localized_image(input_pil_image, actions)
         if vis: output_image = vis
             
-    final_text_output = f"▶️ OUTPUT:\n{raw_response}"
+    final_output = f"▶️ OUTPUT:\n{raw_response}"
     if reasoning_text:
-        final_text_output = f"🧠 THINKING PROCESS:\n{reasoning_text}\n\n" + final_text_output
+        final_output = f"🧠 THINKING:\n{reasoning_text}\n\n" + final_output
 
-    return final_text_output, output_image
+    return final_output, output_image
 
 # -----------------------------------------------------------------------------
 # 6. UI SETUP