diff --git a/boonepeter.github.io-code/spotify-codes-part-2/src/get_heights.py b/boonepeter.github.io-code/spotify-codes-part-2/src/get_heights.py index c5fa784..8070bd1 100644 --- a/boonepeter.github.io-code/spotify-codes-part-2/src/get_heights.py +++ b/boonepeter.github.io-code/spotify-codes-part-2/src/get_heights.py @@ -1 +1,78 @@ -< \ No newline at end of file +from skimage import io +from skimage.color import rgb2gray +from skimage.filters import threshold_otsu +from skimage.measure import label, regionprops +import matplotlib.pyplot as plt +import matplotlib.patches as patches +from skimage import io +from skimage.color import rgb2gray +from skimage.filters import threshold_otsu, gaussian +from skimage.measure import label, regionprops +import matplotlib.pyplot as plt +import matplotlib.patches as patches +def get_heights(filename: str) -> list: + """Open an image and return a list of the bar heights. + Also saves three images: the processed binary image, the original image with drawings, and a smoothed binary image. + """ + image = io.imread(filename) + im = rgb2gray(image) + binary_im = im > threshold_otsu(im) + + # Apply Gaussian smoothing to binary image to make bars look rounder + smooth_im = gaussian(binary_im, sigma=2) + # Threshold again to obtain a smoothed binary + smooth_binary = smooth_im > 0.5 + labeled = label(smooth_binary) + bar_dimensions = [r.bbox for r in regionprops(labeled)] + bar_dimensions.sort(key=lambda x: x[1]) # Sort left to right + # The first object (spotify logo) is the max height of the bars + logo = bar_dimensions[0] + max_height = logo[2] - logo[0] + sequence = [] + # Create figure and axis for drawing + fig, ax = plt.subplots() + ax.imshow(smooth_binary, cmap='gray') + # Draw rectangle around logo + logo_rect = patches.Rectangle( + (logo[1], logo[0]), + logo[3]-logo[1], + logo[2]-logo[0], + linewidth=2, + edgecolor='yellow', + facecolor='none' + ) + ax.add_patch(logo_rect) + # Add 'Logo' text near the rectangle + ax.text(logo[1], logo[0] - 10, 'Logo', color='yellow', fontsize=12, weight='bold') + # Draw bars and center markers + for bar in bar_dimensions[1:]: + height = bar[2] - bar[0] + ratio = height / max_height + ratio *= 8 + ratio //= 1 + val = int(ratio - 1) + sequence.append(val) + # Draw rectangle around bar + rect = patches.Rectangle( + (bar[1], bar[0]), + bar[3]-bar[1], + bar[2]-bar[0], + linewidth=1, + edgecolor='red', + facecolor='none' + ) + ax.add_patch(rect) + # Draw center marker (white circle) + center_x = (bar[1] + bar[3]) / 2 + center_y = (bar[0] + bar[2]) / 2 + center_marker = patches.Circle((center_x, center_y), radius=5, color='black') + ax.add_patch(center_marker) + # Save processed binary image used for calculations + plt.imsave('processed_binary_image.png', binary_im, cmap='gray') + # Save smoothed binary image + plt.imsave('img_smooth.png', smooth_binary, cmap='gray') + # Save image with drawings + plt.axis('off') + plt.savefig('image_with_drawings.png', bbox_inches='tight', pad_inches=0) + plt.close() + return sequence \ No newline at end of file