import cv2

# Vortrainiertes YOLO-Modell laden
model = cv2.dnn_DetectionModel(
    "frozen_inference_graph.pb",
    "graph.pbtxt"
)



model.setInputSize(320, 320)
model.setInputScale(1.0 / 127.5)
model.setInputMean((127.5, 127.5, 127.5))
model.setInputSwapRB(True)

# Kamera oder Videostream
kamera = cv2.VideoCapture(0)

while True:
    ok, bild = kamera.read()
    if not ok:
        break

    klassen, sicherheit, boxen = model.detect(
        bild,
        confThreshold=0.5
    )

    if len(klassen) > 0:
        for klasse, score, box in zip(klassen.flatten(),
                                      sicherheit.flatten(),
                                      boxen):

            x, y, w, h = box

            cv2.rectangle(
                bild,
                (x, y),
                (x + w, y + h),
                (0, 255, 0),
                2
            )

            cv2.putText(
                bild,
                f"ID:{klasse} {score:.2f}",
                (x, y - 10),
                cv2.FONT_HERSHEY_SIMPLEX,
                0.5,
                (0, 255, 0),
                2
            )

    cv2.imshow("Objekterkennung", bild)

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

kamera.release()
cv2.destroyAllWindows()
print('Hello world!')