From 3bed5bf6ccbfbb8b9f549257e60d61ec74ca6782 Mon Sep 17 00:00:00 2001 From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com> Date: Mon, 30 Dec 2024 13:54:40 +0000 Subject: [PATCH] cached detection frames with update to pattern, refactored always_execute --- base/control_base.py | 4 +- base/detector_base.py | 17 ++++++- controls/detector_controls.py | 50 ++++++++++++++----- detectors/detector_nodes.py | 92 +++++++++++++++++------------------ detectors/detectors.py | 78 ++++++++++++++++++++++++++--- 5 files changed, 171 insertions(+), 70 deletions(-) diff --git a/base/control_base.py b/base/control_base.py index f14fada..13eba20 100644 --- a/base/control_base.py +++ b/base/control_base.py @@ -36,8 +36,8 @@ class ControlNodeBase(ABC): return { "required": { "always_execute": ("BOOLEAN", { - "default": True, - "tooltip": "When enabled, the node updates every execution" + "default": False, + "tooltip": "When enabled, the node updates every execution. This is primarily for use INSIDE of ComfyUI. Generally, you should set this to FALSE for real-time applications, like ComfyStream." }), }, } diff --git a/base/detector_base.py b/base/detector_base.py index c8fbf30..e2124e2 100644 --- a/base/detector_base.py +++ b/base/detector_base.py @@ -5,6 +5,17 @@ from enum import Enum class Detector(ABC): """Base class for image region detectors""" + @classmethod + def preprocess(cls, frame: np.ndarray) -> dict: + """ + Optional preprocessing step for the entire frame + Args: + frame: Full frame data as numpy array + Returns: + dict containing preprocessed data to be shared across ROIs + """ + return {} + @abstractmethod def setup(self, **kwargs): """Initialize detector with parameters""" @@ -15,9 +26,11 @@ class Detector(ABC): """ Analyze frame region and return detection value and visualization Args: - current_frame: Current frame data + current_frame: Current frame data for ROI region mask: ROI mask - state: Detector state dict for temporal data + state: Detector state dict containing: + - preprocessed: dict of preprocessed data from preprocess() if available + - Any other temporal or persistent data needed by the detector Returns: (detection_value, visualization_mask) """ diff --git a/controls/detector_controls.py b/controls/detector_controls.py index dfad218..eb24f08 100644 --- a/controls/detector_controls.py +++ b/controls/detector_controls.py @@ -9,16 +9,12 @@ class DetectionControlBase(ControlNodeBase): @classmethod def INPUT_TYPES(cls): - return { - "required": { - "image": ("IMAGE",), - "roi_chain": ("ROI",), - "always_execute": ("BOOLEAN", { - "default": True, - "tooltip": "When enabled, the node updates every execution" - }) - } - } + inputs = super().INPUT_TYPES() # Get base inputs including always_execute + inputs["required"].update({ + "image": ("IMAGE",), + "roi_chain": ("ROI",), + }) + return inputs RETURN_TYPES = ("MASK",) CATEGORY = "real-time/control/detection" @@ -74,6 +70,25 @@ class DetectionControlBase(ControlNodeBase): # Convert image tensor to numpy current_frame = (image[0] * 255).cpu().numpy().astype(np.uint8) detection_mask = np.zeros_like(current_frame[:,:,0], dtype=np.float32) + + # Collect all unique detector types from ROI chain + detector_types = set() + current_roi = roi_chain + while current_roi is not None: + detector = current_roi["detector"] + detector_types.add(detector.__class__) + current_roi = current_roi["next"] + + print("Found detector types:", [d.__name__ for d in detector_types]) + + # Run preprocessing for each detector type + preprocessed_data = {} + for detector_type in detector_types: + print(f"Running preprocess for {detector_type.__name__}") + if hasattr(detector_type, 'preprocess'): + data = detector_type.preprocess(current_frame) + print(f"Preprocessed data keys: {list(data.keys())}") + preprocessed_data[detector_type] = data # Process ROI chain current_roi = roi_chain @@ -82,17 +97,28 @@ class DetectionControlBase(ControlNodeBase): y_min, x_min, y_max, x_max = bounds roi_id = str(bounds) - # Process detection + # Create unique detector state for this ROI detector_state = state["detector_states"].setdefault(roi_id, {}) + detector_state.update({ + 'y_offset': y_min, # Add ROI coordinates to state + 'x_offset': x_min + }) roi_state = state["roi_states"].setdefault(roi_id, { "active": False, "count": 0 }) + # Add preprocessed data to detector state + detector_type = current_roi["detector"].__class__ + if detector_type in preprocessed_data: + print(f"\nProcessing ROI {roi_id}") + print(f"Found preprocessed data: {list(preprocessed_data[detector_type].keys())}") + detector_state["preprocessed"] = preprocessed_data[detector_type] + detection_value, viz_mask = current_roi["detector"].detect( current_frame[y_min:y_max+1, x_min:x_max+1], current_roi["mask"][y_min:y_max+1, x_min:x_max+1], - detector_state + detector_state # Now each ROI has its own state ) # Update visualization mask diff --git a/detectors/detector_nodes.py b/detectors/detector_nodes.py index 8b097fd..2bea2a1 100644 --- a/detectors/detector_nodes.py +++ b/detectors/detector_nodes.py @@ -7,23 +7,23 @@ import numpy as np class RegionOfInterest(ControlNodeBase): @classmethod def INPUT_TYPES(cls): - return { - "required": { - "mask": ("MASK",), - "detector": ("DETECTOR",), # Takes configured detector node output - "action": (list(action.value for action in ROIAction),), - "value": ("FLOAT", {"default": 0.1, "step": 0.01}), - }, - "optional": { - "next_roi": ("ROI",) - } + inputs = super().INPUT_TYPES() # Get base inputs + inputs["required"].update({ # Update with additional inputs + "mask": ("MASK",), + "detector": ("DETECTOR",), # Takes configured detector node output + "action": (list(action.value for action in ROIAction),), + "value": ("FLOAT", {"default": 0.1, "step": 0.01}), + }) + inputs["optional"] = { # Add optional inputs + "next_roi": ("ROI",) } + return inputs RETURN_TYPES = ("ROI",) FUNCTION = "update" CATEGORY = "real-time/control/detection" - def update(self, mask, detector, action, value, next_roi=None): + def update(self, mask, detector, action, value, always_execute=False, next_roi=None): """Implements required update method from ControlNodeBase""" mask_np = mask[0].cpu().numpy() @@ -50,30 +50,30 @@ class MotionDetectorNode(ControlNodeBase): @classmethod def INPUT_TYPES(cls): - return { - "required": { - "threshold": ("FLOAT", { - "default": 0.1, - "min": 0.0, - "max": 1.0, - "step": 0.01, - "tooltip": "Motion detection sensitivity" - }), - "blur_size": ("INT", { - "default": 5, - "min": 1, - "max": 21, - "step": 2, - "tooltip": "Size of blur kernel for noise reduction" - }) - } - } + inputs = super().INPUT_TYPES() # Get base inputs + inputs["required"].update({ + "threshold": ("FLOAT", { + "default": 0.1, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "Motion detection sensitivity" + }), + "blur_size": ("INT", { + "default": 5, + "min": 1, + "max": 21, + "step": 2, + "tooltip": "Size of blur kernel for noise reduction" + }) + }) + return inputs RETURN_TYPES = ("DETECTOR",) FUNCTION = "update" CATEGORY = "real-time/control/detection" - def update(self, threshold, blur_size): + def update(self, threshold, blur_size, always_execute=False): """Implements required update method from ControlNodeBase""" detector = MotionDetector() detector.setup(threshold=threshold, blur_size=blur_size) @@ -85,27 +85,27 @@ class BrightnessDetectorNode(ControlNodeBase): @classmethod def INPUT_TYPES(cls): - return { - "required": { - "threshold": ("FLOAT", { - "default": 0.5, - "min": 0.0, - "max": 1.0, - "step": 0.01, - "tooltip": "Brightness threshold for visualization" - }), - "use_average": ("BOOLEAN", { - "default": True, - "tooltip": "Use average brightness instead of maximum" - }) - } - } + inputs = super().INPUT_TYPES() # Get base inputs + inputs["required"].update({ + "threshold": ("FLOAT", { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "Brightness threshold for visualization" + }), + "use_average": ("BOOLEAN", { + "default": True, + "tooltip": "Use average brightness instead of maximum" + }) + }) + return inputs RETURN_TYPES = ("DETECTOR",) FUNCTION = "update" CATEGORY = "real-time/control/detection" - def update(self, threshold, use_average): + def update(self, threshold, use_average, always_execute=False): """Implements required update method from ControlNodeBase""" detector = BrightnessDetector() detector.setup(threshold=threshold, use_average=use_average) diff --git a/detectors/detectors.py b/detectors/detectors.py index f0b2f8b..efc92cb 100644 --- a/detectors/detectors.py +++ b/detectors/detectors.py @@ -4,29 +4,60 @@ import cv2 from ..base.detector_base import Detector class MotionDetector(Detector): + @classmethod + def preprocess(cls, current_frame): + """Preprocess the entire frame for motion detection""" + curr_gray = cv2.cvtColor(current_frame, cv2.COLOR_RGB2GRAY) + print(f"Preprocessed frame shape: {curr_gray.shape}") + return { + 'gray_frame': curr_gray + } + def setup(self, threshold: float = 0.1, blur_size: int = 5): self.threshold = threshold self.blur_size = blur_size def detect(self, current_frame: np.ndarray, mask: np.ndarray, state: dict) -> tuple[float, np.ndarray]: - # Convert and blur current frame - curr_gray = cv2.cvtColor(current_frame, cv2.COLOR_RGB2GRAY) + # Log input shapes and ROI info + print(f"\nMotion Detector:") + print(f"ROI frame shape: {current_frame.shape}") + print(f"ROI mask shape: {mask.shape}") + print(f"State keys: {list(state.keys())}") + + # Use preprocessed grayscale frame if available + preprocessed = state.get('preprocessed', {}) + if 'gray_frame' in preprocessed: + full_gray = preprocessed['gray_frame'] + print(f"Full preprocessed gray shape: {full_gray.shape}") + + # Get the ROI slice from the preprocessed frame using the correct offset + # The ROI coordinates are stored in the state or can be derived from the frame shape + y_offset = state.get('y_offset', 0) # We need to pass these from the control base + x_offset = state.get('x_offset', 0) + print(f"Using offsets: y={y_offset}, x={x_offset}") + + curr_gray = full_gray[y_offset:y_offset + current_frame.shape[0], + x_offset:x_offset + current_frame.shape[1]] + print(f"Sliced gray ROI shape: {curr_gray.shape}") + else: + print("No preprocessed data found, computing grayscale for ROI") + curr_gray = cv2.cvtColor(current_frame, cv2.COLOR_RGB2GRAY) + + # Apply blur curr_blur = cv2.GaussianBlur(curr_gray, (self.blur_size, self.blur_size), 0) # Get previous frame from state prev_blur = state.get('prev_frame_blurred') if prev_blur is None: + print("First frame, initializing previous frame") state['prev_frame_blurred'] = curr_blur return 0.0, np.zeros_like(curr_gray, dtype=np.float32) - + # Calculate motion diff = cv2.absdiff(curr_blur, prev_blur) _, thresh = cv2.threshold(diff, self.threshold * 255, 255, cv2.THRESH_BINARY) thresh = thresh * (mask > 0.5) - # Update state - state['prev_frame_blurred'] = curr_blur - # Calculate detection value roi_area = np.sum(mask > 0.5) if roi_area > 0: @@ -34,16 +65,47 @@ class MotionDetector(Detector): else: detection = 0.0 + # Update state + state['prev_frame_blurred'] = curr_blur + + print(f"Motion detection value: {detection:.3f}") + print(f"ROI area: {roi_area}") + print(f"Sum of motion pixels: {np.sum(thresh)}") + return detection, thresh / 255.0 class BrightnessDetector(Detector): + @classmethod + def preprocess(cls, current_frame): + """Preprocess the entire frame for brightness detection""" + # Convert to grayscale once + gray = cv2.cvtColor(current_frame, cv2.COLOR_RGB2GRAY) + return { + 'gray_frame': gray + } + def setup(self, threshold: float = 0.5, use_average: bool = True): self.threshold = threshold self.use_average = use_average def detect(self, current_frame: np.ndarray, mask: np.ndarray, state: dict) -> tuple[float, np.ndarray]: - # Convert to grayscale - gray = cv2.cvtColor(current_frame, cv2.COLOR_RGB2GRAY) + # Debug: Log state contents + print(f"Detector state keys: {list(state.keys())}") + if 'preprocessed' in state: + print(f"Preprocessed data keys: {list(state['preprocessed'].keys())}") + + # Use preprocessed grayscale frame if available + preprocessed = state.get('preprocessed', {}) + if 'gray_frame' in preprocessed: + # Use the pre-computed grayscale frame + gray = preprocessed['gray_frame'][ + current_frame.shape[0]:current_frame.shape[0] + current_frame.shape[0], + current_frame.shape[1]:current_frame.shape[1] + current_frame.shape[1] + ] + else: + # Fallback to computing grayscale for this region + gray = cv2.cvtColor(current_frame, cv2.COLOR_RGB2GRAY) + masked_gray = gray * (mask > 0.5) # Calculate brightness