diff --git a/radio.py b/radio.py index b957b91..0abf838 100644 --- a/radio.py +++ b/radio.py @@ -1,22 +1,43 @@ import cv2 +import numpy as np import math import itertools as it DENSITY = 5 RED = [0, 0, 255] +N_COLORS = 3 -def draw_rectangles(frame): - """ Draw some rectangles on top of the image """ +def draw_rectangle(frame, sp, ep): + """ Draw a rectangle on the frame """ + return cv2.rectangle(frame, sp, ep, RED) + + +def analyze_block(frame, sp, ep): + """ Analyze a block """ + block = np.float32(frame[sp[1]:ep[1], sp[0]:ep[0], 0:3]) + block = block.reshape((-1, 3)) + criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, .2) + compactness, labels, palette = cv2.kmeans(block, N_COLORS, None, criteria, + 2, cv2.KMEANS_RANDOM_CENTERS) + unique, counts = np.unique(labels.flatten(), return_counts=True) + for index, count in zip(unique, counts): + print(index, count, palette[index]) + return frame + + +def analyze(frame): + """ Analyze the full frame """ height, width, d = frame.shape n = DENSITY m = math.ceil(n * (height / width)) dx = width / n dy = height / m - for x, y in it.product(range(n), range(n)): + for x, y in it.product(range(n), range(m)): sp = (int(x * dx), int(y * dy)) ep = (int(x * dx + dx), int(y * dy + dy)) - frame = cv2.rectangle(frame, sp, ep, (255, 0, 0)) + frame = draw_rectangle(frame, sp, ep) + frame = analyze_block(frame, sp, ep) return frame @@ -25,7 +46,7 @@ if __name__ == '__main__': while True: ret, frame = camera.read() - frame = draw_rectangles(frame) + frame = analyze(frame) cv2.imshow('Input', frame) c = cv2.waitKey(1)