import cv2
from yolov5processor.infer import ExecuteInference
import glob
import numpy as np
#from google.colab.patches import cv2_imshow

yp = ExecuteInference(weight="E:\buds\best.pt", confidence=0.75, iou= 0.45, agnostic_nms= False, img_size= 416)
#yp = ExecuteInference(weight=r"E:\buds\best.pt")
vid = cv2.VideoCapture("E:\buds\original video\predict.mp4")
frame_width = int(vid.get(3))
frame_height = int(vid.get(4))


#count_line_position=550
detect = []
offset=6
while (True):
    faulty_count = 0
    good_count = 0
    rejected_count = 0
    ret, image = vid.read()

    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=3,
                        fontScale=1, 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=3,
                        fontScale=1, fontFace=cv2.LINE_AA)
                good_count += 1
                    
           else :
                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=3,
                            fontScale=1, fontFace=cv2.LINE_AA)
                rejected_count += 1
                
        print(f"good-{good_count}, faulty-{faulty_count}, rejected-{rejected_count}")
    
        
    
            
        print("good:",good_count)
        print("faulty:",faulty_count)
        print("rejected:",rejected_count)
      
    result = f"Good: {good_count}, Faulty: {faulty_count}, Rejected: {rejected_count}"
    
    cv2.putText(image,result,(450,70),cv2.FONT_HERSHEY_SIMPLEX,2,(0,0,255),5)
           
            
    

    cv2.imshow('frame', cv2.resize(image, (720, 560)))
    #cv2_imshow(cv2.resize(image, (720, 560)))
    if cv2.waitKey(1) & 0xFF == ord('q'):
      break
    

# After the loop release the cap object
vid.release()
# Destroy all the windowsss