Cleaning up

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)