pynode-flow 0.2.0__py3-none-any.whl

pynode/nodes/SliceNode/slice_collector_node.py

@@ -0,0 +1,370 @@
+"""
+Slice Collector Node - collects predictions from multiple image slices and merges them.
+Designed to work with SliceImageNode in a slice-detect-merge workflow.
+"""
+
+from dataclasses import dataclass
+import time
+from typing import Any, Dict, List, Optional
+from pynode.nodes.base_node import BaseNode, Info, MessageKeys
+
+_info = Info()
+_info.add_text("Collects detection predictions from multiple image slices sent as separate messages, then merges them into a unified result.")
+_info.add_header("Inputs")
+_info.add_bullets(("Input 0:", "Individual slice prediction messages with slice metadata (from split output mode)"))
+_info.add_header("Outputs")
+_info.add_bullets(
+    ("Output 0:", "Merged predictions after all slices collected and NMS applied"),
+)
+_info.add_header("Configuration")
+_info.add_bullets(
+    ("Timeout:", "Maximum time to wait for all slices before outputting partial results"),
+    ("NMS IoU Threshold:", "Threshold for Non-Maximum Suppression"),
+    ("Match Metric:", "IoU or IoS for duplicate detection matching"),
+    ("Class Agnostic NMS:", "Apply NMS across all classes vs per-class"),
+)
+_info.add_header("Usage")
+_info.add_text("Designed for workflows: SliceImageNode (split mode) → Inference → SliceCollectorNode")
+
+@dataclass(frozen=True)
+class NodeKeys:
+    BBOX = 'bbox'
+    SLICE_INDEX = 'slice_index'
+    SLICE_OFFSET = 'slice_offset'
+    SLICE_BBOX = 'slice_bbox'
+    SLICES = 'slices'
+    SLICE_COUNT = 'slice_count'
+    COUNT = 'count'
+    EXPECTED_COUNT = 'expected_count'
+    OFFSET = 'offset'
+    IS_FULL_IMAGE = 'is_full_image'
+    ORIGINAL_WIDTH = 'original_width'
+    ORIGINAL_HEIGHT = 'original_height'
+    
+    NMS_THRESHOLD = 'nms_threshold'
+    MATCH_METRIC = 'match_metric'
+    CLASS_AGNOSTIC = 'class_agnostic'
+    TIMEOUT = 'timeout'
+
+
+class SliceCollectorNode(BaseNode):
+    """
+    Slice Collector Node - collects detection predictions from multiple image slices
+    and merges them into a single unified result.
+    
+    This node is designed to pair with SliceImageNode in workflows like:
+    
+    SliceImageNode -> Split -> YOLONode -> SliceCollectorNode
+    
+    It tracks slice messages by their msg_id from the original slice operation,
+    collects all predictions, transforms coordinates, and applies NMS.
+    """
+    display_name = 'Slice Collector'
+    info = str(_info)
+    icon = '📥'
+    category = 'vision'
+    color = '#FFA07A'
+    border_color = '#FF7F50'
+    text_color = '#000000'
+    input_count = 1
+    output_count = 1
+    
+    DEFAULT_CONFIG = {
+        NodeKeys.TIMEOUT: 5.0,
+        NodeKeys.NMS_THRESHOLD: 0.5,
+        NodeKeys.MATCH_METRIC: 'iou',
+        NodeKeys.CLASS_AGNOSTIC: False,
+        MessageKeys.DROP_MESSAGES: False
+    }
+    
+    properties = [
+        {
+            'name': NodeKeys.TIMEOUT,
+            'label': 'Collection Timeout (seconds)',
+            'type': 'number',
+            'default': DEFAULT_CONFIG[NodeKeys.TIMEOUT]
+        },
+        {
+            'name': NodeKeys.NMS_THRESHOLD,
+            'label': 'NMS IoU Threshold',
+            'type': 'number',
+            'default': DEFAULT_CONFIG[NodeKeys.NMS_THRESHOLD]
+        },
+        {
+            'name': NodeKeys.MATCH_METRIC,
+            'label': 'Match Metric',
+            'type': 'select',
+            'options': [
+                {'value': 'iou', 'label': 'IoU (Intersection over Union)'},
+                {'value': 'ios', 'label': 'IoS (Intersection over Smaller)'}
+            ],
+            'default': DEFAULT_CONFIG[NodeKeys.MATCH_METRIC]
+        },
+        {
+            'name': NodeKeys.CLASS_AGNOSTIC,
+            'label': 'Class Agnostic NMS',
+            'type': 'checkbox',
+            'default': DEFAULT_CONFIG[NodeKeys.CLASS_AGNOSTIC]
+        }
+    ]
+    
+    def __init__(self, node_id=None, name="slice_collector"):
+        super().__init__(node_id, name)
+        # Track collections by parent message ID
+        self._collections: Dict[str, Dict] = {}
+    
+    def _calculate_iou(self, box1: List[float], box2: List[float]) -> float:
+        """Calculate IoU between two bounding boxes in [x1, y1, x2, y2] format."""
+        x1_1, y1_1, x2_1, y2_1 = box1
+        x1_2, y1_2, x2_2, y2_2 = box2
+        
+        x1_i = max(x1_1, x1_2)
+        y1_i = max(y1_1, y1_2)
+        x2_i = min(x2_1, x2_2)
+        y2_i = min(y2_1, y2_2)
+        
+        if x2_i <= x1_i or y2_i <= y1_i:
+            return 0.0
+        
+        intersection = (x2_i - x1_i) * (y2_i - y1_i)
+        area1 = (x2_1 - x1_1) * (y2_1 - y1_1)
+        area2 = (x2_2 - x1_2) * (y2_2 - y1_2)
+        union = area1 + area2 - intersection
+        
+        return intersection / union if union > 0 else 0.0
+    
+    def _calculate_ios(self, box1: List[float], box2: List[float]) -> float:
+        """Calculate IoS (Intersection over Smaller) between two boxes."""
+        x1_1, y1_1, x2_1, y2_1 = box1
+        x1_2, y1_2, x2_2, y2_2 = box2
+        
+        x1_i = max(x1_1, x1_2)
+        y1_i = max(y1_1, y1_2)
+        x2_i = min(x2_1, x2_2)
+        y2_i = min(y2_1, y2_2)
+        
+        if x2_i <= x1_i or y2_i <= y1_i:
+            return 0.0
+        
+        intersection = (x2_i - x1_i) * (y2_i - y1_i)
+        area1 = (x2_1 - x1_1) * (y2_1 - y1_1)
+        area2 = (x2_2 - x1_2) * (y2_2 - y1_2)
+        smaller = min(area1, area2)
+        
+        return intersection / smaller if smaller > 0 else 0.0
+    
+    def _transform_detections(self, detections: List[Dict], offset: List[int]) -> List[Dict]:
+        """Transform detection coordinates from slice-local to original image coords."""
+        transformed = []
+        offset_x, offset_y = offset
+        
+        for det in detections:
+            new_det = det.copy()
+            bbox = det.get(NodeKeys.BBOX, [])
+            
+            if len(bbox) >= 4:
+                x1, y1, x2, y2 = bbox[:4]
+                new_det[NodeKeys.BBOX] = [x1 + offset_x, y1 + offset_y, x2 + offset_x, y2 + offset_y]
+                new_det[NodeKeys.SLICE_OFFSET] = offset
+            
+            transformed.append(new_det)
+        
+        return transformed
+    
+    def _nms(self, detections: List[Dict], nms_threshold: float, 
+             match_metric: str = NodeKeys.MATCH_METRIC, class_agnostic: bool = False) -> List[Dict]:
+        """Apply Non-Maximum Suppression to remove duplicate detections."""
+        if not detections:
+            return []
+        
+        sorted_dets = sorted(detections, key=lambda x: x.get('confidence', 0), reverse=True)
+        metric_fn = self._calculate_iou if match_metric == 'iou' else self._calculate_ios
+        
+        if class_agnostic:
+            groups = {'all': sorted_dets}
+        else:
+            groups: Dict[Any, List[Dict]] = {}
+            for det in sorted_dets:
+                class_id = det.get('class_id', det.get('class_name', 'unknown'))
+                if class_id not in groups:
+                    groups[class_id] = []
+                groups[class_id].append(det)
+        
+        kept = []
+        for class_id, group in groups.items():
+            keep_indices = []
+            
+            for i, det in enumerate(group):
+                should_keep = True
+                bbox_i = det.get(NodeKeys.BBOX, [])
+                
+                if len(bbox_i) < 4:
+                    continue
+                
+                for j in keep_indices:
+                    bbox_j = group[j].get(NodeKeys.BBOX, [])
+                    if len(bbox_j) < 4:
+                        continue
+                    
+                    if metric_fn(bbox_i, bbox_j) > nms_threshold:
+                        should_keep = False
+                        break
+                
+                if should_keep:
+                    keep_indices.append(i)
+            
+            kept.extend([group[i] for i in keep_indices])
+        
+        return kept
+    
+    def _cleanup_expired(self):
+        """Remove expired collections."""
+        timeout = self.get_config_float(NodeKeys.TIMEOUT, self.DEFAULT_CONFIG[NodeKeys.TIMEOUT])
+        current_time = time.time()
+        expired = [k for k, v in self._collections.items() 
+                   if current_time - v['start_time'] > timeout]
+        for k in expired:
+            del self._collections[k]
+    
+    def _process_collection(self, collection_id: str) -> Optional[Dict]:
+        """Process a complete collection and return merged result."""
+        collection = self._collections.get(collection_id)
+        if not collection:
+            return None
+        
+        nms_threshold = self.get_config_float(NodeKeys.NMS_THRESHOLD, self.DEFAULT_CONFIG[NodeKeys.NMS_THRESHOLD])
+        match_metric = self.config.get(NodeKeys.MATCH_METRIC, self.DEFAULT_CONFIG[NodeKeys.MATCH_METRIC])
+        class_agnostic = self.get_config_bool(NodeKeys.CLASS_AGNOSTIC, self.DEFAULT_CONFIG[NodeKeys.CLASS_AGNOSTIC])
+        
+        all_detections = []
+        full_image_detections = []
+        
+        for slice_data in collection[NodeKeys.SLICES].values():
+            offset = slice_data.get(NodeKeys.OFFSET, [0, 0])
+            detections = slice_data.get(MessageKeys.CV.DETECTIONS, [])
+            is_full_image = slice_data.get(NodeKeys.IS_FULL_IMAGE, False)
+            
+            if is_full_image:
+                full_image_detections.extend(detections)
+            else:
+                transformed = self._transform_detections(detections, offset)
+                all_detections.extend(transformed)
+        
+        # Apply NMS
+        merged = self._nms(all_detections, nms_threshold, match_metric, class_agnostic)
+        
+        # Combine with full image detections if present
+        if full_image_detections:
+            combined = merged + full_image_detections
+            final_detections = self._nms(combined, nms_threshold, match_metric, class_agnostic)
+        else:
+            final_detections = merged
+        
+        # Sort by confidence
+        final_detections = sorted(final_detections, key=lambda x: x.get(MessageKeys.CV.CONFIDENCE, 0), reverse=True)
+        
+        return {
+            MessageKeys.CV.DETECTIONS: final_detections,
+            MessageKeys.CV.DETECTION_COUNT: len(final_detections),
+            NodeKeys.ORIGINAL_WIDTH: collection.get(NodeKeys.ORIGINAL_WIDTH, 0),
+            NodeKeys.ORIGINAL_HEIGHT: collection.get(NodeKeys.ORIGINAL_HEIGHT, 0),
+            NodeKeys.SLICE_COUNT: len(collection[NodeKeys.SLICES]),
+            MessageKeys.CV.BBOX_FORMAT: 'xyxy',
+            MessageKeys.IMAGE.PATH: collection.get(MessageKeys.IMAGE.PATH)
+        }
+    
+    def on_input(self, msg: Dict[str, Any], input_index: int = 0):
+        """
+        Collect slice predictions and merge when complete.
+        
+        Expected msg format (from Split of SliceImageNode output, processed by YOLO):
+        {
+            'payload': {
+                'image': ...,
+                'detections': [...],
+                ...
+            },
+            'parts': {
+                'index': int,
+                'count': int,
+                'id': str  # Parent message ID
+            },
+            'slice_offset': [x, y],  # From SliceImageNode
+            'slice_index': int,
+            'is_full_image': bool,
+            'original_width': int,
+            'original_height': int
+        }
+        """
+        self._cleanup_expired()
+        
+        payload = msg.get(MessageKeys.PAYLOAD, {})
+        parts = msg.get('parts', {})
+        
+        # Get collection ID (from parts or generate one)
+        collection_id = parts.get('id', msg.get(MessageKeys.MSG_ID, 'default'))
+        expected_count = parts.get('count', 1)
+        slice_index = parts.get('index', msg.get(NodeKeys.SLICE_INDEX, 0))
+        
+        # Get slice metadata (check message level first, then payload)
+        offset = msg.get( NodeKeys.OFFSET, payload.get(NodeKeys.OFFSET, [0, 0]) if isinstance(payload, dict) else [0, 0])
+        slice_bbox = msg.get(NodeKeys.BBOX, payload.get(NodeKeys.BBOX, None) if isinstance(payload, dict) else None)
+        is_full_image = msg.get(NodeKeys.IS_FULL_IMAGE, payload.get(NodeKeys.IS_FULL_IMAGE, False) if isinstance(payload, dict) else False)
+        original_width = msg.get(NodeKeys.ORIGINAL_WIDTH, payload.get(NodeKeys.ORIGINAL_WIDTH, 0) if isinstance(payload, dict) else 0)
+        original_height = msg.get(NodeKeys.ORIGINAL_HEIGHT, payload.get(NodeKeys.ORIGINAL_HEIGHT, 0) if isinstance(payload, dict) else 0)
+        
+        # Get detections
+        if isinstance(payload, dict):
+            detections = payload.get(MessageKeys.CV.DETECTIONS, [])
+            image = payload.get(MessageKeys.IMAGE.PATH)
+        else:
+            detections = []
+            image = None
+        
+        # Initialize collection if needed
+        if collection_id not in self._collections:
+            self._collections[collection_id] = {
+                'start_time': time.time(),
+                NodeKeys.EXPECTED_COUNT: expected_count,
+                NodeKeys.SLICES: {},
+                NodeKeys.ORIGINAL_WIDTH: original_width,
+                NodeKeys.ORIGINAL_HEIGHT: original_height,
+                MessageKeys.ORIGINAL_MSG: msg.copy(),
+                MessageKeys.IMAGE.PATH: None
+            }
+        
+        collection = self._collections[collection_id]
+        
+        # Store slice data
+        collection[NodeKeys.SLICES][slice_index] = {
+            MessageKeys.CV.DETECTIONS: detections,
+            NodeKeys.OFFSET: offset,
+            NodeKeys.IS_FULL_IMAGE: is_full_image
+        }
+        
+        # Store full image if this is the full image slice
+        if is_full_image and image is not None:
+            collection[MessageKeys.IMAGE.PATH] = image
+        
+        # Update dimensions if available
+        if original_width > 0:
+            collection[NodeKeys.ORIGINAL_WIDTH] = original_width
+        if original_height > 0:
+            collection[NodeKeys.ORIGINAL_HEIGHT] = original_height
+        
+        # Check if complete
+        if len(collection[NodeKeys.SLICES]) >= collection[NodeKeys.EXPECTED_COUNT]:
+            result = self._process_collection(collection_id)
+            
+            if result:
+                out_msg = collection[MessageKeys.ORIGINAL_MSG].copy()
+                out_msg[MessageKeys.PAYLOAD] = result
+                out_msg[MessageKeys.TOPIC] = out_msg.get(MessageKeys.TOPIC, 'merged_predictions')
+                
+                # Clean up parts since we've merged
+                out_msg.pop('parts', None)
+                
+                self.send(out_msg)
+            
+            # Clean up
+            del self._collections[collection_id]