Introduction to the ESP32-S31 with Vision: Assignment 5

from machine import Pin, bitstream
import display
import espdl
import sensor
import time


MODEL = "/espdet_pico_224_224_hand.espdl"
LED_PIN = 37
SCORE = 0.5
NMS = 0.7

ZONE_X = 0.60
ZONE_Y = 0.15
ZONE_WIDTH = 0.30
ZONE_HEIGHT = 0.70
MIN_OVERLAP = 0.15

MAX_DISTANCE_SQ = 80 * 80
MAX_MISSED = 3
ENTER_FRAMES = 2
EXIT_FRAMES = 5

GREEN = (0, 24, 0)
RED = (32, 0, 0)
OFF = (0, 0, 0)
WS2812_TIMING = (400, 850, 800, 450)


def set_led(pin, color):
    red, green, blue = color
    bitstream(pin, 0, WS2812_TIMING, bytes((green, red, blue)))
    time.sleep_us(80)


def make_zone(frame):
    return (
        int(frame.width() * ZONE_X),
        int(frame.height() * ZONE_Y),
        int(frame.width() * ZONE_WIDTH),
        int(frame.height() * ZONE_HEIGHT),
    )


def center(detection):
    x, y, w, h, score, category = detection
    return x + w // 2, y + h // 2


def update_tracks(detections, tracks, next_id):
    updated = {}
    used_ids = set()
    output = []

    for detection in detections:
        cx, cy = center(detection)
        best_id = None
        best_distance = MAX_DISTANCE_SQ + 1

        for track_id, state in tracks.items():
            if track_id in used_ids:
                continue
            px, py, missed = state
            distance = (cx - px) ** 2 + (cy - py) ** 2
            if distance < best_distance:
                best_id = track_id
                best_distance = distance

        if best_id is None or best_distance > MAX_DISTANCE_SQ:
            best_id = next_id
            next_id += 1

        used_ids.add(best_id)
        updated[best_id] = (cx, cy, 0)
        output.append((best_id, detection))

    for track_id, state in tracks.items():
        if track_id in used_ids:
            continue
        px, py, missed = state
        if missed < MAX_MISSED:
            updated[track_id] = (px, py, missed + 1)

    return output, updated, next_id


def overlap_ratio(hand, zone):
    hx, hy, hw, hh, score, category = hand
    zx, zy, zw, zh = zone
    left = max(hx, zx)
    top = max(hy, zy)
    right = min(hx + hw, zx + zw)
    bottom = min(hy + hh, zy + zh)
    intersection = max(0, right - left) * max(0, bottom - top)
    return intersection / max(1, hw * hh)


def update_warning(raw_intrusion, warning, enter_count, exit_count):
    if raw_intrusion:
        enter_count += 1
        exit_count = 0
        if enter_count >= ENTER_FRAMES:
            warning = True
    else:
        enter_count = 0
        exit_count += 1
        if exit_count >= EXIT_FRAMES:
            warning = False
    return warning, enter_count, exit_count

def main():
    led = Pin(LED_PIN, Pin.OUT)
    detector = None
    lcd = None
    tracks = {}
    next_id = 1
    warning = False
    enter_count = 0
    exit_count = EXIT_FRAMES
    last_led_state = None

    sensor.reset()
    try:
        lcd = display.Display()
        sensor.set_pixformat(sensor.RGB565)
        sensor.set_framesize(sensor.QVGA)
        sensor.skip_frames(time=1000)
        detector = espdl.ESPDet(MODEL, score=SCORE, nms=NMS)
        print("HANDS_OFF:READY")

        while True:
            frame = sensor.snapshot()
            zone = make_zone(frame)
            detections = detector.detect(frame)
            tracked, tracks, next_id = update_tracks(
                detections, tracks, next_id)

            raw_intrusion = False
            states = []
            for track_id, detection in tracked:
                ratio = overlap_ratio(detection, zone)
                inside = ratio >= MIN_OVERLAP
                raw_intrusion = raw_intrusion or inside
                states.append((track_id, detection, ratio, inside))

            warning, enter_count, exit_count = update_warning(
                raw_intrusion, warning, enter_count, exit_count)

            for track_id, detection, ratio, inside in states:
                x, y, w, h, score, category = detection
                color = (255, 0, 0) if inside else (0, 255, 0)
                frame.draw_rectangle(
                    x, y, w, h, color=color, thickness=2)
                frame.draw_string(
                    x, max(0, y - 12),
                    "hand %d %.0f%%" % (track_id, ratio * 100),
                    color=color)

            zx, zy, zw, zh = zone
            zone_color = (255, 0, 0) if warning else (0, 255, 0)
            frame.draw_rectangle(
                zx, zy, zw, zh, color=zone_color, thickness=3)
            frame.draw_string(
                8, 8,
                "KEEP HANDS OFF!" if warning else "ZONE CLEAR",
                color=zone_color)

            if warning != last_led_state:
                set_led(led, RED if warning else GREEN)
                print("HANDS_OFF:WARNING" if warning
                      else "HANDS_OFF:CLEAR")
                last_led_state = warning

            lcd.write(frame)
            time.sleep_ms(20)
    finally:
        set_led(led, OFF)
        if detector is not None:
            detector.deinit()
        sensor.shutdown()
        if lcd is not None:
            lcd.deinit()


main()

HANDS_OFF:READY
HANDS_OFF:CLEAR