nodes/ryanontheinside-ComfyUI_RealtimeNodes
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

cached detection frames with update to pattern, refactored always_execute

Browse Source
This commit is contained in:
RyanOnTheInside
2024-12-30 13:54:40 +00:00
parent c162b36a01
commit 3bed5bf6cc
5 changed files with 171 additions and 70 deletions
Download Patch File Download Diff File Expand all files Collapse all files
base/control_base.py
+2 -2
View File
@@ -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."
}),
},
}
base/detector_base.py
+15 -2
View File
@@ -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)
"""
controls/detector_controls.py
+38 -12
View File
@@ -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
detectors/detector_nodes.py
+46 -46
View File
@@ -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)
detectors/detectors.py
+70 -8
View File
@@ -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