added function to generate the isbead image from tractrac's hdf5 output file

main.py

@@ -3,6 +3,8 @@ import numpy as np
 import cv2
 from pathlib import Path
 import math
+import h5py
+import time
 
 
 def create_slope_image(image_width: int, image_height: int, p1_angle: float, p1_radius: float, bg_color: float = 0.0, fg_color: float = 1.0):
@@ -29,9 +31,84 @@ def create_slope_image(image_width: int, image_height: int, p1_angle: float, p1_
     return slope_image
 
 
+class TracData():
+
+    def __init__(self):
+        self.pts = None
+        self.th = None
+        self.num_frames = 0
+        self.image_size = None
+
+
+def trac_data_to_hdf5(trac_data: TracData, hdf5_file_path: Path):
+    version = '3.0 (22/05/2021) | J. Heyman'
+    f = h5py.File(hdf5_file_path, 'w')
+    f.attrs['version'] = 'HDF5 file made with TracTrac Python v' + version
+    f.attrs['date'] = time.strftime("%d/%m/%Y")
+    f.attrs['nFrames'] = trac_data.num_frames
+    f.attrs['size'] = trac_data.image_size
+    for items in trac_data.th.keys():
+        f.attrs['th:' + items] = trac_data.th[items]
+    f.create_dataset("Frame", data=np.uint16(trac_data.pts[:, 0]))
+    f.create_dataset("Id", data=np.uint16(trac_data.pts[:, 1]))
+    f.create_dataset("x", data=np.float32(trac_data.pts[:, 2:4]))
+    f.create_dataset("u", data=np.float32(trac_data.pts[:, 4:6]))
+    f.create_dataset("a", data=np.float32(trac_data.pts[:, 6:8]))
+    f.create_dataset("u_motion", data=np.float32(trac_data.pts[:, 8:10]))
+    f.create_dataset("err_motion", data=np.uint32(trac_data.pts[:, 10]))
+    f.close()
+
+
+def hdf5_to_trac_data(hdf5_file_path: Path):
+
+    trac_data = TracData()
+    f = h5py.File(hdf5_file_path, 'r')
+
+    version = f.attrs['version']  # noqa
+    date = f.attrs['date']  # noqa
+    trac_data.num_frames = f.attrs['nFrames']
+    trac_data.image_size = f.attrs['size']
+    print(type(trac_data.image_size))
+    # f.attrs['size'] = I0f.shape
+    # for items in th[0].keys():
+    #     f.attrs['th:' + items] = th[0][items]
+    num_pts = f['Frame'].shape[0]
+    trac_data.pts = np.ndarray(shape=(num_pts, 11), dtype=float)
+    trac_data.pts[:, 0] = f['Frame']
+    trac_data.pts[:, 1] = f['Id']
+    trac_data.pts[:, 2:4] = f['x']
+    trac_data.pts[:, 4:6] = f['u']
+    trac_data.pts[:, 6:8] = f['a']
+    trac_data.pts[:, 8:10] = f['u_motion']
+    trac_data.pts[:, 10] = f['err_motion']
+    f.close()
+    return trac_data
+
+
+def create_image(trac_data: TracData, frame_index: int, particle_radius: float):
+    image = np.zeros(shape=trac_data.image_size, dtype=float)
+    print(trac_data.pts)
+    frame_pts = trac_data.pts[trac_data.pts[:, 0] == frame_index]
+    for x in frame_pts[:, 2:4]:
+        center_coordinates = (int(x[0]), int(x[1]))
+        # center_coordinates = (10,10)
+        print(center_coordinates)
+        cv2.circle(image, center_coordinates, int(particle_radius), color=1.0, thickness=-1)
+    return image
+
+
 def main():
-    slope_image = create_slope_image(image_width=128, image_height=256, p1_angle=0.7, p1_radius=3.0)
-    cv2.imwrite(str(Path('toto.tif')), slope_image)
+    # python3 ./tractrac.git/Python/tractrac.py -f ./grassloper.git/samples/sample002.avi --output 1 -a --saveplot
+    trac_data = hdf5_to_trac_data('./grassloper.git/samples/TracTrac/sample002_track.hdf5')
+    kernel = np.ones((15, 15), np.uint8)
+    for frame_index in range(1, 2):
+        isbead_image = create_image(trac_data, frame_index=frame_index, particle_radius=11.0)
+        closing = cv2.morphologyEx(isbead_image, cv2.MORPH_CLOSE, kernel)
+        cv2.imwrite('isbead_%04d.tif' % frame_index, isbead_image)
+        cv2.imwrite('closing_%04d.tif' % frame_index, closing)
+
+    # slope_image = create_slope_image(image_width=128, image_height=256, p1_angle=0.7, p1_radius=3.0)
+    # cv2.imwrite(str(Path('toto.tif')), slope_image)
 
 
 if __name__ == '__main__':