from ..base.control_base import ControlNodeBase from enum import Enum import numpy as np import torch import cv2 import time class ROIAction(Enum): # Behavioral actions TOGGLE = "toggle" # Toggles between min/max values MOMENTARY = "momentary" # Outputs max while motion detected TRIGGER = "trigger" # Triggers once per motion event COUNTER = "counter" # Counts motion events # Mathematical actions ADD = "add" # Add value when motion detected SUBTRACT = "subtract" # Subtract value when motion detected MULTIPLY = "multiply" # Multiply by value when motion detected DIVIDE = "divide" # Divide by value when motion detected SET = "set" # Set to specific value when motion detected class ROINode(ControlNodeBase): """Defines a single region of interest and its action""" @classmethod def INPUT_TYPES(cls): return { "required": { "mask": ("MASK",), # ROI defined by mask "action": (list(action.value for action in ROIAction),), "value": ("FLOAT", { "default": 0.1, "step": 0.01, "tooltip": "Value to use for mathematical operations" }), }, "optional": { "next_roi": ("ROI",) # Chain to next ROI } } RETURN_TYPES = ("ROI",) FUNCTION = "define_roi" CATEGORY = "real-time/control/motion" def define_roi(self, mask, action, value, next_roi=None): # Convert mask to numpy for bounding box calculation mask_np = mask[0].cpu().numpy() # Find non-zero coordinates in mask coords = np.nonzero(mask_np) if len(coords[0]) == 0: # Empty mask y_min = x_min = y_max = x_max = 0 else: y_min, y_max = coords[0].min(), coords[0].max() x_min, x_max = coords[1].min(), coords[1].max() roi_data = { "mask": mask_np, "bounds": (y_min, x_min, y_max, x_max), "action": action, "value": value, # Store the operation value "next": next_roi } return (roi_data,) def update(self, *args, **kwargs): return self.define_roi(*args, **kwargs) class MotionController(ControlNodeBase): """Processes motion detection for a chain of ROIs and manages their states""" @classmethod def INPUT_TYPES(cls): inputs = super().INPUT_TYPES() inputs["required"].update({ "image": ("IMAGE",), "roi_chain": ("ROI",), "threshold": ("FLOAT", { "default": 0.1, "min": 0.0, "max": 1.0, "step": 0.01, "tooltip": "Motion detection threshold" }), "blur_size": ("INT", { "default": 5, "min": 1, "max": 21, "step": 2, "tooltip": "Size of Gaussian blur kernel" }), "minimum_value": ("FLOAT", { "default": 0.0, "tooltip": "Minimum output value" }), "maximum_value": ("FLOAT", { "default": 1.0, "tooltip": "Maximum output value" }), "starting_value": ("FLOAT", { "default": 0.0, "tooltip": "Initial output value" }) }) return inputs RETURN_TYPES = ("FLOAT", "MASK") FUNCTION = "process_motion" CATEGORY = "real-time/control/motion" def process_motion(self, image, roi_chain, threshold, blur_size, minimum_value, maximum_value, starting_value, always_execute=True): # Get or initialize state with new caching fields state = self.get_state({ "prev_frame": None, "prev_frame_blurred": None, # Cache blurred previous frame "roi_states": {}, "current_value": starting_value, "last_cleanup": time.time() # Track last state cleanup }) # Convert image tensor to numpy array and prepare current frame current_frame = (image[0] * 255).cpu().numpy().astype(np.uint8) current_gray = cv2.cvtColor(current_frame, cv2.COLOR_RGB2GRAY) current_blurred = cv2.GaussianBlur(current_gray, (blur_size, blur_size), 0) # Initialize motion mask motion_mask = np.zeros_like(current_gray, dtype=np.float32) if state["prev_frame"] is not None: # Collect active ROIs for state cleanup active_rois = set() # Process each ROI in the chain current_roi = roi_chain while current_roi is not None: mask = current_roi["mask"] bounds = current_roi["bounds"] action = current_roi["action"] roi_id = str(bounds) active_rois.add(roi_id) # Extract ROI regions once y_min, x_min, y_max, x_max = bounds roi_region = slice(y_min, y_max+1), slice(x_min, x_max+1) # Process ROI using cached blurred frames roi_current = current_blurred[roi_region] roi_prev = state["prev_frame_blurred"][roi_region] roi_mask = mask[roi_region] # Detect motion in ROI diff = cv2.absdiff(roi_current, roi_prev) _, thresh = cv2.threshold(diff, threshold * 255, 255, cv2.THRESH_BINARY) thresh = thresh * (roi_mask > 0.5) # Update motion mask motion_mask[roi_region] = np.maximum( motion_mask[roi_region], thresh / 255.0 ) # Calculate motion magnitude roi_area = np.sum(roi_mask > 0.5) motion_detected = np.sum(thresh) / (255 * roi_area) > threshold if roi_area > 0 else False # Get or initialize ROI state roi_state = state["roi_states"].get(roi_id, {"active": False, "count": 0}) value = current_roi["value"] # Process actions with optimized value updates if motion_detected and not roi_state["active"]: roi_state["active"] = True # Update current value based on action type if action == ROIAction.ADD.value: state["current_value"] = min(maximum_value, state["current_value"] + value) elif action == ROIAction.SUBTRACT.value: state["current_value"] = max(minimum_value, state["current_value"] - value) elif action == ROIAction.MULTIPLY.value: state["current_value"] = min(maximum_value, state["current_value"] * value) elif action == ROIAction.DIVIDE.value: if value != 0: state["current_value"] = max(minimum_value, state["current_value"] / value) elif action == ROIAction.SET.value: state["current_value"] = max(minimum_value, min(maximum_value, value)) elif action == ROIAction.TOGGLE.value: state["current_value"] = maximum_value if state["current_value"] == minimum_value else minimum_value elif action == ROIAction.TRIGGER.value: state["current_value"] = maximum_value elif action == ROIAction.COUNTER.value: roi_state["count"] += 1 state["current_value"] = minimum_value + ( roi_state["count"] % (int((maximum_value - minimum_value) + 1)) ) elif action == ROIAction.MOMENTARY.value: state["current_value"] = maximum_value if motion_detected else minimum_value elif not motion_detected: roi_state["active"] = False if action == ROIAction.TRIGGER.value: state["current_value"] = minimum_value state["roi_states"][roi_id] = roi_state current_roi = current_roi["next"] # Cleanup inactive ROI states periodically (every 60 seconds) if time.time() - state["last_cleanup"] > 60: state["roi_states"] = {k: v for k, v in state["roi_states"].items() if k in active_rois} state["last_cleanup"] = time.time() # Update state with current frame data state["prev_frame"] = current_gray state["prev_frame_blurred"] = current_blurred self.set_state(state) # Convert mask to tensor mask_tensor = torch.from_numpy(motion_mask).unsqueeze(0) return (state["current_value"], mask_tensor) def update(self, *args, **kwargs): return self.process_motion(*args, **kwargs) class RoundingMode(Enum): ROUND = "round" # Standard rounding FLOOR = "floor" # Round down CEIL = "ceil" # Round up TRUNCATE = "truncate" # Remove decimal part class IntegerMotionController(MotionController): """Processes motion detection for ROIs and outputs integer values""" @classmethod def INPUT_TYPES(cls): inputs = super().INPUT_TYPES() # Convert float parameters to int inputs["required"]["minimum_value"] = ("INT", { "default": 0, "tooltip": "Minimum output value" }) inputs["required"]["maximum_value"] = ("INT", { "default": 100, "tooltip": "Maximum output value" }) inputs["required"]["starting_value"] = ("INT", { "default": 0, "tooltip": "Initial output value" }) # Add rounding mode parameter inputs["required"]["rounding_mode"] = ( [mode.value for mode in RoundingMode], {"default": RoundingMode.ROUND.value} ) return inputs RETURN_TYPES = ("INT", "MASK") FUNCTION = "process_motion" CATEGORY = "real-time/control/motion" def _apply_rounding(self, value, mode): """Apply the specified rounding mode to convert float to int""" if mode == RoundingMode.ROUND.value: return int(round(value)) elif mode == RoundingMode.FLOOR.value: return int(np.floor(value)) elif mode == RoundingMode.CEIL.value: return int(np.ceil(value)) else: # TRUNCATE return int(value) def process_motion(self, image, roi_chain, threshold, blur_size, minimum_value, maximum_value, starting_value, rounding_mode, always_execute=True): # Call parent class method with float conversion of bounds float_result, mask = super().process_motion( image, roi_chain, threshold, blur_size, float(minimum_value), float(maximum_value), float(starting_value), always_execute ) # Convert result to integer using specified rounding mode int_result = self._apply_rounding(float_result, rounding_mode) # Ensure result stays within bounds int_result = max(minimum_value, min(maximum_value, int_result)) return (int_result, mask)