
from yolov5processor.infer import ExecuteInference
import glob
import numpy as np

yp = ExecuteInference(weight="E:\buds\best.pt")


faulty_count = 0
good_count = 0
rejected_count = 0
image = cv2.imread("E:\buds\original-predicted images\IMG_20220305_223916_jpg.rf.522aaff60984faaa45ddf1b14bc5bb9f.jpg")

pred = yp.predict(image)
print(pred)
for det in pred:
    conf = "{:.2f}".format(int(det['conf'] * 100))
    if int(det['conf'] * 100) > 50:
        # print(det['class'])
        label = det['class']  # print(label.item())
        det = det['points']
        if label == "faulty":
            cv2.rectangle(image, (int(det[0]), int(det[1])), (int(det[2]), int(det[3])), (0, 0, 255), 2)

            cv2.putText(img=image, text=conf + label, org=(int(det[0]), int(det[1]) - 10), color=(0, 0, 225),
                        thickness=1,
                        fontScale=0.5, fontFace=cv2.LINE_AA)
            faulty_count += 1
        elif label =="good" :
            cv2.rectangle(image, (int(det[0]), int(det[1])), (int(det[2]), int(det[3])), (0, 0, 255), 2)

            cv2.putText(img=image, text=conf + label, org=(int(det[0]), int(det[1]) - 10), color=(0, 0, 225),
                        thickness=1,
                        fontScale=0.5, fontFace=cv2.LINE_AA)
            good_count += 1
        elif label =="rejected" :
            cv2.rectangle(image, (int(det[0]), int(det[1])), (int(det[2]), int(det[3])), (0, 0, 255), 2)

            cv2.putText(img=image, text=conf + label, org=(int(det[0]), int(det[1]) - 10), color=(0, 0, 225),
                        thickness=1,
                        fontScale=0.5, fontFace=cv2.LINE_AA)
            rejected_count += 1
print(f"good-{good_count}, faulty-{faulty_count}, rejected-{rejected_count}")
#from google.colab.patches import cv2_imshow

cv2.imshow("image",image)