UNIHIKER: Flame detection

from datetime import datetime
from sys import exit
from threading import Thread
from time import sleep

from numpy import ndarray
from pinpong.board import Board, Pin, NeoPixel
import cv2 as cv
import numpy as np


DISPLAY_WIDTH: int = 240
DISPLAY_HEIGHT: int = 320

DURATION: float = .5
LED_BLACK: tuple = (0, 0, 0)
LED_RED: tuple = (100, 0, 0)

LOWER_YELLOW: ndarray = np.array([18, 100, 100])
UPPER_YELLOW: ndarray = np.array([35, 255, 255])
LOWER_ORANGE: ndarray = np.array([5, 100, 100])
UPPER_ORANGE: ndarray = np.array([18, 255, 255])
LOWER_RED_1: ndarray = np.array([0, 100, 100])
UPPER_RED_1: ndarray = np.array([5, 255, 255])
LOWER_RED_2: ndarray = np.array([170, 100, 100])
UPPER_RED_2: ndarray = np.array([180, 255, 255])


def is_flame_detected(image: np.ndarray, sensitivity: int) -> bool:
    """
    Detects if a flame is present in the given image.

    :param image: The image to detect the flame in.
    :type image: np.ndarray
    :param sensitivity: The sensitivity level for detecting the flame (smaller values will detect smaller flames).
    :type sensitivity: int
    :return: True if a flame is detected, False otherwise.
    :rtype: bool
    """
    hsv = cv.cvtColor(image, cv.COLOR_BGR2HSV)

    mask_yellow = cv.inRange(hsv, LOWER_YELLOW, UPPER_YELLOW)
    mask_orange = cv.inRange(hsv, LOWER_ORANGE, UPPER_ORANGE)
    mask_red1 = cv.inRange(hsv, LOWER_RED_1, UPPER_RED_1)
    mask_red2 = cv.inRange(hsv, LOWER_RED_2, UPPER_RED_2)

    mask = mask_yellow | mask_orange | mask_red1 | mask_red2

    kernel = np.ones((5, 5), np.uint8)
    mask = cv.morphologyEx(mask, cv.MORPH_OPEN, kernel)
    mask = cv.morphologyEx(mask, cv.MORPH_CLOSE, kernel)

    contours, _ = cv.findContours(mask, cv.RETR_EXTERNAL, cv.CHAIN_APPROX_SIMPLE)

    for contour in contours:
        area = cv.contourArea(contour)
        if area > sensitivity:
            x, y, w, h = cv.boundingRect(contour)
            aspect_ratio = float(w) / h
            if 0.2 < aspect_ratio < 1.0:
                hull = cv.convexHull(contour)
                hull_area = cv.contourArea(hull)
                solidity = float(area) / hull_area
                if solidity > 0.5:
                    return True

    return False


def on_trackbar(val: int) -> None:
    """
    Callback function for the trackbar to update the sensitivity threshold.

    :param val: The new value of the trackbar.
    :type val: int
    :return: None
    """
    global sensitivity_threshold

    sensitivity_threshold = val


def set_led_color(duration: float, neopixel: NeoPixel) -> None:
    """
    Sets the color of the LEDs in a NeoPixel strip for a specified duration.

    :param duration: The duration (in seconds) for which the LEDs will display the color.
    :type duration: float
    :param neopixel: The NeoPixel strip object.
    :type neopixel: NeoPixel
    :return: None
    """
    for led in range(3):
        neopixel[led] = LED_RED

    sleep(duration)

    for led in range(3):
        neopixel[led] = LED_BLACK


def run_led_thread(duration: float, neopixel: NeoPixel) -> None:
    """
    Run LED Thread

    :param duration: The duration of the LED thread in seconds.
    :type duration: float
    :param neopixel: The neopixel object.
    :type neopixel: NeoPixel
    :return: None
    """
    led_thread = Thread(target=set_led_color, args=(duration, neopixel))
    led_thread.start()


if __name__ == '__main__':
    Board("UNIHIKER").begin()

    np_led = NeoPixel(pin_obj=Pin(Pin.P13), num=3)

    sensitivity_threshold = 2500

    cap = cv.VideoCapture(0)
    if not cap.isOpened():
        print('[ERROR] Could not open camera')
        exit(1)

    cap.set(cv.CAP_PROP_FRAME_WIDTH, DISPLAY_WIDTH)
    cap.set(cv.CAP_PROP_FRAME_HEIGHT, DISPLAY_HEIGHT)
    cap.set(cv.CAP_PROP_BUFFERSIZE, 1)

    cv.namedWindow('display_window', cv.WND_PROP_FULLSCREEN)
    cv.setWindowProperty('display_window', cv.WND_PROP_FULLSCREEN, cv.WINDOW_FULLSCREEN)

    cv.createTrackbar('Sensitivity', 'display_window', sensitivity_threshold, 5000, on_trackbar)

    while True:
        ret, frame = cap.read()

        if not ret:
            break

        if cv.waitKey(1) & 0xff == ord('q'):
            break

        flipped_frame = cv.flip(frame, 1)
        flipped_frame.flags.writeable = False

        flame_detected = is_flame_detected(flipped_frame, sensitivity_threshold)

        if flame_detected:
            now = datetime.now()
            current_time = now.strftime("%Y-%m-%d %H:%M:%S")
            print(f"Flame detected at {current_time}")
            run_led_thread(duration=DURATION, neopixel=np_led)

        cv.imshow('display_window', flipped_frame)

    cap.release()
    cv.destroyAllWindows()