radio/radio.py

import cv2
import numpy as np
import math
import itertools as it

DENSITY = 4
RED = [0, 0, 255]
N_COLORS = 3


def draw_rectangle(frame, sp, ep):
    """ Draw a rectangle on the frame """
    return cv2.rectangle(frame, sp, ep, RED)


def analyze_block(frame, index, sp, ep):
    """ Analyze a block """
    block = frame[sp[1]:ep[1], sp[0]:ep[0], 0:3]
    average_color = np.average(block, (0, 1))
    sp2 = tuple(int(x) for x in (2 * np.array(sp) + np.array(ep)) / 3)
    ep2 = tuple(int(x) for x in (np.array(sp) + 2 * np.array(ep)) / 3)
    average_color = [int(x) for x in average_color]
    frame = cv2.rectangle(frame, sp2, ep2, average_color, 5)
    average_color = np.uint8([[average_color]])
    h, s, v = cv2.cvtColor(average_color, cv2.COLOR_BGR2HSV)[0][0]
    print(index, h, s, v)
    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 index, (x, y) in enumerate(it.product(range(n), range(m))):
        sp = (int(x * dx), int(y * dy))
        ep = (int(x * dx + dx), int(y * dy + dy))
        frame = draw_rectangle(frame, sp, ep)
        frame = analyze_block(frame, index, sp, ep)
    return frame


if __name__ == '__main__':
    camera = cv2.VideoCapture(0)

    while True:
        ret, frame = camera.read()
        frame = analyze(frame)
        cv2.imshow('Input', frame)

        c = cv2.waitKey(1)
        if c == 27:
            break

    camera.release()
    cv2.destroyAllWindows()