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

performance

Browse Source
This commit is contained in:
RyanOnTheInside
2024-12-30 16:43:20 +00:00
parent 3bed5bf6cc
commit 3f5d903cda
4 changed files with 149 additions and 117 deletions
Download Patch File Download Diff File Expand all files Collapse all files
base/detector_base.py
+40 -3
View File
@@ -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"
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"
}
controls/detector_controls.py
+19 -19
View File
@@ -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)
detectors/detector_nodes.py
+26 -7
View File
@@ -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
detectors/detectors.py
+64 -88
View File
@@ -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