83 lines
2.8 KiB
Python
83 lines
2.8 KiB
Python
import os
|
|
import h5py
|
|
import numpy as np
|
|
import time
|
|
import json
|
|
|
|
from bec_lib.core import RedisConnector, BECMessage
|
|
|
|
|
|
def load_data() -> tuple:
|
|
"""
|
|
|
|
Returns
|
|
"""
|
|
proj_nr = 180
|
|
basedir = f"/Users/janwyzula/PSI/test_data_reduced/projection_{proj_nr:06d}/"
|
|
|
|
metadata_name = f"/Users/janwyzula/PSI/test_data_reduced/projection_{proj_nr:06d}.json"
|
|
with open(metadata_name) as file:
|
|
metadata = json.load(file)
|
|
|
|
filenames = [fname for fname in os.listdir(basedir) if fname.endswith(".h5")]
|
|
filenames.sort()
|
|
|
|
for ii, fname in enumerate(filenames):
|
|
with h5py.File(os.path.join(basedir, fname), "r") as h5file:
|
|
if ii == 0:
|
|
q = h5file["q"][...].T.squeeze()
|
|
norm_sum = h5file["norm_sum"][...]
|
|
data = np.zeros((len(filenames), *h5file["I_all"][...].shape))
|
|
data[ii, ...] = h5file["I_all"][...]
|
|
|
|
return data, q, norm_sum, metadata
|
|
|
|
|
|
def _get_projection_keys(producer):
|
|
keys = producer.keys("px_stream/projection_*")
|
|
if not keys:
|
|
return []
|
|
return keys
|
|
|
|
|
|
def send_data(data, q, norm_sum, bec_producer, metadata, proj_nr) -> None:
|
|
""""""
|
|
start = time.time()
|
|
|
|
keys = _get_projection_keys(bec_producer)
|
|
pipe = bec_producer.pipeline()
|
|
proj_numbers = set(key.decode().split("px_stream/projection_")[1].split("/")[0] for key in keys)
|
|
if len(proj_numbers) > 5:
|
|
for entry in sorted(proj_numbers)[0:-5]:
|
|
for key in bec_producer.keys(f"px_stream/projection_{entry}/*"):
|
|
bec_producer.delete(topic=key, pipe=pipe)
|
|
print(f"Deleting {key}")
|
|
|
|
# Add new data
|
|
return_dict = {"metadata": metadata, "q": q, "norm_sum": norm_sum}
|
|
msg = BECMessage.DeviceMessage(signals=return_dict).dumps()
|
|
bec_producer.set_and_publish(f"px_stream/projection_{proj_nr}/metadata", msg=msg, pipe=pipe)
|
|
return_dict = {"proj_nr": proj_nr}
|
|
msg = BECMessage.DeviceMessage(signals=return_dict).dumps()
|
|
bec_producer.set_and_publish(f"px_stream/proj_nr", msg=msg, pipe=pipe)
|
|
pipe.execute()
|
|
for line in range(data.shape[0]):
|
|
return_dict = {"data": data[line, ...]}
|
|
msg = BECMessage.DeviceMessage(signals=return_dict).dumps()
|
|
print(f"Sending line {line}")
|
|
bec_producer.rpush(topic=f"px_stream/projection_{proj_nr}/data", msgs=msg)
|
|
print(f"Time to send {time.time()-start} seconds")
|
|
print(f"Rate {data.shape[0]/(time.time()-start)} Hz")
|
|
print(f"Data volume {data.nbytes/1e6} MB")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
data, q, norm_sum, metadata = load_data()
|
|
bec_producer = RedisConnector(["localhost:6379"]).producer()
|
|
proj_nr = 180
|
|
while True:
|
|
send_data(data, q, norm_sum, bec_producer, metadata, proj_nr=proj_nr)
|
|
time.sleep(30)
|
|
bec_producer.delete(topic=f"px_stream/projection_{proj_nr}/data:val")
|
|
# proj_nr = proj_nr + 1
|