added a simple slope finder

it's functional but needs cleaning and improving
This commit is contained in:
Guillaume Raffy 2022-11-30 16:09:04 +01:00
parent 4d3de3c6f9
commit 71fdb927f7
2 changed files with 70 additions and 15 deletions

View File

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

83
main.py
View File

@ -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)