Spaces:
Running
Running
File size: 1,589 Bytes
3bb804c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 |
import __main__
BaseNode = __main__.BaseNode
QtGui = __main__.QtGui
import cv2
import numpy as np
class DisplayNode(BaseNode):
"""
Displays an image input directly.
"""
NODE_CATEGORY = "Output"
NODE_COLOR = QtGui.QColor(100, 100, 150)
def __init__(self):
super().__init__()
self.node_title = "Display"
self.inputs = {'image_in': 'image'}
self.outputs = {}
self.image = np.zeros((256, 256, 3), dtype=np.uint8)
def step(self):
img_in = self.get_blended_input('image_in', 'first')
if img_in is not None:
# Convert to 0-255 uint8 BGR
if img_in.dtype == np.float32 or img_in.dtype == np.float64:
img = (np.clip(img_in, 0, 1) * 255).astype(np.uint8)
else:
img = img_in.astype(np.uint8)
# Handle grayscale
if img.ndim == 2:
img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
# Handle RGB
elif img.shape[2] == 3:
# Assuming input is RGB, convert to BGR for display
# img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
# ^ Let's assume the host handles RGB, if not, uncomment this
pass
self.image = cv2.resize(img, (256, 256))
else:
self.image = (self.image * 0.9).astype(np.uint8) # Fade out
def get_display_image(self):
return QtGui.QImage(self.image.data, 256, 256, 256*3, QtGui.QImage.Format.Format_RGB888) |