import cv2
import numpy as np
from tracker import *

# Create tracker object
tracker = EuclideanDistTracker()

cap = cv2.VideoCapture("Video Source")

count = 0

while True:
    ret, frame = cap.read()

    print("\n")
    
    if(count != 0):
        print("Frame Count: ", count)
    frame = cv2.resize(frame, (0, 0), fx = 1.5, fy = 1.5)
    height, width, channels = frame.shape
 
    # 1. Object Detection
    hsvFrame = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
    sensitivity = 20
    white_lower = np.array([0,0,255-sensitivity])
    white_upper = np.array([255,sensitivity,255])
    white_mask = cv2.inRange(hsvFrame, white_lower, white_upper)

    # Morphological Transform, Dilation
    kernal = np.ones((3, 3), "uint8")
    c = 1

    contours_w, hierarchy_w = cv2.findContours(white_mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
    detections = [] 

    for contour_w in contours_w:
        area = cv2.contourArea(contour_w)
        if area > 200 and area < 100000:
            cv2.drawContours(frame, [contour_w], -1, (0, 0, 0), 1)
            x,y,w,h = cv2.boundingRect(contour_w)
            detections.append([x, y, w, h])

    # 2. Object Tracking
    boxes_ids = tracker.update(detections)
    for box_id in boxes_ids:
        x, y, w, h, id = box_id
        cv2.putText(frame, str(id), (x - 8, y + 8), cv2.FONT_HERSHEY_PLAIN, 2, (0, 0, 0), 2)
        cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 0, 0), 1)

    cv2.imshow("Frame", frame) 
    count += 1
    key = cv2.waitKey(0)
    if key == 27:
        break

cap.release()
cv2.destroyAllWindows()