import argparse
import cv2
from models import ClassificationModel

cls_model = ClassificationModel()

def main(args):
    cap = cv2.VideoCapture(args.video_path)
    object_detector = cv2.createBackgroundSubtractorMOG2(history=100, varThreshold=50)
    while True:
        ret, frame = cap.read()
        if(ret):
            roi = frame[100: 900,330:1900]
            mask = object_detector.apply(roi)
            _, mask = cv2.threshold(mask,254,255, cv2.THRESH_BINARY)
            conturs, _ =cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

            for cnt in conturs:
                area = cv2.contourArea(cnt)
                if area > 200:
                    cv2.drawContours(roi,[cnt],-1,(0,255,0),2)
                    x,y,w,h = cv2.boundingRect(cnt)
                    rectangle = cv2.rectangle(roi,(x,y),(x+w,y+h),(0,255,0),3)
                    image_to_predict = roi[y:y+h+10,x:x+w+10]
                    label = cls_model.predict(image_to_predict)
                    cv2.putText(rectangle, label, (x, y-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (36,255,12), 1)
            
            roi = cv2.resize(roi, (960, 540)) 
            cv2.imshow("roi", roi)

            key = cv2.waitKey(30)
            if key == 27:
                break
        else:
            break
    cap.release()
    cv2.destroyAllWindows()

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--video_path", type=str, default='./test_videos/rybki2.mp4')
    args = parser.parse_args()
    main(args)