from machine import Pin
import time


INPUT_WIDTH = 480
INPUT_HEIGHT = 480


predictions_raw = [
    [196.875, 87.805, 71.25, 110.875, 0.8712],
    [319.453, 115.007, 64.219, 108.809, 0.8425],
    [406.641, 60.258, 43.594, 62.669, 0.6628],
    [89.062, 148.407, 35.625, 44.763, 0.5688],
]

predictions = []
for raw in predictions_raw:
    x_center, y_center, width, height, conf = raw
    x_center_norm = x_center / INPUT_WIDTH
    y_center_norm = y_center / INPUT_HEIGHT
    width_norm = width / INPUT_WIDTH
    height_norm = height / INPUT_HEIGHT
    predictions.append([x_center_norm, y_center_norm, width_norm, height_norm, conf])

threshold = 0.5

led = Pin(25, Pin.OUT)  

def xywh_to_xyxy(box):
    x, y, w, h = box
    return [
        x - w / 2,  # x1
        y - h / 2,  # y1
        x + w / 2,  # x2
        y + h / 2   # y2
    ]

def iou(box1, box2):
    # box = [x1, y1, x2, y2]
    x1 = max(box1[0], box2[0])
    y1 = max(box1[1], box2[1])
    x2 = min(box1[2], box2[2])
    y2 = min(box1[3], box2[3])

    inter_area = max(0, x2 - x1) * max(0, y2 - y1)
    area1 = (box1[2] - box1[0]) * (box1[3] - box1[1])
    area2 = (box2[2] - box2[0]) * (box2[3] - box2[1])
    union_area = area1 + area2 - inter_area
    if union_area == 0:
        return 0
    return inter_area / union_area

def nms(boxes, scores, iou_threshold=0.5):
    keep = []
    idxs = list(range(len(boxes)))
    idxs.sort(key=lambda i: scores[i], reverse=True)

    while idxs:
        current = idxs.pop(0)
        keep.append(current)
        idxs = [i for i in idxs if iou(boxes[current], boxes[i]) < iou_threshold]
    return keep

def main():

    filtered = [p for p in predictions if p[4] > threshold]
    if not filtered:
        print("No detections above threshold")
        led.value(0)
        return

    boxes = [xywh_to_xyxy(p[:4]) for p in filtered]
    scores = [p[4] for p in filtered]

    keep_indices = nms(boxes, scores)

    final_boxes = [boxes[i] for i in keep_indices]
    final_scores = [scores[i] for i in keep_indices]

    print("Detections after NMS:")
    for i, box in enumerate(final_boxes):
        print(f"Box {i}: {box}, Confidence: {final_scores[i]:.2f}")


    led.value(1)

if __name__ == "__main__":
    while True:
        main()
        time.sleep(3)