diff --git a/base/detector_base.py b/base/detector_base.py index e2124e2..0ea1c53 100644 --- a/base/detector_base.py +++ b/base/detector_base.py @@ -1,16 +1,39 @@ from abc import ABC, abstractmethod import numpy as np from enum import Enum +import cv2 + +class SharedProcessing: + """Handles common image processing operations used by multiple detectors""" + + @staticmethod + def get_shared_data(frame: np.ndarray, blur_size: int = 5) -> dict: + """ + Performs common image processing operations once + Args: + frame: RGB numpy array + blur_size: Size of Gaussian blur kernel + Returns: + dict of preprocessed data + """ + gray = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY) + blurred = cv2.GaussianBlur(gray, (blur_size, blur_size), 0) + return { + "gray": gray, + "blurred": blurred, + "frame": frame # Original frame + } class Detector(ABC): """Base class for image region detectors""" @classmethod - def preprocess(cls, frame: np.ndarray) -> dict: + def preprocess(cls, frame: np.ndarray, shared_data: dict = None) -> dict: """ Optional preprocessing step for the entire frame Args: frame: Full frame data as numpy array + shared_data: Dict containing shared preprocessed data Returns: dict containing preprocessed data to be shared across ROIs """ @@ -29,7 +52,8 @@ class Detector(ABC): current_frame: Current frame data for ROI region mask: ROI mask state: Detector state dict containing: - - preprocessed: dict of preprocessed data from preprocess() if available + - preprocessed: dict of preprocessed data from preprocess() + - shared: dict of shared preprocessed data - Any other temporal or persistent data needed by the detector Returns: (detection_value, visualization_mask) @@ -47,4 +71,17 @@ class ROIAction(Enum): SUBTRACT = "subtract" # Subtract value when detected MULTIPLY = "multiply" # Multiply by value when detected DIVIDE = "divide" # Divide by value when detected - SET = "set" \ No newline at end of file + SET = "set" # Set to value when detected + +# Define descriptions as a separate dictionary +ROIAction.descriptions = { + ROIAction.TOGGLE.value: "Toggles between min/max values on each detection", + ROIAction.MOMENTARY.value: "Outputs max value while detection is active, min otherwise", + ROIAction.TRIGGER.value: "Outputs max value once per detection event", + ROIAction.COUNTER.value: "Counts the number of detection events", + ROIAction.ADD.value: "Adds the specified value when detection occurs", + ROIAction.SUBTRACT.value: "Subtracts the specified value when detection occurs", + ROIAction.MULTIPLY.value: "Multiplies by the specified value when detection occurs", + ROIAction.DIVIDE.value: "Divides by the specified value when detection occurs", + ROIAction.SET.value: "Sets output directly to the specified value when detection occurs" +} \ No newline at end of file diff --git a/controls/detector_controls.py b/controls/detector_controls.py index eb24f08..1c3e6f3 100644 --- a/controls/detector_controls.py +++ b/controls/detector_controls.py @@ -1,5 +1,5 @@ from ..base.control_base import ControlNodeBase -from ..base.detector_base import ROIAction +from ..base.detector_base import ROIAction, SharedProcessing import time import numpy as np import torch @@ -71,6 +71,9 @@ class DetectionControlBase(ControlNodeBase): current_frame = (image[0] * 255).cpu().numpy().astype(np.uint8) detection_mask = np.zeros_like(current_frame[:,:,0], dtype=np.float32) + # Generate shared preprocessed data once + shared_data = SharedProcessing.get_shared_data(current_frame) + # Collect all unique detector types from ROI chain detector_types = set() current_roi = roi_chain @@ -78,17 +81,12 @@ class DetectionControlBase(ControlNodeBase): 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 + preprocessed_data[detector_type] = detector_type.preprocess(current_frame, shared_data) # Process ROI chain current_roi = roi_chain @@ -97,28 +95,30 @@ class DetectionControlBase(ControlNodeBase): y_min, x_min, y_max, x_max = bounds roi_id = str(bounds) - # Create unique detector state for this ROI + # Get detector state detector_state = state["detector_states"].setdefault(roi_id, {}) detector_state.update({ - 'y_offset': y_min, # Add ROI coordinates to state - 'x_offset': x_min + 'y_offset': y_min, + 'x_offset': x_min, + 'shared': shared_data # Add shared data }) + + # Add preprocessed data if available + detector_type = current_roi["detector"].__class__ + if detector_type in preprocessed_data: + detector_state["preprocessed"] = preprocessed_data[detector_type] + + # Get ROI state 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] - + # Run detection 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 # Now each ROI has its own state + detector_state ) # Update visualization mask @@ -135,7 +135,7 @@ class DetectionControlBase(ControlNodeBase): ) current_roi = current_roi["next"] - + self.set_state(state) return state["current_value"], torch.from_numpy(detection_mask).unsqueeze(0) diff --git a/detectors/detector_nodes.py b/detectors/detector_nodes.py index 2bea2a1..b386780 100644 --- a/detectors/detector_nodes.py +++ b/detectors/detector_nodes.py @@ -8,14 +8,33 @@ class RegionOfInterest(ControlNodeBase): @classmethod def INPUT_TYPES(cls): 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}), + + # Format tooltip using ROIAction descriptions + action_tooltip = "Available actions:\n" + "\n".join( + f"• {action.value}: {ROIAction.descriptions[action.value]}" + for action in ROIAction + ) + + inputs["required"].update({ + "mask": ("MASK", { + "tooltip": "Binary mask defining the region of interest" + }), + "detector": ("DETECTOR", { + "tooltip": "Configured detector node (Motion or Brightness) to process this region" + }), + "action": (list(action.value for action in ROIAction), { + "tooltip": action_tooltip + }), + "value": ("FLOAT", { + "default": 0.1, + "step": 0.01, + "tooltip": "Value to use in the selected action (amount to add/subtract/multiply/divide/set)" + }), }) - inputs["optional"] = { # Add optional inputs - "next_roi": ("ROI",) + inputs["optional"] = { + "next_roi": ("ROI", { + "tooltip": "Optional connection to chain multiple ROIs together" + }) } return inputs diff --git a/detectors/detectors.py b/detectors/detectors.py index efc92cb..49c87f7 100644 --- a/detectors/detectors.py +++ b/detectors/detectors.py @@ -4,108 +4,84 @@ 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): + def setup(self, threshold=0.1, blur_size=5): self.threshold = threshold self.blur_size = blur_size + self.prev_frame = None - def detect(self, current_frame: np.ndarray, mask: np.ndarray, state: dict) -> tuple[float, np.ndarray]: - # 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) - - # Calculate detection value - roi_area = np.sum(mask > 0.5) - if roi_area > 0: - detection = np.sum(thresh) / (255.0 * roi_area) - 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) + def preprocess(cls, frame: np.ndarray, shared_data: dict) -> dict: + """Use shared grayscale and blurred data""" return { - 'gray_frame': gray + "blurred": shared_data["blurred"] # Reuse shared blurred frame } + def detect(self, current_frame: np.ndarray, mask: np.ndarray, state: dict) -> tuple[float, np.ndarray]: + # Get preprocessed data + preprocessed = state.get("preprocessed", {}) + shared = state.get("shared", {}) + y_offset = state.get('y_offset', 0) + x_offset = state.get('x_offset', 0) + + # Get the blurred ROI from preprocessed data + full_blurred = preprocessed["blurred"] + # Extract the correct ROI region from the full blurred image + current_blurred = full_blurred[y_offset:y_offset+current_frame.shape[0], + x_offset:x_offset+current_frame.shape[1]] + + # Initialize or update previous frame for this ROI + roi_key = f"{y_offset},{x_offset}" + if "prev_blurred" not in state: + state["prev_blurred"] = {} + if roi_key not in state["prev_blurred"]: + state["prev_blurred"][roi_key] = current_blurred + return 0.0, np.zeros_like(current_blurred, dtype=np.float32) + + # Calculate frame difference + diff = cv2.absdiff(current_blurred, state["prev_blurred"][roi_key]) + _, thresh = cv2.threshold(diff, self.threshold * 255, 255, cv2.THRESH_BINARY) + + # Apply ROI mask + thresh = thresh * (mask > 0.5) + + # Update previous frame + state["prev_blurred"][roi_key] = current_blurred.copy() + + # Calculate motion magnitude + roi_area = np.sum(mask > 0.5) + if roi_area > 0: + motion_magnitude = np.sum(thresh) / (255.0 * roi_area) + else: + motion_magnitude = 0.0 + + return motion_magnitude, thresh.astype(np.float32) / 255.0 + +class BrightnessDetector(Detector): def setup(self, threshold: float = 0.5, use_average: bool = True): self.threshold = threshold self.use_average = use_average + @classmethod + def preprocess(cls, frame: np.ndarray, shared_data: dict) -> dict: + """Use shared grayscale data""" + return { + "gray": shared_data["gray"] # Reuse shared grayscale frame + } + def detect(self, current_frame: np.ndarray, mask: np.ndarray, state: dict) -> tuple[float, np.ndarray]: - # 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())}") + # Get preprocessed data + preprocessed = state.get("preprocessed", {}) + shared = state.get("shared", {}) + y_offset = state.get('y_offset', 0) + x_offset = state.get('x_offset', 0) - # 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) + # Get the grayscale ROI from preprocessed data + full_gray = preprocessed["gray"] + # Extract the correct ROI region from the full grayscale image + gray = full_gray[y_offset:y_offset+current_frame.shape[0], + x_offset:x_offset+current_frame.shape[1]] + # Apply mask masked_gray = gray * (mask > 0.5) # Calculate brightness