You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
125 lines
4.1 KiB
125 lines
4.1 KiB
import torch
|
|
import SpoutGL
|
|
|
|
from itertools import islice, cycle, repeat
|
|
import array
|
|
from random import randint
|
|
import time
|
|
from OpenGL import GL
|
|
|
|
from multiprocessing import Queue
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
|
|
TARGET_FPS = 30
|
|
SEND_WIDTH = 512
|
|
SEND_HEIGHT = 512
|
|
|
|
alpha_cache = np.full((512, 512, 1), 255, dtype=np.uint8)
|
|
|
|
|
|
|
|
def spout_buffer_to_tensor(buffer, width, height):
|
|
# np_buffer = np.asarray(buffer, dtype=np.uint8)
|
|
np_buffer=np.frombuffer(buffer, dtype=np.uint8)
|
|
image_bgra = np_buffer.reshape((height, width, 4))
|
|
|
|
image_rgb = image_bgra[..., [2, 1, 0]]
|
|
image_float = image_rgb.astype(np.float32) / 255.0
|
|
# image_normalized = (image_float * 2.0) - 1.0
|
|
tensor = torch.from_numpy(image_float).permute(2, 0, 1)
|
|
|
|
del np_buffer # Free memory
|
|
del image_bgra # Free memory
|
|
del image_rgb # Free memory
|
|
del image_float # Free memory
|
|
|
|
|
|
return tensor.unsqueeze(0)
|
|
|
|
|
|
def get_spout_image(queue, wwidth: int, wheight: int) -> None:
|
|
with SpoutGL.SpoutReceiver() as receiver:
|
|
receiver.setReceiverName("Spout DX11 Sender")
|
|
image_bgra = np.zeros((SEND_HEIGHT, SEND_WIDTH, 4), dtype=np.uint8)
|
|
|
|
while True:
|
|
result = receiver.receiveImage(image_bgra, GL.GL_RGBA, False, 0)
|
|
# print("Receive result", result)
|
|
|
|
if receiver.isUpdated():
|
|
continue
|
|
# width = receiver.getSenderWidth()
|
|
# height = receiver.getSenderHeight()
|
|
# image_bgra = array.array('B', [0] * (width * height * 4)) # Correctly reallocate buffer with updated size
|
|
# print("Spout Receiver updated, Buffer size", width, height)
|
|
|
|
# if buffer and result and not SpoutGL.helpers.isBufferEmpty(buffer):
|
|
if SpoutGL.helpers.isBufferEmpty(image_bgra):
|
|
continue
|
|
# pixels=spout_buffer_to_tensor(buffer, width, height)
|
|
# print("get_spout_image", pixels.shape)
|
|
|
|
image_rgb_array= image_bgra[:, :, [2, 1, 0]]
|
|
pixels=Image.fromarray(image_rgb_array, 'RGB')
|
|
queue.put(pixels, block=False)
|
|
|
|
|
|
# Wait until the next frame is ready
|
|
# Wait time is in milliseconds; note that 0 will return immediately
|
|
# receiver.waitFrameSync("SpoutSender", 10000)
|
|
|
|
|
|
|
|
|
|
def randcolor():
|
|
return randint(0, 255)
|
|
|
|
|
|
def tensor_to_spout_image(tensor):
|
|
image = tensor.squeeze(0)
|
|
if image.device.type != "cpu":
|
|
image = image.cpu()
|
|
image = image.permute(1, 2, 0).numpy()
|
|
|
|
if image.min() < 0:
|
|
image = (image + 1) / 2 # Scale from [-1, 1] to [0, 1]
|
|
image = np.clip(image * 255, 0, 255).astype(np.uint8)
|
|
|
|
# h, w, _ = image_np.shape
|
|
# alpha = np.full((h, w, 1), 255, dtype=np.uint8)
|
|
image_rgba = np.concatenate((image, alpha_cache), axis=-1)
|
|
image_bgra = image_rgba[..., [2, 1, 0, 3]]
|
|
|
|
del image # Free memory
|
|
|
|
return np.ascontiguousarray(image_bgra) # Ensure the array is contiguous in memory
|
|
|
|
def send_spout_image(queue: Queue, width: int, height: int)->None:
|
|
|
|
with SpoutGL.SpoutSender() as sender:
|
|
sender.setSenderName("StreamDiffusion")
|
|
|
|
while True:
|
|
|
|
# Check if there are images in the queue
|
|
if not queue.empty():
|
|
output_image = queue.get(block=False)
|
|
# pixels = tensor_to_spout_image(image)
|
|
|
|
output_bgr_array = np.array(output_image, dtype=np.uint8)[:, :, ::-1]
|
|
output_bgra_array = np.zeros((SEND_HEIGHT, SEND_WIDTH, 4), dtype=np.uint8)
|
|
output_bgra_array[:, :, :3] = output_bgr_array
|
|
output_bgra_array[:, :, 3] = 255
|
|
buffer = output_bgra_array
|
|
|
|
|
|
result = sender.sendImage(buffer, width, height, GL.GL_RGBA, False, 0)
|
|
# print("Send result", result)
|
|
|
|
# Indicate that a frame is ready to read
|
|
sender.setFrameSync("StreamDiffusion")
|
|
|
|
# Wait for next send attempt
|
|
# time.sleep(1./TARGET_FPS) |