diff --git a/README.md b/README.md index 4b3116d..18bd7db 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,4 @@ # grassloper an application to estimate the slope of a granular surface flow + + diff --git a/main.py b/main.py index d1a232e..43b877e 100755 --- a/main.py +++ b/main.py @@ -5,6 +5,7 @@ from pathlib import Path import math import h5py import time +import scipy.stats 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): @@ -85,27 +86,79 @@ def hdf5_to_trac_data(hdf5_file_path: Path): 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 +class SlopeFinder(): + + def __init__(self, beads_radius: float): + self.beads_radius = beads_radius + + def find_slope(self, trac_data: TracData, frame_index: int): + isbead_image = SlopeFinder.create_isbead_image(trac_data, frame_index, self.beads_radius) + cv2.imwrite('isbead_%04d.tif' % frame_index, isbead_image) + if False: + kernel = np.ones((15, 15), np.uint8) + closing = cv2.morphologyEx(isbead_image, cv2.MORPH_CLOSE, kernel) + cv2.imwrite('closing_%04d.tif' % frame_index, closing) + # remove the jumping beads + # apply connected component analysis to the thresholded image + # thresh = cv2.threshold(isbead_image, 0, 255, cv2.THRESH_BINARY_INV | cv2.THRESH_OTSU)[1] + thresh = cv2.compare(isbead_image, 0.5, cv2.CMP_GT) + connectivity = 4 + output = cv2.connectedComponentsWithStats(thresh, connectivity, cv2.CV_32S) + (num_labels, labels_image, stats, centroids) = output + print('num_labels: ', num_labels) + print('labels_image: ', labels_image) + print('stats: ', stats) + print('centroids: ', centroids) + cv2.imwrite('labels_%04d.tif' % frame_index, labels_image) + is_non_jumping_bead_image = np.zeros(shape=labels_image.shape, dtype=np.uint8) + bead_area = math.pi * self.beads_radius * self.beads_radius + for label in range(1, num_labels): # ignore the first label, as it's the background + area = stats[label][4] + if area > bead_area * 1.1: + is_non_jumping_bead_image[labels_image == label] = 255 + cv2.imwrite('non_jumping_beads_%04d.tif' % frame_index, is_non_jumping_bead_image) + # extract the surface points + surface_y = np.ndarray(shape=(is_non_jumping_bead_image.shape[1],), dtype=int) + surface_y.fill(is_non_jumping_bead_image.shape[0]) + contours, hierarchy = cv2.findContours(is_non_jumping_bead_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE) + print('num contours: ', len(contours)) + assert(len(contours) == 1) + contour = contours[0] + for point in contour: + x, y = point[0] + surface_y[x] = min(surface_y[x], y) + print(surface_y) + surface_pts = [] + for x in range(len(surface_y)): + y = surface_y[x] + if y != is_non_jumping_bead_image.shape[0]: + surface_pts.append((x, y)) + print(surface_pts) + x = [pt[0] for pt in surface_pts] + y = [pt[1] for pt in surface_pts] + lin_regress_result = scipy.stats.linregress(x, y) + print(lin_regress_result) + + + @staticmethod + def create_isbead_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(): # 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) + slope_finder = SlopeFinder(beads_radius=11.0) 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_finder.find_slope(trac_data, frame_index=frame_index) # 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)