from ultralytics import YOLO
"""
Created on Wed Oct  4 16:44:12 2023

@author: lin
"""
import glob
import sys
sys.path.append('../../..')
import os
import cv2
import json
import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
# from utils.bbox_op import non_max_supression

def plot_prediction(image_np, bboxes, classes, scores, label_map):
    color=(255,0,0)
    thickness=5
    font_scale=3

    for i, box in enumerate(bboxes):
        box = bboxes[i, :]
        
        x0, y0, x1, y1 = box

        xmin = int(x0)
        ymin = int(y0)
        xmax = int(x1)
        ymax = int(y1)

        print(xmin, ymin, xmax, ymax)
        
        image_np = cv2.rectangle(image_np, (xmin, ymin), (xmax, ymax), color=color, thickness=thickness)
        text_x = xmin - 10 if xmin > 20 else xmin + 10
        text_y = ymin - 10 if ymin > 20 else ymin + 10
        display_str = label_map[str(int(classes))]

        cv2.putText(
            image_np,
            display_str,
            (text_x, text_y),
            cv2.FONT_HERSHEY_SIMPLEX,
            font_scale,
            color,
            thickness,
        )
    plt.imshow(image_np)
    plt.show()

if __name__ == "__main__":
    min_th = 0.1
    labels_json = "coco_labels.json"
    with open(labels_json) as f:
        label_map = json.load(f)
    img_path = "test_images"
    saved_tflite = "tflite_model.tflite"
    # load model
    
    model = YOLO("yolov8n.pt")
    images = glob.glob(os.path.join(img_path, "*"))
    for img in images:
        
        image_np = cv2.imread(img) # gpreprocess(img)
        image_np = cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB)
        # print(image_np.shape)
    
        # image_np = np.array(Image.open(image_path))
        results = model(img)
        boxes = results[0].boxes

        print(boxes.xyxy)
        print(boxes.cls)
        print(boxes.conf)
        plot_prediction(image_np, boxes.xyxy.numpy(), boxes.cls.numpy(), boxes.conf.numpy(), label_map)