
import time
from imutils.video.pivideostream import PiVideoStream
import tensorflow as tf
import numpy as np
import cv2
import RPi.GPIO as GPIO
import requests as req

GPIO.setwarnings(False)
GPIO.setmode(GPIO.BOARD)
GPIO.setup(10, GPIO.IN, pull_up_down=GPIO.PUD_DOWN)

interpreter = tf.lite.Interpreter(model_path="model_unquant.tflite")
interpreter.allocate_tensors()

input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

target_height = input_details[0]["shape"][1]
target_width = input_details[0]["shape"][2]

f = open("labels.txt", "r")
lines = f.readlines()
f.close()
classes = {}
for line in lines:
    pair = line.strip().split(maxsplit=1)
    classes[int(pair[0])] = pair[1].strip()


def detect(frame):
    # Prepare input data
    resized = cv2.resize(frame, (target_width, target_height))
    input_data = np.expand_dims(resized, axis=0)
    input_data = (np.float32(input_data) - 127.5) / 127.5
    interpreter.set_tensor(input_details[0]["index"], input_data)

    interpreter.invoke()
    detection = interpreter.get_tensor(output_details[0]["index"])
    return detection


def draw_detection(frame, detection):
    for i, s in enumerate(detection[0]):
        tag = f"{classes[i]}: {s*100:.2f}%"
        cv2.putText(frame, tag, (10, 20 + 20 * i),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1)
    return frame


def main():
    camera = PiVideoStream(resolution=(512, 400)).start()
    time.sleep(2)

    while True:
        frame = camera.read()
        detection = detect(frame)
        value = classes[detection.tolist()[0].index(
            max(detection.tolist()[0]))]
        drawn = draw_detection(frame, detection)
        cv2.imshow("frame", drawn)
        if GPIO.input(10) == GPIO.HIGH:
            request = 'ESP32_IP_Address' + '/' + value
            response = req.get(request)
        if cv2.waitKey(1) & 0xFF == ord("q"):
            break

    camera.stop()
    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()
