| import re |
| from typing import Any, Dict, Iterable, List, Optional, Tuple |
|
|
| import numpy as np |
| import cv2 |
| import copy |
|
|
|
|
| def process_segmentation( |
| segmentation: np.ndarray, |
| segmentation_id_map: Optional[Dict[int, Any]], |
| color_map: Dict[int, List[int]], |
| current_segment: Any, |
| current_subgoal_segment: Optional[str], |
| previous_subgoal_segment: Optional[str], |
| current_task_name: str, |
| existing_points: Optional[List[List[int]]] = None, |
| existing_subgoal_filled: Optional[str] = None, |
| ) -> Dict[str, Any]: |
| """ |
| Shared helper to compute segmentation filtering and grounded subgoal text. |
| |
| Returns a dict with: |
| - segmentation_result: segmentation mask filtered to visible ids |
| - segmentation_result_2d: squeezed version of segmentation_result |
| - segmentation_points: cached center points for current targets |
| - current_subgoal_segment_filled: subgoal string with centers filled in |
| - no_object_flag: whether the target ids are missing in the mask |
| - updated_previous_subgoal_segment: equals current_subgoal_segment for caller caching |
| - vis_obj_id_list: ids kept in segmentation_result |
| """ |
| segmentation_2d = segmentation.squeeze() if segmentation.ndim > 2 else segmentation |
|
|
| if isinstance(current_segment, (list, tuple)): |
| active_segments = list(current_segment) |
| elif current_segment is None: |
| active_segments = [] |
| else: |
| active_segments = [current_segment] |
|
|
| segment_ids_by_index = {idx: [] for idx in range(len(active_segments))} |
| vis_obj_id_list: List[int] = [] |
| if isinstance(segmentation_id_map, dict): |
| for obj_id, obj in sorted(segmentation_id_map.items()): |
| if active_segments: |
| for idx, target in enumerate(active_segments): |
| if obj is target: |
| vis_obj_id_list.append(obj_id) |
| segment_ids_by_index[idx].append(obj_id) |
| break |
| if getattr(obj, "name", None) == "table-workspace": |
| color_map[obj_id] = [0, 0, 0] |
|
|
| segmentation_result = np.where( |
| np.isin(segmentation_2d, vis_obj_id_list), segmentation_2d, 0 |
| ) |
| segmentation_result_2d = segmentation_result.squeeze() |
|
|
| segmentation_points = existing_points or [] |
| current_subgoal_segment_filled = existing_subgoal_filled |
| no_object_flag = False |
|
|
| if current_subgoal_segment != previous_subgoal_segment: |
|
|
| def compute_center_from_ids(segmentation_mask: np.ndarray, ids: Iterable[int]): |
| nonlocal no_object_flag |
| ids = list(ids) |
| if not ids: |
| return None |
| mask = np.isin(segmentation_mask, ids) |
| if not np.any(mask): |
| no_object_flag = True |
| return None |
| coords = np.argwhere(mask) |
| if coords.size == 0: |
| return None |
| center_y = int(coords[:, 0].mean()) |
| center_x = int(coords[:, 1].mean()) |
| return [center_y, center_x] |
|
|
| segment_centers: List[Optional[List[int]]] = [] |
| if active_segments: |
| for idx in range(len(active_segments)): |
| segment_centers.append( |
| compute_center_from_ids( |
| segmentation_2d, segment_ids_by_index.get(idx, []) |
| ) |
| ) |
| else: |
| segment_centers.append( |
| compute_center_from_ids(segmentation_2d, vis_obj_id_list) |
| ) |
|
|
| segmentation_points = [center for center in segment_centers if center is not None] |
|
|
| if current_subgoal_segment: |
| normalized_centers: List[Optional[str]] = [] |
| for center in segment_centers: |
| if center is None: |
| normalized_centers.append(None) |
| continue |
| center_y, center_x = center |
| normalized_centers.append(f"<{center_y}, {center_x}>") |
|
|
| placeholder_pattern = re.compile(r"<[^>]*>") |
| placeholders = list(placeholder_pattern.finditer(current_subgoal_segment)) |
| placeholder_count = len(placeholders) |
| if placeholder_count > 0 and normalized_centers: |
| replacements = normalized_centers.copy() |
| if len(replacements) == 1 and placeholder_count > 1: |
| replacements = replacements * placeholder_count |
| elif len(replacements) < placeholder_count: |
| replacements.extend([None] * (placeholder_count - len(replacements))) |
|
|
| missing_placeholder = False |
| new_text_parts: List[str] = [] |
| last_idx = 0 |
| for idx, match in enumerate(placeholders): |
| new_text_parts.append( |
| current_subgoal_segment[last_idx : match.start()] |
| ) |
| replacement_text = replacements[idx] |
| if replacement_text is None: |
| missing_placeholder = True |
| else: |
| new_text_parts.append(replacement_text) |
| last_idx = match.end() |
| new_text_parts.append(current_subgoal_segment[last_idx:]) |
| current_subgoal_segment_filled = ( |
| current_task_name if missing_placeholder else "".join(new_text_parts) |
| ) |
| else: |
| current_subgoal_segment_filled = current_subgoal_segment |
| else: |
| current_subgoal_segment_filled = current_subgoal_segment |
|
|
| return { |
| "segmentation_result": segmentation_result, |
| "segmentation_result_2d": segmentation_result_2d, |
| "segmentation_points": segmentation_points, |
| "current_subgoal_segment_filled": current_subgoal_segment_filled, |
| "no_object_flag": no_object_flag, |
| "updated_previous_subgoal_segment": current_subgoal_segment, |
| "vis_obj_id_list": vis_obj_id_list, |
| } |
|
|
|
|
| def create_segmentation_visuals( |
| segmentation: np.ndarray, |
| segmentation_result: np.ndarray, |
| base_frame: np.ndarray, |
| color_map: Dict[int, List[int]], |
| segmentation_points: List[List[int]], |
| ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: |
| """ |
| Build colored segmentation visualizations and target overlay for video export. |
| |
| Returns (segmentation_vis, segmentation_result_vis, target_for_video). |
| """ |
| segmentation_for_video = copy.deepcopy(segmentation) |
| segmentation_result_for_video = copy.deepcopy(segmentation_result) |
|
|
| segmentation_vis = np.zeros( |
| (*segmentation_for_video.shape[:2], 3), dtype=np.uint8 |
| ) |
| segmentation_result_vis = np.zeros( |
| (*segmentation_result_for_video.shape[:2], 3), dtype=np.uint8 |
| ) |
|
|
| seg_2d = ( |
| segmentation_for_video.squeeze() |
| if segmentation_for_video.ndim > 2 |
| else segmentation_for_video |
| ) |
| seg_result_2d = ( |
| segmentation_result_for_video.squeeze() |
| if segmentation_result_for_video.ndim > 2 |
| else segmentation_result_for_video |
| ) |
|
|
| for seg_id in np.unique(seg_2d): |
| if seg_id > 0: |
| color = color_map.get(seg_id, [255, 255, 255]) |
| mask = seg_2d == seg_id |
| segmentation_vis[mask] = color |
|
|
| for seg_id in np.unique(seg_result_2d): |
| if seg_id > 0: |
| color = color_map.get(seg_id, [255, 255, 255]) |
| mask = seg_result_2d == seg_id |
| segmentation_result_vis[mask] = color |
|
|
| target_for_video = copy.deepcopy(base_frame) |
|
|
| if segmentation_vis.shape[:2] != base_frame.shape[:2]: |
| segmentation_vis = cv2.resize( |
| segmentation_vis, |
| (base_frame.shape[1], base_frame.shape[0]), |
| interpolation=cv2.INTER_NEAREST, |
| ) |
|
|
| if segmentation_result_vis.shape[:2] != base_frame.shape[:2]: |
| segmentation_result_vis = cv2.resize( |
| segmentation_result_vis, |
| (base_frame.shape[1], base_frame.shape[0]), |
| interpolation=cv2.INTER_NEAREST, |
| ) |
|
|
| if segmentation_points: |
| diameter = 5 |
| for center_y, center_x in segmentation_points: |
| cv2.circle(target_for_video, (center_x, center_y), diameter, (255, 0, 0), -1) |
|
|
| return segmentation_vis, segmentation_result_vis, target_for_video |
|
|
