60 lines
2.1 KiB
Python
60 lines
2.1 KiB
Python
import os
|
|
import time
|
|
import redis
|
|
import io
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
# Connect to Redis
|
|
redis_host = os.getenv("redis_host", "localhost")
|
|
redis_port = int(os.getenv("redis_port", 6379))
|
|
redis_conn = redis.Redis(host=redis_host, port=redis_port, db=0, decode_responses=False)
|
|
|
|
# Environment variables
|
|
stream_label = os.getenv("stream_label", "cropped_stream")
|
|
out_folder = os.getenv("out_folder", "/app/out_folder")
|
|
stream_label_queue = f"{stream_label}_cubes"
|
|
|
|
# Ensure output folder exists
|
|
os.makedirs(out_folder, exist_ok=True)
|
|
|
|
def fromRedisList(queue_label):
|
|
compressed_data = None
|
|
retry = 0
|
|
while compressed_data == None:
|
|
compressed_data = redis_conn.lpop(queue_label)
|
|
if compressed_data:
|
|
retry = 0
|
|
print(f"Popped data from queue: {queue_label}")
|
|
buffer = io.BytesIO(compressed_data)
|
|
loaded_data = np.load(buffer)
|
|
data_list = [loaded_data[key] for key in loaded_data.files]
|
|
return data_list
|
|
else:
|
|
retry += 1
|
|
if retry % 50 == 0:
|
|
print(f"Queue {queue_label} empty for {retry/50} seconds")
|
|
time.sleep(1/50.0)
|
|
if retry > 1000:
|
|
raise(f'Queue {queue_label} 20s empty')
|
|
|
|
def save_frames_as_gif(frames, output_filename, duration=1):
|
|
images = [Image.fromarray(frame[:, :, ::-1]) for frame in frames]
|
|
images[0].save(output_filename, save_all=True, append_images=images[1:], duration=duration, loop=0)
|
|
|
|
if __name__ == "__main__":
|
|
counter = 0
|
|
print(f"[INFO] Starting GIF writer for queue: {stream_label}")
|
|
while True:
|
|
frames = fromRedisList(stream_label_queue)
|
|
if frames:
|
|
buffer = os.path.join(out_folder, f"{stream_label}_{counter:05}.npz")
|
|
filename = os.path.join(out_folder, f"{stream_label}_{counter:05}.gif")
|
|
save_frames_as_gif(frames, filename)
|
|
np.savez_compressed(buffer, *frames) # Use *data_list to unpack the list into arguments for savez_compressed
|
|
|
|
print(f"[INFO] Saved as npz: {buffer}")
|
|
counter += 1
|
|
else:
|
|
time.sleep(0.1)
|