124 lines
4.4 KiB
Python
124 lines
4.4 KiB
Python
from ..base.control_base import ControlNodeBase
|
|
import time
|
|
import numpy as np
|
|
import cv2
|
|
import torch
|
|
|
|
class FPSMonitor(ControlNodeBase):
|
|
"""Generates an FPS overlay as an image and mask"""
|
|
|
|
RETURN_TYPES = ("IMAGE", "MASK")
|
|
FUNCTION = "update"
|
|
CATEGORY = "real-time/control/utility"
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
inputs = super().INPUT_TYPES()
|
|
inputs["required"].update({
|
|
"width": ("INT", {
|
|
"default": 512,
|
|
"min": 64,
|
|
"max": 4096,
|
|
"step": 1
|
|
}),
|
|
"height": ("INT", {
|
|
"default": 512,
|
|
"min": 64,
|
|
"max": 4096,
|
|
"step": 1
|
|
}),
|
|
"text_color": ("INT", {
|
|
"default": 255,
|
|
"min": 0,
|
|
"max": 255,
|
|
"step": 1
|
|
}),
|
|
"text_size": ("FLOAT", {
|
|
"default": 1.0,
|
|
"min": 0.1,
|
|
"max": 5.0,
|
|
"step": 0.1
|
|
}),
|
|
"window_size": ("INT", {
|
|
"default": 60,
|
|
"min": 1,
|
|
"max": 300,
|
|
"step": 1,
|
|
"tooltip": "Number of frames to average over"
|
|
})
|
|
})
|
|
return inputs
|
|
|
|
def update(self, width: int, height: int, text_color: int, text_size: float, window_size: int, always_execute=True):
|
|
current_time = time.time()
|
|
state = self.get_state({
|
|
"last_time": current_time,
|
|
"frame_times": [],
|
|
"last_fps": 0,
|
|
"last_avg_fps": 0,
|
|
"cached_image": None,
|
|
"cached_mask": None,
|
|
"last_dims": (0, 0),
|
|
"last_color": 0,
|
|
"last_size": 0
|
|
})
|
|
|
|
# Calculate FPS
|
|
delta_time = current_time - state["last_time"]
|
|
current_fps = 1.0 / delta_time if delta_time > 0 else 0
|
|
|
|
# Update rolling window
|
|
state["frame_times"].append(delta_time)
|
|
if len(state["frame_times"]) > window_size:
|
|
state["frame_times"].pop(0)
|
|
|
|
# Calculate average FPS
|
|
avg_delta = sum(state["frame_times"]) / len(state["frame_times"])
|
|
average_fps = 1.0 / avg_delta if avg_delta > 0 else 0
|
|
|
|
# Check if we need to redraw
|
|
needs_redraw = (
|
|
abs(current_fps - state["last_fps"]) > 1.0 or # FPS changed significantly
|
|
abs(average_fps - state["last_avg_fps"]) > 0.5 or # Avg FPS changed significantly
|
|
state["last_dims"] != (width, height) or # Dimensions changed
|
|
state["last_color"] != text_color or # Color changed
|
|
state["last_size"] != text_size or # Text size changed
|
|
state["cached_image"] is None # No cache exists
|
|
)
|
|
|
|
if needs_redraw:
|
|
# Create text
|
|
fps_text = f"FPS: {current_fps:.1f} (avg: {average_fps:.1f})"
|
|
|
|
# Create image and mask using numpy/cv2
|
|
image = np.zeros((height, width, 3), dtype=np.uint8)
|
|
mask = np.zeros((height, width), dtype=np.uint8)
|
|
|
|
# Position text in top-left with padding
|
|
padding = 10
|
|
pos = (padding, padding + 20) # Add vertical offset for cv2's text positioning
|
|
|
|
# Draw text - much simpler with cv2!
|
|
cv2.putText(image, fps_text, pos, cv2.FONT_HERSHEY_SIMPLEX, text_size, (text_color, text_color, text_color), 1)
|
|
cv2.putText(mask, fps_text, pos, cv2.FONT_HERSHEY_SIMPLEX, text_size, 255, 1)
|
|
|
|
# Convert to float32 and normalize
|
|
image = image.astype(np.float32) / 255.0
|
|
mask = mask.astype(np.float32) / 255.0
|
|
|
|
# Convert to torch tensors with correct shapes
|
|
state["cached_image"] = torch.from_numpy(image[None, ...]) # BHWC
|
|
state["cached_mask"] = torch.from_numpy(mask[None, ...]) # BHW
|
|
|
|
# Update cache parameters
|
|
state["last_fps"] = current_fps
|
|
state["last_avg_fps"] = average_fps
|
|
state["last_dims"] = (width, height)
|
|
state["last_color"] = text_color
|
|
state["last_size"] = text_size
|
|
|
|
state["last_time"] = current_time
|
|
self.set_state(state)
|
|
|
|
return (state["cached_image"], state["cached_mask"])
|