diff --git a/main.py b/main.py index f633012..d1a232e 100755 --- a/main.py +++ b/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__':