added a simple slope finder

it's functional but needs cleaning and improving

README.md

@@ -1,2 +1,4 @@
 # grassloper
 an application to estimate the slope of a granular surface flow
+
+

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,7 +86,62 @@ def hdf5_to_trac_data(hdf5_file_path: Path):
     return trac_data
 
 
-def create_image(trac_data: TracData, frame_index: int, particle_radius: float):
+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]
@@ -100,12 +156,9 @@ def create_image(trac_data: TracData, frame_index: int, particle_radius: float):
 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)