parent
51638877ac
commit
bccc70b371
7 changed files with 0 additions and 557 deletions
@ -1,200 +0,0 @@ |
|||||||
import os |
|
||||||
import sys |
|
||||||
import time |
|
||||||
|
|
||||||
import torch |
|
||||||
from diffusers import AutoencoderTiny, StableDiffusionPipeline |
|
||||||
from diffusers.utils import load_image |
|
||||||
|
|
||||||
sys.path.insert(0, os.path.abspath('../StreamDiffusion')) |
|
||||||
|
|
||||||
from streamdiffusion import StreamDiffusion |
|
||||||
from streamdiffusion.image_utils import postprocess_image |
|
||||||
|
|
||||||
from utils.viewer import receive_images |
|
||||||
|
|
||||||
from utils.wrapper import StreamDiffusionWrapper |
|
||||||
from threading import Thread |
|
||||||
|
|
||||||
|
|
||||||
from multiprocessing import Process, Queue, get_context |
|
||||||
|
|
||||||
from perlin import perlin_2d, rand_perlin_2d, rand_perlin_2d_octaves, perlin_2d_octaves |
|
||||||
from scene_prompt import surreal_prompt_parts |
|
||||||
from scene_prompt import surreal_prompts |
|
||||||
from scene_prompt import regret_prompts |
|
||||||
|
|
||||||
from spout_util import send_spout_image, get_spout_image |
|
||||||
|
|
||||||
from osc import start_osc_server |
|
||||||
|
|
||||||
|
|
||||||
import fire |
|
||||||
|
|
||||||
def image_generation_process( |
|
||||||
queue: Queue, |
|
||||||
fps_queue: Queue, |
|
||||||
prompt_queue: Queue, |
|
||||||
input_queue: Queue, |
|
||||||
# prompt: str, |
|
||||||
model_id_or_path: str, |
|
||||||
)-> None: |
|
||||||
# stream = StreamDiffusionWrapper( |
|
||||||
# model_id_or_path=model_id_or_path, |
|
||||||
# lora_dict=None, |
|
||||||
# t_index_list=[0, 16, 32, 45], |
|
||||||
# frame_buffer_size=1, |
|
||||||
# width=512, |
|
||||||
# height=512, |
|
||||||
# warmup=10, |
|
||||||
# acceleration="xformers", |
|
||||||
# mode="txt2img", |
|
||||||
# use_denoising_batch=False, |
|
||||||
# cfg_type="none", |
|
||||||
# seed=2, |
|
||||||
# ) |
|
||||||
stream = StreamDiffusionWrapper( |
|
||||||
model_id_or_path=model_id_or_path, |
|
||||||
t_index_list=[0], |
|
||||||
frame_buffer_size=1, |
|
||||||
warmup=10, |
|
||||||
acceleration="tensorrt", |
|
||||||
use_lcm_lora=False, |
|
||||||
mode="img2img", |
|
||||||
cfg_type="none", |
|
||||||
use_denoising_batch=True, |
|
||||||
output_type="pil", |
|
||||||
) |
|
||||||
|
|
||||||
start_prompt = "A glowing, vintage phone booth standing in surreal landscapes across different scene" |
|
||||||
# Prepare the stream |
|
||||||
stream.prepare( |
|
||||||
prompt=start_prompt, |
|
||||||
num_inference_steps=50, |
|
||||||
) |
|
||||||
|
|
||||||
# Prepare image |
|
||||||
# init_image = load_image("example.png").resize((512, 512)) |
|
||||||
|
|
||||||
# Warmup >= len(t_index_list) x frame_buffer_size |
|
||||||
# for _ in range(stream.batch_size - 1): |
|
||||||
# stream() |
|
||||||
|
|
||||||
previous_output = None |
|
||||||
idx=0 |
|
||||||
last_time = time.time() |
|
||||||
|
|
||||||
while True: |
|
||||||
# try: |
|
||||||
start_time = time.time() |
|
||||||
# x_output = stream(image=previous_output) |
|
||||||
# x_output=stream.stream.txt2img_sd_turbo(1).cpu() |
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
input_image= input_queue.get(block=True) |
|
||||||
# input_image = stream.preprocess_image('input.png') |
|
||||||
|
|
||||||
|
|
||||||
# Check if a new prompt is available in the prompt_queue |
|
||||||
if not prompt_queue.empty(): |
|
||||||
new_prompt = prompt_queue.get(block=False) |
|
||||||
if new_prompt: |
|
||||||
x_output = stream.img2img(image=input_image, prompt=new_prompt) |
|
||||||
print(f"Received new prompt from queue: {new_prompt}") |
|
||||||
else: |
|
||||||
# Use the current prompt if no new prompt is available |
|
||||||
x_output = stream.img2img(image=input_image) |
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
# preprocessed_image =stream.postprocess_image(x_output) |
|
||||||
|
|
||||||
queue.put(x_output, block=False) |
|
||||||
|
|
||||||
# queue.put(preprocessed_image, block=False) |
|
||||||
|
|
||||||
# Calculate FPS |
|
||||||
elapsed_time = time.time() - start_time |
|
||||||
fps = 1 / elapsed_time if elapsed_time > 0 else float('inf') |
|
||||||
fps_queue.put(fps) |
|
||||||
|
|
||||||
# x_output = (x_output + 1) / 2 # Scale from [-1, 1] to [0, 1] |
|
||||||
# x_output = torch.clamp(x_output, 0, 1) |
|
||||||
# previous_output = x_output |
|
||||||
|
|
||||||
# except KeyboardInterrupt: |
|
||||||
# print(f"fps: {fps}") |
|
||||||
# return |
|
||||||
|
|
||||||
|
|
||||||
def main()-> None: |
|
||||||
|
|
||||||
try: |
|
||||||
ctx = get_context('spawn') |
|
||||||
queue = Queue() |
|
||||||
fps_queue = Queue() |
|
||||||
# noise_queue = Queue() |
|
||||||
spout_in_queue = Queue() |
|
||||||
|
|
||||||
# prompt = "A surreal landscapes" |
|
||||||
# prompt=regret_prompts[0] |
|
||||||
|
|
||||||
prompt_queue = Queue() |
|
||||||
|
|
||||||
# model_id_or_path = "KBlueLeaf/kohaku-v2.1" |
|
||||||
model_id_or_path = "stabilityai/sd-turbo" |
|
||||||
|
|
||||||
|
|
||||||
# start_osc_server(prompt_queue) |
|
||||||
process_osc = ctx.Process( |
|
||||||
target=start_osc_server, |
|
||||||
args=(prompt_queue,) |
|
||||||
) |
|
||||||
process_osc.start() |
|
||||||
|
|
||||||
print("Starting spout input process") |
|
||||||
process_spout_in = ctx.Process( |
|
||||||
target=get_spout_image, |
|
||||||
args=(spout_in_queue, 512, 512), |
|
||||||
) |
|
||||||
process_spout_in.start() |
|
||||||
|
|
||||||
|
|
||||||
print("Starting image generation process") |
|
||||||
process_gen= ctx.Process( |
|
||||||
target=image_generation_process, |
|
||||||
args=(queue, fps_queue, prompt_queue, spout_in_queue, model_id_or_path), |
|
||||||
) |
|
||||||
process_gen.start() |
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
# process_show=ctx.Process(target=receive_images, args=(queue, fps_queue)) |
|
||||||
# process_show.start() |
|
||||||
|
|
||||||
# print("Starting spout output process") |
|
||||||
process_spout_out=ctx.Process(target=send_spout_image, args=(queue, 512, 512)) |
|
||||||
process_spout_out.start() |
|
||||||
|
|
||||||
|
|
||||||
process_gen.join() |
|
||||||
# process_spout_in.join() |
|
||||||
process_spout_out.join() |
|
||||||
process_osc.join() |
|
||||||
|
|
||||||
|
|
||||||
except KeyboardInterrupt: |
|
||||||
print("Process interrupted") |
|
||||||
|
|
||||||
process_gen.terminate() |
|
||||||
# process_spout_in.terminate() |
|
||||||
process_spout_out.terminate() |
|
||||||
process_osc.terminate() |
|
||||||
|
|
||||||
return |
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
if __name__ == "__main__": |
|
||||||
fire.Fire(main) |
|
||||||
@ -1,23 +0,0 @@ |
|||||||
import argparse |
|
||||||
import math |
|
||||||
|
|
||||||
from pythonosc.dispatcher import Dispatcher |
|
||||||
from pythonosc import osc_server |
|
||||||
|
|
||||||
OSC_PORT = 8787 |
|
||||||
|
|
||||||
|
|
||||||
def start_osc_server(queue): |
|
||||||
|
|
||||||
def onReceivePrompt(address, *args): |
|
||||||
prompt = " ".join(args) |
|
||||||
print(f"Received prompt: {prompt}") |
|
||||||
queue.put(prompt) |
|
||||||
|
|
||||||
dispatcher = Dispatcher() |
|
||||||
dispatcher.map("/prompt", onReceivePrompt) |
|
||||||
|
|
||||||
|
|
||||||
server = osc_server.ThreadingOSCUDPServer(("localhost", OSC_PORT), dispatcher) |
|
||||||
print(f"OSC server is running on port {OSC_PORT}") |
|
||||||
server.serve_forever() |
|
||||||
@ -1,69 +0,0 @@ |
|||||||
import torch |
|
||||||
import math |
|
||||||
|
|
||||||
def rand_perlin_2d(shape, res, fade = lambda t: 6*t**5 - 15*t**4 + 10*t**3): |
|
||||||
delta = (res[0] / shape[0], res[1] / shape[1]) |
|
||||||
d = (shape[0] // res[0], shape[1] // res[1]) |
|
||||||
|
|
||||||
grid = torch.stack(torch.meshgrid(torch.arange(0, res[0], delta[0]), torch.arange(0, res[1], delta[1])), dim = -1) % 1 |
|
||||||
angles = 2*math.pi*torch.rand(res[0]+1, res[1]+1) |
|
||||||
gradients = torch.stack((torch.cos(angles), torch.sin(angles)), dim = -1) |
|
||||||
|
|
||||||
tile_grads = lambda slice1, slice2: gradients[slice1[0]:slice1[1], slice2[0]:slice2[1]].repeat_interleave(d[0], 0).repeat_interleave(d[1], 1) |
|
||||||
dot = lambda grad, shift: (torch.stack((grid[:shape[0],:shape[1],0] + shift[0], grid[:shape[0],:shape[1], 1] + shift[1] ), dim = -1) * grad[:shape[0], :shape[1]]).sum(dim = -1) |
|
||||||
|
|
||||||
n00 = dot(tile_grads([0, -1], [0, -1]), [0, 0]) |
|
||||||
n10 = dot(tile_grads([1, None], [0, -1]), [-1, 0]) |
|
||||||
n01 = dot(tile_grads([0, -1],[1, None]), [0, -1]) |
|
||||||
n11 = dot(tile_grads([1, None], [1, None]), [-1,-1]) |
|
||||||
t = fade(grid[:shape[0], :shape[1]]) |
|
||||||
return math.sqrt(2) * torch.lerp(torch.lerp(n00, n10, t[..., 0]), torch.lerp(n01, n11, t[..., 0]), t[..., 1]) |
|
||||||
|
|
||||||
def rand_perlin_2d_octaves(shape, res, octaves=1, persistence=0.5): |
|
||||||
noise = torch.zeros(shape) |
|
||||||
frequency = 1 |
|
||||||
amplitude = 1 |
|
||||||
for _ in range(octaves): |
|
||||||
noise += amplitude * rand_perlin_2d(shape, (frequency*res[0], frequency*res[1])) |
|
||||||
frequency *= 2 |
|
||||||
amplitude *= persistence |
|
||||||
return noise |
|
||||||
|
|
||||||
def perlin_2d(shape, res, seed, fade=lambda t: 6*t**5 - 15*t**4 + 10*t**3): |
|
||||||
delta = (res[0] / shape[0], res[1] / shape[1]) |
|
||||||
d = (shape[0] // res[0], shape[1] // res[1]) |
|
||||||
|
|
||||||
grid = torch.stack(torch.meshgrid(torch.arange(0, res[0], delta[0]), torch.arange(0, res[1], delta[1])), dim=-1) % 1 |
|
||||||
base_seed = int(seed) |
|
||||||
frac_seed = seed - base_seed |
|
||||||
|
|
||||||
torch.manual_seed(base_seed) |
|
||||||
angles_base = 2 * math.pi * torch.rand(res[0] + 1, res[1] + 1) |
|
||||||
gradients_base = torch.stack((torch.cos(angles_base), torch.sin(angles_base)), dim=-1) |
|
||||||
|
|
||||||
torch.manual_seed(base_seed + 1) |
|
||||||
angles_next = 2 * math.pi * torch.rand(res[0] + 1, res[1] + 1) |
|
||||||
gradients_next = torch.stack((torch.cos(angles_next), torch.sin(angles_next)), dim=-1) |
|
||||||
|
|
||||||
gradients = (1 - frac_seed) * gradients_base + frac_seed * gradients_next |
|
||||||
|
|
||||||
tile_grads = lambda slice1, slice2: gradients[slice1[0]:slice1[1], slice2[0]:slice2[1]].repeat_interleave(d[0], 0).repeat_interleave(d[1], 1) |
|
||||||
dot = lambda grad, shift: (torch.stack((grid[:shape[0], :shape[1], 0] + shift[0], grid[:shape[0], :shape[1], 1] + shift[1]), dim=-1) * grad[:shape[0], :shape[1]]).sum(dim=-1) |
|
||||||
|
|
||||||
n00 = dot(tile_grads([0, -1], [0, -1]), [0, 0]) |
|
||||||
n10 = dot(tile_grads([1, None], [0, -1]), [-1, 0]) |
|
||||||
n01 = dot(tile_grads([0, -1], [1, None]), [0, -1]) |
|
||||||
n11 = dot(tile_grads([1, None], [1, None]), [-1, -1]) |
|
||||||
t = fade(grid[:shape[0], :shape[1]]) |
|
||||||
return math.sqrt(2) * torch.lerp(torch.lerp(n00, n10, t[..., 0]), torch.lerp(n01, n11, t[..., 0]), t[..., 1]) |
|
||||||
|
|
||||||
|
|
||||||
def perlin_2d_octaves(shape, res, seed, octaves=1, persistence=0.5, fade=lambda t: 6*t**5 - 15*t**4 + 10*t**3): |
|
||||||
noise = torch.zeros(shape) |
|
||||||
frequency = 1 |
|
||||||
amplitude = 1 |
|
||||||
for i in range(octaves): |
|
||||||
noise += amplitude * perlin_2d(shape, (frequency * res[0], frequency * res[1]), seed + i, fade) |
|
||||||
frequency *= 2 |
|
||||||
amplitude *= persistence |
|
||||||
return noise |
|
||||||
@ -1,42 +0,0 @@ |
|||||||
surreal_prompts = [ |
|
||||||
"a surreal landscape of floating islands under a glowing sky", |
|
||||||
"an ethereal valley where waterfalls rise into the clouds", |
|
||||||
"a dreamlike desert with mirrored sand and hovering stones", |
|
||||||
"an endless ocean reflecting fractured moons and stars", |
|
||||||
"a neon-lit canyon with levitating ruins and glowing mist", |
|
||||||
"a twilight forest where the trees grow upside-down", |
|
||||||
"a luminous terrain with bioluminescent plants and crystal arches", |
|
||||||
"a gravity-defying mountain range spiraling into the void", |
|
||||||
"a shattered realm of glass bridges and hovering towers", |
|
||||||
"an alien world lit by pulsating constellations and fluid geometry" |
|
||||||
] |
|
||||||
|
|
||||||
surreal_prompt_parts = [ |
|
||||||
"a surreal landscape", |
|
||||||
"with floating islands", |
|
||||||
"glowing waterfalls", |
|
||||||
"neon-colored skies", |
|
||||||
"mirror-like desert ground", |
|
||||||
"levitating rocks", |
|
||||||
"upside-down trees", |
|
||||||
"ancient ruins suspended in air", |
|
||||||
"bioluminescent flora", |
|
||||||
"shattered moons overhead", |
|
||||||
"a path of glass tiles", |
|
||||||
"crystal towers emitting soft hums", |
|
||||||
"gravity-defying rivers", |
|
||||||
"alien constellations glowing brightly" |
|
||||||
] |
|
||||||
|
|
||||||
regret_prompts = [ |
|
||||||
"a lone figure standing in a vast, empty desert at dusk", |
|
||||||
"fractured mirrors scattered across the sand, reflecting different memories", |
|
||||||
"a withered tree growing upside down from the sky, its roots dripping ink", |
|
||||||
"floating clocks melting into the horizon, ticking backwards", |
|
||||||
"ghostly silhouettes walking in reverse, retracing forgotten steps", |
|
||||||
"a house half-submerged in water, its windows glowing faintly with past laughter", |
|
||||||
"the sky opens into a tunnel of old photographs slowly burning at the edges", |
|
||||||
"giant stone hands reaching out from the earth, trying to grasp something lost", |
|
||||||
"an ocean made of letters never sent, waves crashing with whispered apologies", |
|
||||||
"a child version of the figure stands alone, staring at the adult with distant eyes" |
|
||||||
] |
|
||||||
@ -1,125 +0,0 @@ |
|||||||
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) |
|
||||||
@ -1,98 +0,0 @@ |
|||||||
import os |
|
||||||
import sys |
|
||||||
import threading |
|
||||||
import time |
|
||||||
import tkinter as tk |
|
||||||
from multiprocessing import Queue |
|
||||||
from typing import List |
|
||||||
from PIL import Image, ImageTk |
|
||||||
from streamdiffusion.image_utils import postprocess_image |
|
||||||
|
|
||||||
sys.path.append(os.path.join(os.path.dirname(__file__), "..", "..")) |
|
||||||
|
|
||||||
|
|
||||||
def update_image(image_data: Image.Image, label: tk.Label) -> None: |
|
||||||
""" |
|
||||||
Update the image displayed on a Tkinter label. |
|
||||||
|
|
||||||
Parameters |
|
||||||
---------- |
|
||||||
image_data : Image.Image |
|
||||||
The image to be displayed. |
|
||||||
label : tk.Label |
|
||||||
The labels where the image will be updated. |
|
||||||
""" |
|
||||||
width = 512 |
|
||||||
height = 512 |
|
||||||
tk_image = ImageTk.PhotoImage(image_data, size=width) |
|
||||||
label.configure(image=tk_image, width=width, height=height) |
|
||||||
label.image = tk_image # keep a reference |
|
||||||
|
|
||||||
def _receive_images( |
|
||||||
queue: Queue, fps_queue: Queue, label: tk.Label, fps_label: tk.Label |
|
||||||
) -> None: |
|
||||||
""" |
|
||||||
Continuously receive images from a queue and update the labels. |
|
||||||
|
|
||||||
Parameters |
|
||||||
---------- |
|
||||||
queue : Queue |
|
||||||
The queue to receive images from. |
|
||||||
fps_queue : Queue |
|
||||||
The queue to put the calculated fps. |
|
||||||
label : tk.Label |
|
||||||
The label to update with images. |
|
||||||
fps_label : tk.Label |
|
||||||
The label to show fps. |
|
||||||
""" |
|
||||||
while True: |
|
||||||
try: |
|
||||||
if not queue.empty(): |
|
||||||
label.after( |
|
||||||
0, |
|
||||||
update_image, |
|
||||||
postprocess_image(queue.get(block=False), output_type="pil")[0], |
|
||||||
label, |
|
||||||
) |
|
||||||
if not fps_queue.empty(): |
|
||||||
fps_label.config(text=f"FPS: {fps_queue.get(block=False):.2f}") |
|
||||||
|
|
||||||
time.sleep(0.0005) |
|
||||||
except KeyboardInterrupt: |
|
||||||
return |
|
||||||
|
|
||||||
|
|
||||||
def receive_images(queue: Queue, fps_queue: Queue) -> None: |
|
||||||
""" |
|
||||||
Setup the Tkinter window and start the thread to receive images. |
|
||||||
|
|
||||||
Parameters |
|
||||||
---------- |
|
||||||
queue : Queue |
|
||||||
The queue to receive images from. |
|
||||||
fps_queue : Queue |
|
||||||
The queue to put the calculated fps. |
|
||||||
""" |
|
||||||
root = tk.Tk() |
|
||||||
root.title("Image Viewer") |
|
||||||
label = tk.Label(root) |
|
||||||
fps_label = tk.Label(root, text="FPS: 0") |
|
||||||
label.grid(column=0) |
|
||||||
fps_label.grid(column=1) |
|
||||||
|
|
||||||
def on_closing(): |
|
||||||
print("window closed") |
|
||||||
root.quit() # stop event loop |
|
||||||
return |
|
||||||
|
|
||||||
thread = threading.Thread( |
|
||||||
target=_receive_images, args=(queue, fps_queue, label, fps_label), daemon=True |
|
||||||
) |
|
||||||
thread.start() |
|
||||||
|
|
||||||
try: |
|
||||||
root.protocol("WM_DELETE_WINDOW", on_closing) |
|
||||||
root.mainloop() |
|
||||||
except KeyboardInterrupt: |
|
||||||
return |
|
||||||
|
|
||||||
Loading…
Reference in new issue