46 lines
1.6 KiB
Python
46 lines
1.6 KiB
Python
import numpy as np
|
|
import matplotlib.pyplot as plt
|
|
from skimage import io
|
|
from skimage.filters import threshold_otsu, gaussian
|
|
from skimage.morphology import closing, square
|
|
from skimage.measure import label, regionprops
|
|
from skimage.color import rgb2gray
|
|
import matplotlib.patches as patches
|
|
from sklearn.cluster import KMeans
|
|
|
|
def get_heights(filename: str):
|
|
image = io.imread(filename)
|
|
im = rgb2gray(image)
|
|
binary_im = im > threshold_otsu(im)
|
|
smooth_im = gaussian(binary_im.astype(float), sigma=1)
|
|
morph_im = closing(smooth_im > 0.5, square(3))
|
|
|
|
labeled = label(morph_im)
|
|
bar_dims = [r.bbox for r in regionprops(labeled)]
|
|
bar_dims.sort(key=lambda x: x[1]) # left to right
|
|
|
|
bars = bar_dims[1:] # skip logo
|
|
bar_heights_raw = []
|
|
print(len(bars))
|
|
for bar in bars:
|
|
top, left, bottom, right = bar
|
|
effective_height = bottom - top # use bounding box height directly
|
|
bar_heights_raw.append(effective_height)
|
|
|
|
# Cluster measured heights to 8 clusters representing discrete bar levels
|
|
bar_heights_raw_np = np.array(bar_heights_raw).reshape(-1, 1)
|
|
|
|
kmeans = KMeans(n_clusters=8, random_state=0).fit(bar_heights_raw_np)
|
|
cluster_centers = np.sort(kmeans.cluster_centers_.flatten())
|
|
|
|
# Assign each bar height to closest cluster center index (0 to 7)
|
|
predicted_levels = []
|
|
for h in bar_heights_raw:
|
|
diffs = np.abs(cluster_centers - h)
|
|
closest_cluster = np.argmin(diffs)
|
|
predicted_levels.append(int(closest_cluster))
|
|
|
|
|
|
print((predicted_levels))
|
|
return predicted_levels
|