proj_2_cv.py

# -*- coding: utf-8 -*-
"""Proj#2-CV.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1XpHDlyt2K_uFCebW-0E59VSxaSvzcNf3

# **Imports**
"""

#1
#Check CPU and RAM specifications
!cat /proc/cpuinfo
!cat /proc/meminfo

#2
!pip list

#3
import numpy as np 
import time
from google.colab import files, auth, drive
import cv2 as cv
from google.colab.patches import cv2_imshow

#4
# Mount
drive.mount('/content/gdrive', force_remount=True)
data_dir_drive ='/content/gdrive/My Drive/Colab Notebooks/CV/Proj2' #change to the project folder on the drive
data_dir ='/content' # When read from unziped file

#5
# Unzip dataset to /content
start = time.time()

!unzip -q '/content/gdrive/My Drive/Colab Notebooks/CV/Proj2/Project2.zip' -d '/content/gdrive/My Drive/Colab Notebooks/CV/Proj2/data'

print('Took', (time.time() - start), ' secundes to unzip')

#6
!ls '/content/gdrive/My Drive/Colab Notebooks/CV/Proj2/data'

#7
folder_in = data_dir_drive + '/data/training_data/'
folder_out = data_dir_drive + '/output_data/lupascu_marian_407/'

"""# **Task 1**
You receive a test set of 57 images of a snooker table seen from above at a random time of a game (for example the black ball might be potted). The task is to count the number of balls on the table and specify their color for each of the 57 images in the test set.
"""

#1
def get_bottom_right(img, show_results = True):
    w, h = img.shape[::-1]

    methods = ['cv.TM_CCOEFF', 'cv.TM_CCOEFF_NORMED', 'cv.TM_CCORR',
            'cv.TM_CCORR_NORMED', 'cv.TM_SQDIFF', 'cv.TM_SQDIFF_NORMED']

    meth = methods[1]

    template1 = cv.imread(data_dir_drive + '/data/templates/tem1.jpg', 0)
    template2 = cv.imread(data_dir_drive + '/data/templates/tem11.jpg', 0)
    w1, h1 = template1.shape[::-1]
    w2, h2 = template2.shape[::-1]

    method = eval(meth)

    # Apply template Matching
    res1 = cv.matchTemplate(img[h//2:, w//2:],template1,method)
    res2 = cv.matchTemplate(img[h//2:, w//2:],template2,method)
    min_val1, max_val1, min_loc1, max_loc1 = cv.minMaxLoc(res1)
    min_val2, max_val2, min_loc2, max_loc2 = cv.minMaxLoc(res2)

    # If the method is TM_SQDIFF or TM_SQDIFF_NORMED, take minimum
    if method in [cv.TM_SQDIFF, cv.TM_SQDIFF_NORMED]:
        top_left1 = min_loc1
        top_left2 = min_loc2
    else:
        top_left1 = max_loc1
        top_left2 = max_loc2

    #print(res1[top_left1[1], top_left1[0]])
    #print(res2[top_left2[1], top_left2[0]])

    if res1[top_left1[1], top_left1[0]] > 0.6 or res2[top_left2[1], top_left2[0]] > 0.6:
        if res1[top_left1[1], top_left1[0]] > res2[top_left2[1], top_left2[0]]:
            top_left1 = (top_left1[0] + w//2, top_left1[1] + h//2)
            bottom_right1 = (top_left1[0] + w1, top_left1[1] + h1)
            if show_results:         
                cv.rectangle(img,top_left1, bottom_right1, 255, 2)
                cv2_imshow(cv.resize(img, (0, 0), fx=0.5, fy=0.5))
            return top_left1, bottom_right1, ((top_left1[0] + bottom_right1[0]) // 2, (top_left1[1] + bottom_right1[1]) // 2), img
        else:
            top_left2 = (top_left2[0] + w//2, top_left2[1] + h//2)
            bottom_right2 = (top_left2[0] + w2, top_left2[1] + h2)
            if show_results:         
                cv.rectangle(img,top_left2, bottom_right2, 255, 2)
                cv2_imshow(cv.resize(img, (0, 0), fx=0.5, fy=0.5))
            return top_left2, bottom_right2, ((top_left2[0] + bottom_right2[0]) // 2, (top_left2[1] + bottom_right2[1]) // 2), img
    else:
        return None, None, None, img

#2
def get_bottom_left(img, show_results = True):
    w, h = img.shape[::-1]

    methods = ['cv.TM_CCOEFF', 'cv.TM_CCOEFF_NORMED', 'cv.TM_CCORR',
            'cv.TM_CCORR_NORMED', 'cv.TM_SQDIFF', 'cv.TM_SQDIFF_NORMED']

    meth = methods[1]

    template1 = cv.imread(data_dir_drive + '/data/templates/tem2.jpg', 0)
    template2 = cv.imread(data_dir_drive + '/data/templates/tem22.jpg', 0)
    w1, h1 = template1.shape[::-1]
    w2, h2 = template2.shape[::-1]

    method = eval(meth)

    # Apply template Matching
    res1 = cv.matchTemplate(img[h//2:, :w//2],template1,method)
    res2 = cv.matchTemplate(img[h//2:, :w//2],template2,method)
    min_val1, max_val1, min_loc1, max_loc1 = cv.minMaxLoc(res1)
    min_val2, max_val2, min_loc2, max_loc2 = cv.minMaxLoc(res2)

    # If the method is TM_SQDIFF or TM_SQDIFF_NORMED, take minimum
    if method in [cv.TM_SQDIFF, cv.TM_SQDIFF_NORMED]:
        top_left1 = min_loc1
        top_left2 = min_loc2
    else:
        top_left1 = max_loc1
        top_left2 = max_loc2

    #print(res1[top_left1[1], top_left1[0]])
    #print(res2[top_left2[1], top_left2[0]])

    if res1[top_left1[1], top_left1[0]] > 0.6 or res2[top_left2[1], top_left2[0]] > 0.6:
        if res1[top_left1[1], top_left1[0]] > res2[top_left2[1], top_left2[0]]:
            top_left1 = (top_left1[0], top_left1[1] + h//2)
            bottom_right1 = (top_left1[0] + w1, top_left1[1] + h1)
            if show_results:  
                cv.rectangle(img,top_left1, bottom_right1, 255, 2)       
                cv2_imshow(cv.resize(img, (0, 0), fx=0.5, fy=0.5))
            return top_left1, bottom_right1, ((top_left1[0] + bottom_right1[0]) // 2, (top_left1[1] + bottom_right1[1]) // 2), img
        else:
            top_left2 = (top_left2[0], top_left2[1] + h//2)
            bottom_right2 = (top_left2[0] + w2, top_left2[1] + h2)
            if show_results:         
                cv.rectangle(img,top_left2, bottom_right2, 255, 2)
                cv2_imshow(cv.resize(img, (0, 0), fx=0.5, fy=0.5))
            return top_left2, bottom_right2, ((top_left2[0] + bottom_right2[0]) // 2, (top_left2[1] + bottom_right2[1]) // 2), img
    else:
        return None, None, None, img

#3
def get_top_left(img, show_results = True):
    w, h = img.shape[::-1]

    methods = ['cv.TM_CCOEFF', 'cv.TM_CCOEFF_NORMED', 'cv.TM_CCORR',
            'cv.TM_CCORR_NORMED', 'cv.TM_SQDIFF', 'cv.TM_SQDIFF_NORMED']

    meth = methods[1]

    template1 = cv.imread(data_dir_drive + '/data/templates/tem3.jpg', 0)
    template2 = cv.imread(data_dir_drive + '/data/templates/tem33.jpg', 0)
    w1, h1 = template1.shape[::-1]
    w2, h2 = template2.shape[::-1]

    method = eval(meth)

    # Apply template Matching
    res1 = cv.matchTemplate(img[:h//2, :w//2],template1,method)
    res2 = cv.matchTemplate(img[:h//2, :w//2],template2,method)
    min_val1, max_val1, min_loc1, max_loc1 = cv.minMaxLoc(res1)
    min_val2, max_val2, min_loc2, max_loc2 = cv.minMaxLoc(res2)

    # If the method is TM_SQDIFF or TM_SQDIFF_NORMED, take minimum
    if method in [cv.TM_SQDIFF, cv.TM_SQDIFF_NORMED]:
        top_left1 = min_loc1
        top_left2 = min_loc2
    else:
        top_left1 = max_loc1
        top_left2 = max_loc2

    #print(res1[top_left1[1], top_left1[0]])
    #print(res2[top_left2[1], top_left2[0]])

    if res1[top_left1[1], top_left1[0]] > 0.6 or res2[top_left2[1], top_left2[0]] > 0.6:
        if res1[top_left1[1], top_left1[0]] > res2[top_left2[1], top_left2[0]]:
            bottom_right1 = (top_left1[0] + w1, top_left1[1] + h1)
            if show_results:         
                cv.rectangle(img,top_left1, bottom_right1, 255, 2)
                cv2_imshow(cv.resize(img, (0, 0), fx=0.5, fy=0.5))
            return top_left1, bottom_right1, ((top_left1[0] + bottom_right1[0]) // 2, (top_left1[1] + bottom_right1[1]) // 2), img
        else:
            bottom_right2 = (top_left2[0] + w2, top_left2[1] + h2)
            if show_results:  
                cv.rectangle(img,top_left2, bottom_right2, 255, 2)       
                cv2_imshow(cv.resize(img, (0, 0), fx=0.5, fy=0.5))
            return top_left2, bottom_right2, ((top_left2[0] + bottom_right2[0]) // 2, (top_left2[1] + bottom_right2[1]) // 2), img
    else:
        return None, None, None, img

#4
def get_top_right(img, show_results = True):
    w, h = img.shape[::-1]

    methods = ['cv.TM_CCOEFF', 'cv.TM_CCOEFF_NORMED', 'cv.TM_CCORR',
            'cv.TM_CCORR_NORMED', 'cv.TM_SQDIFF', 'cv.TM_SQDIFF_NORMED']

    meth = methods[1]

    template1 = cv.imread(data_dir_drive + '/data/templates/tem4.jpg', 0)
    template2 = cv.imread(data_dir_drive + '/data/templates/tem44.jpg', 0)
    w1, h1 = template1.shape[::-1]
    w2, h2 = template2.shape[::-1]

    method = eval(meth)

    # Apply template Matching
    res1 = cv.matchTemplate(img[:h//2, w//2:],template1,method)
    res2 = cv.matchTemplate(img[:h//2, w//2:],template2,method)
    min_val1, max_val1, min_loc1, max_loc1 = cv.minMaxLoc(res1)
    min_val2, max_val2, min_loc2, max_loc2 = cv.minMaxLoc(res2)

    # If the method is TM_SQDIFF or TM_SQDIFF_NORMED, take minimum
    if method in [cv.TM_SQDIFF, cv.TM_SQDIFF_NORMED]:
        top_left1 = min_loc1
        top_left2 = min_loc2
    else:
        top_left1 = max_loc1
        top_left2 = max_loc2

    #print(res1[top_left1[1], top_left1[0]])
    #print(res2[top_left2[1], top_left2[0]])

    if res1[top_left1[1], top_left1[0]] > 0.6 or res2[top_left2[1], top_left2[0]] > 0.6:
        if res1[top_left1[1], top_left1[0]] > res2[top_left2[1], top_left2[0]]:
            top_left1 = (top_left1[0] + w//2, top_left1[1])
            bottom_right1 = (top_left1[0] + w1, top_left1[1] + h1)
            if show_results:      
                cv.rectangle(img,top_left1, bottom_right1, 255, 2)   
                cv2_imshow(cv.resize(img, (0, 0), fx=0.5, fy=0.5))
            return top_left1, bottom_right1, ((top_left1[0] + bottom_right1[0]) // 2, (top_left1[1] + bottom_right1[1]) // 2), img
        else:
            top_left2 = (top_left2[0] + w//2, top_left2[1])
            bottom_right2 = (top_left2[0] + w2, top_left2[1] + h2)
            if show_results:
                cv.rectangle(img,top_left2, bottom_right2, 255, 2)
                cv2_imshow(cv.resize(img, (0, 0), fx=0.5, fy=0.5))
            return top_left2, bottom_right2, ((top_left2[0] + bottom_right2[0]) // 2, (top_left2[1] + bottom_right2[1]) // 2), img
    else:
        return None, None, None, img

#5
def get_middle_right(img, show_results = True):
    w, h = img.shape[::-1]

    methods = ['cv.TM_CCOEFF', 'cv.TM_CCOEFF_NORMED', 'cv.TM_CCORR',
            'cv.TM_CCORR_NORMED', 'cv.TM_SQDIFF', 'cv.TM_SQDIFF_NORMED']

    meth = methods[1]

    template1 = cv.imread(data_dir_drive + '/data/templates/tem5.jpg', 0)
    template2 = cv.imread(data_dir_drive + '/data/templates/tem55.jpg', 0)
    template3 = cv.imread(data_dir_drive + '/data/templates/tem555.jpg', 0)
    w1, h1 = template1.shape[::-1]
    w2, h2 = template2.shape[::-1]
    w3, h3 = template3.shape[::-1]

    method = eval(meth)

    # Apply template Matching
    res1 = cv.matchTemplate(img[:, w//2:],template1,method)
    res2 = cv.matchTemplate(img[:, w//2:],template2,method)
    res3 = cv.matchTemplate(img[:, w//2:],template3,method)
    min_val1, max_val1, min_loc1, max_loc1 = cv.minMaxLoc(res1)
    min_val2, max_val2, min_loc2, max_loc2 = cv.minMaxLoc(res2)
    min_val3, max_val3, min_loc3, max_loc3 = cv.minMaxLoc(res3)

    # If the method is TM_SQDIFF or TM_SQDIFF_NORMED, take minimum
    if method in [cv.TM_SQDIFF, cv.TM_SQDIFF_NORMED]:
        top_left1 = min_loc1
        top_left2 = min_loc2
        top_left3 = min_loc3
    else:
        top_left1 = max_loc1
        top_left2 = max_loc2
        top_left3 = max_loc3

    #print(res1[top_left1[1], top_left1[0]])
    #print(res2[top_left2[1], top_left2[0]])
    #print(res3[top_left3[1], top_left3[0]])

    if res1[top_left1[1], top_left1[0]] > 0.75 or res2[top_left2[1], top_left2[0]] > 0.75 or res3[top_left3[1], top_left3[0]] > 0.75:
        L = [res1[top_left1[1], top_left1[0]], res2[top_left2[1], top_left2[0]], res3[top_left3[1], top_left3[0]]]
        i = np.argmax(np.array(L))
        if i == 0:
            top_left1 = (top_left1[0] + w//2, top_left1[1])
            bottom_right1 = (top_left1[0] + w1, top_left1[1] + h1)
            if show_results:      
                cv.rectangle(img,top_left1, bottom_right1, 255, 2)   
                cv2_imshow(cv.resize(img, (0, 0), fx=0.5, fy=0.5))
            return top_left1, bottom_right1, ((top_left1[0] + bottom_right1[0]) // 2, (top_left1[1] + bottom_right1[1]) // 2), img

        if i == 1:
            top_left2 = (top_left2[0] + w//2, top_left2[1])
            bottom_right2 = (top_left2[0] + w2, top_left2[1] + h2)
            if show_results:      
                cv.rectangle(img,top_left2, bottom_right2, 255, 2)   
                cv2_imshow(cv.resize(img, (0, 0), fx=0.5, fy=0.5))
            return top_left2, bottom_right2, ((top_left2[0] + bottom_right2[0]) // 2, (top_left2[1] + bottom_right2[1]) // 2), img

        if i == 2:
            top_left3 = (top_left3[0] + w//2, top_left3[1])
            bottom_right3 = (top_left3[0] + w3, top_left3[1] + h3)
            if show_results:      
                cv.rectangle(img,top_left3, bottom_right3, 255, 2)   
                cv2_imshow(cv.resize(img, (0, 0), fx=0.5, fy=0.5))
            return top_left3, bottom_right3, ((top_left3[0] + bottom_right3[0]) // 2, (top_left3[1] + bottom_right3[1]) // 2), img

    else:
        return None, None, None, img

#6
def get_middle_left(img, show_results = True):
    w, h = img.shape[::-1]

    methods = ['cv.TM_CCOEFF', 'cv.TM_CCOEFF_NORMED', 'cv.TM_CCORR',
            'cv.TM_CCORR_NORMED', 'cv.TM_SQDIFF', 'cv.TM_SQDIFF_NORMED']

    meth = methods[1]

    template1 = cv.imread(data_dir_drive + '/data/templates/tem6.jpg', 0)
    template2 = cv.imread(data_dir_drive + '/data/templates/tem66.jpg', 0)
    template3 = cv.imread(data_dir_drive + '/data/templates/tem666.jpg', 0)
    w1, h1 = template1.shape[::-1]
    w2, h2 = template2.shape[::-1]
    w3, h3 = template3.shape[::-1]

    method = eval(meth)

    # Apply template Matching
    res1 = cv.matchTemplate(img[:, :w//2],template1,method)
    res2 = cv.matchTemplate(img[:, :w//2],template2,method)
    res3 = cv.matchTemplate(img[:, :w//2],template3,method)
    min_val1, max_val1, min_loc1, max_loc1 = cv.minMaxLoc(res1)
    min_val2, max_val2, min_loc2, max_loc2 = cv.minMaxLoc(res2)
    min_val3, max_val3, min_loc3, max_loc3 = cv.minMaxLoc(res3)

    # If the method is TM_SQDIFF or TM_SQDIFF_NORMED, take minimum
    if method in [cv.TM_SQDIFF, cv.TM_SQDIFF_NORMED]:
        top_left1 = min_loc1
        top_left2 = min_loc2
        top_left3 = min_loc3
    else:
        top_left1 = max_loc1
        top_left2 = max_loc2
        top_left3 = max_loc3

    #print(res1[top_left1[1], top_left1[0]])
    #print(res2[top_left2[1], top_left2[0]])
    #print(res3[top_left3[1], top_left3[0]])

    if res1[top_left1[1], top_left1[0]] > 0.8 or res2[top_left2[1], top_left2[0]] > 0.8 or res3[top_left3[1], top_left3[0]] > 0.8:
        L = [res1[top_left1[1], top_left1[0]], res2[top_left2[1], top_left2[0]], res3[top_left3[1], top_left3[0]]]
        i = np.argmax(np.array(L))
        if i == 0:
            top_left1 = (top_left1[0], top_left1[1])
            bottom_right1 = (top_left1[0] + w1, top_left1[1] + h1)
            if show_results:      
                cv.rectangle(img,top_left1, bottom_right1, 255, 2)   
                cv2_imshow(cv.resize(img, (0, 0), fx=0.5, fy=0.5))
            return top_left1, bottom_right1, ((top_left1[0] + bottom_right1[0]) // 2, (top_left1[1] + bottom_right1[1]) // 2), img

        if i == 1:
            top_left2 = (top_left2[0], top_left2[1])
            bottom_right2 = (top_left2[0] + w2, top_left2[1] + h2)
            if show_results:      
                cv.rectangle(img,top_left2, bottom_right2, 255, 2)   
                cv2_imshow(cv.resize(img, (0, 0), fx=0.5, fy=0.5))
            return top_left2, bottom_right2, ((top_left2[0] + bottom_right2[0]) // 2, (top_left2[1] + bottom_right2[1]) // 2), img

        if i == 2:
            top_left3 = (top_left3[0], top_left3[1])
            bottom_right3 = (top_left3[0] + w3, top_left3[1] + h3)
            if show_results:      
                cv.rectangle(img,top_left3, bottom_right3, 255, 2)   
                cv2_imshow(cv.resize(img, (0, 0), fx=0.5, fy=0.5))
            return top_left3, bottom_right3, ((top_left3[0] + bottom_right3[0]) // 2, (top_left3[1] + bottom_right3[1]) // 2), img

    else:
        return None, None, None, img

#7
def get_corners(img, show_results = True):
    _, _, bottom_left_corner, _ = get_bottom_left(img, False)
    _, _, bottom_right_corner, _ = get_bottom_right(img, False)
    _, _, top_right_corner, _ = get_top_right(img, False)
    _, _, top_left_corner, _ = get_top_left(img, False)

    if bottom_left_corner is None or bottom_right_corner is None or top_right_corner is None or top_left_corner is None:

        if bottom_left_corner is None and bottom_right_corner is None:
            print("I can't calculate the top pockets")
        elif bottom_left_corner is None or bottom_right_corner is None:
            if top_right_corner is None or top_left_corner is None:
                print("I can't calculate the top pockets")
            else:
                middle = (top_right_corner[0] + top_left_corner[0]) // 2
                if bottom_left_corner is None:
                    bottom_left_corner = (bottom_right_corner[0] - 2 *(bottom_right_corner[0] - middle), bottom_right_corner[1])
                else:
                    bottom_right_corner = (bottom_left_corner[0] + 2 *(middle - bottom_left_corner[0]), bottom_left_corner[1])

        if top_right_corner is None and top_left_corner is None:
            print("I can't calculate the bottom pockets")
        elif top_right_corner is None or top_left_corner is None:
            if bottom_left_corner is None or bottom_right_corner is None:
                print("I can't calculate the top pockets")
            else:
                middle = (bottom_left_corner[0] + bottom_right_corner[0]) // 2
                if top_left_corner is None:
                    top_left_corner = (top_right_corner[0] - 2 *(top_right_corner[0] - middle), top_right_corner[1])
                else:
                    top_right_corner = (top_left_corner[0] + 2 *(middle - top_left_corner[0]), top_left_corner[1])

    if show_results:
        img = cv.circle(img, bottom_left_corner, 5, 255, 2) 
        img = cv.circle(img, bottom_right_corner, 5, 255, 2) 
        img = cv.circle(img, top_right_corner, 5, 255, 2) 
        img = cv.circle(img, top_left_corner, 5, 255, 2) 

        cv2_imshow(cv.resize(img, (0, 0), fx=0.5, fy=0.5))

    return bottom_left_corner, bottom_right_corner, top_right_corner, top_left_corner

#8
for i in range(1, 51):
    print(i)
    imgColor = cv.imread(folder_in + 'Task1/'+ str(i) +'.jpg')
    img = cv.cvtColor(imgColor, cv.COLOR_BGR2GRAY)
    
    bottom_left_corner, bottom_right_corner, top_right_corner, top_left_corner = get_corners(img.copy(), True)

#9
def get_balls_template_matching(img, template, threshold):
    methods = ['cv.TM_CCOEFF', 'cv.TM_CCOEFF_NORMED', 'cv.TM_CCORR',
            'cv.TM_CCORR_NORMED', 'cv.TM_SQDIFF', 'cv.TM_SQDIFF_NORMED']
    meth = methods[1]
    method = eval(meth)

    # Apply template Matching
    res = cv.matchTemplate(img,template,method)

    min_val, max_val, min_loc, max_loc = cv.minMaxLoc(res)

    # If the method is TM_SQDIFF or TM_SQDIFF_NORMED, take minimum
    if method in [cv.TM_SQDIFF, cv.TM_SQDIFF_NORMED]:
        top_left = min_loc
    else:
        top_left = max_loc

    loc = np.where(res >= threshold)
    return loc[::-1]

#10
# Malisiewicz et al.
def non_max_suppression_fast(boxes, overlapThresh):
	# if there are no boxes, return an empty list
	if len(boxes) == 0:
		return []
	# if the bounding boxes integers, convert them to floats --
	# this is important since we'll be doing a bunch of divisions
	if boxes.dtype.kind == "i":
		boxes = boxes.astype("float")
	# initialize the list of picked indexes	
	pick = []
	# grab the coordinates of the bounding boxes
	x1 = boxes[:,0]
	y1 = boxes[:,1]
	x2 = boxes[:,2]
	y2 = boxes[:,3]
	# compute the area of the bounding boxes and sort the bounding
	# boxes by the bottom-right y-coordinate of the bounding box
	area = (x2 - x1 + 1) * (y2 - y1 + 1)
	idxs = np.argsort(y2)
	# keep looping while some indexes still remain in the indexes
	# list
	while len(idxs) > 0:
		# grab the last index in the indexes list and add the
		# index value to the list of picked indexes
		last = len(idxs) - 1
		i = idxs[last]
		pick.append(i)
		# find the largest (x, y) coordinates for the start of
		# the bounding box and the smallest (x, y) coordinates
		# for the end of the bounding box
		xx1 = np.maximum(x1[i], x1[idxs[:last]])
		yy1 = np.maximum(y1[i], y1[idxs[:last]])
		xx2 = np.minimum(x2[i], x2[idxs[:last]])
		yy2 = np.minimum(y2[i], y2[idxs[:last]])
		# compute the width and height of the bounding box
		w = np.maximum(0, xx2 - xx1 + 1)
		h = np.maximum(0, yy2 - yy1 + 1)
		# compute the ratio of overlap
		overlap = (w * h) / area[idxs[:last]]
		# delete all indexes from the index list that have
		idxs = np.delete(idxs, np.concatenate(([last],
			np.where(overlap > overlapThresh)[0])))
	# return only the bounding boxes that were picked using the
	# integer data type
	return boxes[pick].astype("int")

#11
def false_pozitive(img, box, threshold):
    patch = img[box[1]:box[3], box[0]:box[2]]

    low_green = (46, 100, 0) #46
    high_green = (65, 255, 255) #65

    patch = cv.cvtColor(patch, cv.COLOR_BGR2HSV)
    # get the mask for green
    mask_hsv = cv.inRange(patch, low_green, high_green) // 255
    
    return np.sum(mask_hsv) // (mask_hsv.shape[0] * mask_hsv.shape[1]) > threshold

    return False

#12
def get_box_color(img, box, show_results = True):

    img = cv.copyMakeBorder(img, 10, 10, 10, 10, cv.BORDER_REFLECT)
    
    patch = img[box[1]+5:box[3]+15, box[0]+5:box[2]+15]
    
    if show_results:
        cv2_imshow(patch)

    low_green = (0, 100, 0)
    high_green = (60, 255, 60)

    low_white = (80, 140, 65)
    high_white = (255, 255, 255)
 
    low_blue = (91, 0, 0)
    high_blue = (180, 135, 55)

    low_dark_green = (37, 60, 0)
    high_dark_green = (60, 255, 30)

    low_yellow = (0, 107, 68)
    high_yellow = (60, 255, 255)

    low_red = (0, 0, 50)
    high_red = (30, 50, 255)

    low_pink = (45, 0, 166)
    high_pink = (255, 162, 255)

    low_brown = (0, 52, 37)
    high_brown = (45, 105, 166)

    low_black = (0, 0, 0)
    high_black = (45, 38, 55)

    #patch_hsv = cv.cvtColor(patch, cv.COLOR_BGR2HSV)

    colors_dict = {
        0: "green",
        1: "white",
        2: "red",
        3: "blue",
        4: "pink",
        5: "dark_green",
        6: "brown",
        7: "yellow",
        8: "black"
    }

    hist = np.zeros(9)

    height, width, depth = patch.shape
    for i in range(0, height):
        for j in range(0, width):
            score_green = np.sum(cv.inRange(patch, low_green, high_green))
            score_white = np.sum(cv.inRange(patch, low_white, high_white))
            score_red = np.sum(cv.inRange(patch,low_red,high_red))
            score_blue = np.sum(cv.inRange(patch, low_blue, high_blue))
            score_pink = np.sum(cv.inRange(patch, low_pink, high_pink))
            score_dark_green = np.sum(cv.inRange(patch, low_dark_green, high_dark_green))
            score_brown = np.sum(cv.inRange(patch, low_brown, high_brown))
            score_yellow = np.sum(cv.inRange(patch, low_yellow, high_yellow)) 
            score_black = np.sum(cv.inRange(patch, low_black, high_black))

            scores = [score_green, score_white, score_red, score_blue, score_pink, score_dark_green, score_brown, score_yellow, score_black]
            hist = np.add(hist, scores)
            #idx = np.argmax(np.array(scores))
        
    idxs = np.argsort(hist)

    if show_results:
        print(colors_dict[idxs[8]], colors_dict[idxs[7]], colors_dict[idxs[6]])
        print(hist[idxs[8]], hist[idxs[7]], hist[idxs[6]])
    return [colors_dict[idxs[8]], colors_dict[idxs[7]], colors_dict[idxs[6]]]

#13
def process_colors(colors):
    result = {
        "white" : 0,
        "red" : 0,
        "blue" : 0,
        "pink" : 0,
        "dark_green" : 0,
        "brown" : 0,
        "yellow" : 0,
        "black" : 0
    }

    for color in colors:
        if color[0] != "green":
            result[color[0]] += 1
        else:
            if color[1] != "pink" and color[2] == "pink":
                result["pink"] += 1
            else:
                if color[1] == "white" and color[2] == "dark_green":
                    result["dark_green"] += 1
                else:
                    result[color[1]] += 1
    
    if result["white"] > 1:
        result["white"] = 1
    if result["black"] > 1:
        result["black"] = 1
    if result["yellow"] > 1:
        result["red"] = result["red"] + result["yellow"] - 1
        result["yellow"] = 1
    if result["brown"] > 1:
        result["red"] = result["red"] + result["brown"] - 1
        result["brown"] = 1
    if result["pink"] > 1:
        result["red"] = result["red"] + result["pink"] - 1
        result["pink"] = 1

    if result["dark_green"] == 2 and result["blue"] == 0:
        result["dark_green"] = 1
        result["blue"] = 1
    if result["dark_green"] == 0 and result["blue"] == 2:
        result["dark_green"] = 1
        result["blue"] = 1

    if result["dark_green"] > 1:
        result["dark_green"] = 1

    if result["blue"] > 1:
        result["blue"] = 1

    num_balls = 0
    for _, number in result.items(): 
        num_balls += number

    return result, num_balls

#14
results = []
acc_colors = 0
acc_numbers = 0

for i in range(1, 51):
    print(i)
    imgColor = cv.imread(folder_in + 'Task1/'+ str(i) +'.jpg')
    img = cv.cvtColor(imgColor, cv.COLOR_BGR2GRAY)
    
    bottom_left_corner, bottom_right_corner, top_right_corner, top_left_corner = get_corners(img.copy(), False)

    #points_template = np.float32([(2000, 2000), (0, 2000), (0, 0), (2000, 0)]) # the points of the 4 corners (table) in the template image
    #points_query = np.float32([bottom_left_corner, bottom_right_corner, top_right_corner, top_left_corner]) # the points of the 4 corners (table) in the test image
    #homography = cv.getPerspectiveTransform(points_query, points_template)
    #height, width = 2000, 2000 # the shape with respect to the template image
    #aligned_image2 = cv.warpPerspective(imgColor, homography, (width, height), flags=cv.INTER_NEAREST)

    mask = np.zeros(imgColor.shape, dtype=np.uint8)
    
    corners = np.array([[bottom_left_corner, bottom_right_corner, top_right_corner, top_left_corner]], dtype=np.int32)
    channel_count = imgColor.shape[2]  # i.e. 3 or 4 depending on your image
    ignore_mask_color = (255,)*channel_count
    cv.fillPoly(mask, corners, ignore_mask_color)
    # apply the mask
    masked_image = cv.bitwise_and(imgColor, mask)

    masked_image = masked_image[top_right_corner[1]:bottom_left_corner[1], bottom_left_corner[0]:bottom_right_corner[0]]

    ####################

    imgColor = masked_image
    img = cv.cvtColor(imgColor, cv.COLOR_BGR2GRAY)

    w, h = img.shape[::-1]

    thresholds = [0.7, 0.7, 0.7, 0.8, 0.8, 0.8, 0.7]
    locations = []

    for template_numer in range(1, 8):
        template = cv.imread(data_dir_drive + '/data/templates/ball' + str(template_numer) + '.jpg', 0)
        w_T, h_T = template.shape[::-1]

        threshold = thresholds[template_numer-1]
        loc = get_balls_template_matching(img, template, threshold)
        for pt in zip(*loc):
            locations.append((pt[0], pt[1], pt[0] + w_T, pt[1] + h_T))

    boxes = non_max_suppression_fast(np.array(locations), 0.3)

    copyImgColor = imgColor.copy()
    colors = []

    for box in boxes:
        if not false_pozitive(imgColor, box, 0.5):
            color = get_box_color(imgColor, box)
            colors.append(color)
            cv.rectangle(copyImgColor, (box[0], box[1]), (box[2], box[3]), (0,0,255), 2) 

    balls, num_balls = process_colors(colors)
    
    # load ground truth
    ground_truth = {}
    num_balls_truth  = 0
    with open(folder_in + 'Task1/ground-truth/' + str(i) + '_gt.txt') as f:
        num_balls_truth = int(f.readline())
        lines = f.readlines()
        for line in lines:
            info = line.replace('\n', '').split(' ')
            nr = int(info[0])
            color = info[1]

            ground_truth[color] = nr

    ok = True
    nr = 0
    for color, number in balls.items(): 
        nr += number
        if color == 'dark_green':
            if ground_truth['green'] != number:
                ok = False
        else:
            if ground_truth[color] != number:
                ok = False
                print("greseala - ground_truth:", ground_truth[color], "me: ", number, color)

    if ok:
        acc_colors += 1
    else:

        print(num_balls_truth, num_balls)     
        print(ground_truth)
        print(balls)   
        print(ok)

    if num_balls_truth == nr:
        acc_numbers += 1

        ####################

        cv2_imshow(copyImgColor)
    #cv.imwrite(data_dir_drive + '/data/aux/' + str(i) + '.jpg', masked_image)

print('acc numbers = ' + str(acc_numbers * 2) + '%')
print('acc colors = ' + str(acc_colors * 2) + '%')

#15
def task1(imgColor):

    img = cv.cvtColor(imgColor, cv.COLOR_BGR2GRAY)
    
    bottom_left_corner, bottom_right_corner, top_right_corner, top_left_corner = get_corners(img.copy(), False)

    #points_template = np.float32([(2000, 2000), (0, 2000), (0, 0), (2000, 0)]) # the points of the 4 corners (table) in the template image
    #points_query = np.float32([bottom_left_corner, bottom_right_corner, top_right_corner, top_left_corner]) # the points of the 4 corners (table) in the test image
    #homography = cv.getPerspectiveTransform(points_query, points_template)
    #height, width = 2000, 2000 # the shape with respect to the template image
    #aligned_image2 = cv.warpPerspective(imgColor, homography, (width, height), flags=cv.INTER_NEAREST)

    mask = np.zeros(imgColor.shape, dtype=np.uint8)
    
    corners = np.array([[bottom_left_corner, bottom_right_corner, top_right_corner, top_left_corner]], dtype=np.int32)
    channel_count = imgColor.shape[2]  # i.e. 3 or 4 depending on your image
    ignore_mask_color = (255,)*channel_count
    cv.fillPoly(mask, corners, ignore_mask_color)
    # apply the mask
    masked_image = cv.bitwise_and(imgColor, mask)

    masked_image = masked_image[top_right_corner[1]:bottom_left_corner[1], bottom_left_corner[0]:bottom_right_corner[0]]

    ####################

    imgColor = masked_image
    img = cv.cvtColor(imgColor, cv.COLOR_BGR2GRAY)

    w, h = img.shape[::-1]

    thresholds = [0.7, 0.7, 0.7, 0.8, 0.8, 0.8, 0.7]
    locations = []

    for template_numer in range(1, 8):
        template = cv.imread(data_dir_drive + '/data/templates/ball' + str(template_numer) + '.jpg', 0)
        w_T, h_T = template.shape[::-1]

        threshold = thresholds[template_numer-1]
        loc = get_balls_template_matching(img, template, threshold)
        for pt in zip(*loc):
            locations.append((pt[0], pt[1], pt[0] + w_T, pt[1] + h_T))

    boxes = non_max_suppression_fast(np.array(locations), 0.3)

    copyImgColor = imgColor.copy()
    colors = []

    for box in boxes:
        if not false_pozitive(imgColor, box, 0.5):
            color = get_box_color(imgColor, box, False)
            colors.append(color)
            #cv.rectangle(copyImgColor, (box[0], box[1]), (box[2], box[3]), (0,0,255), 2) 

    balls, num_balls = process_colors(colors)
    

    return balls, num_balls

#16
img = cv.imread(folder_in + 'Task1/1.jpg')
balls, num_balls = task1(img)
print(balls, num_balls)

"""# **Task 2**
You receive a test set of 25 videos containing a player making a shot. In each of the 25 videos the snooker table is seen from above. The task is to decide whether a ball was potted into a pocket and if so recognize the color of the potted ball and the pocket (one of the six pockets) where the ball was potted. By solving correctly this task for a video you will earn 0.04 points as follows: if the correct answer is NO (no ball was poted) you will receive 0.04 points if your output is NO, else you will receive 0.02 points if your output is YES, 0.01 points for specifying the correct pocket (1 - top left corner, 2 - top right corner, 3 - bottom left corner, 4 - bottom right corner, 5 - middle left corner, 6 - middle right corner) and 0.01 points for specifying the correct color of the potted ball. In each of the 25 videos there will be a maximum of one ball being potted.
"""

#1
def get_pockets(img, show_results = True):
    _, _, bottom_left_corner, _ = get_bottom_left(img, False)
    _, _, bottom_right_corner, _ = get_bottom_right(img, False)
    _, _, top_right_corner, _ = get_top_right(img, False)
    _, _, top_left_corner, _ = get_top_left(img, False)
    _, _, middle_right_corner, _ = get_middle_right(img, False)
    _, _, middle_left_corner, _ = get_middle_left(img, False)

    if bottom_left_corner is None or bottom_right_corner is None or top_right_corner is None or top_left_corner is None or middle_right_corner is None or middle_left_corner is None:

        if bottom_left_corner is None and bottom_right_corner is None:
            print("I can't calculate the top pockets")
        elif bottom_left_corner is None or bottom_right_corner is None:
            if top_right_corner is None or top_left_corner is None:
                print("I can't calculate the top pockets")
            else:
                middle = (top_right_corner[0] + top_left_corner[0]) // 2
                if bottom_left_corner is None:
                    bottom_left_corner = (bottom_right_corner[0] - 2 *(bottom_right_corner[0] - middle), bottom_right_corner[1])
                else:
                    bottom_right_corner = (bottom_left_corner[0] + 2 *(middle - bottom_left_corner[0]), bottom_left_corner[1])

        if top_right_corner is None and top_left_corner is None:
            print("I can't calculate the bottom pockets")
        elif top_right_corner is None or top_left_corner is None:
            if bottom_left_corner is None or bottom_right_corner is None:
                print("I can't calculate the top pockets")
            else:
                middle = (bottom_left_corner[0] + bottom_right_corner[0]) // 2
                if top_left_corner is None:
                    top_left_corner = (top_right_corner[0] - 2 *(top_right_corner[0] - middle), top_right_corner[1])
                else:
                    top_right_corner = (top_left_corner[0] + 2 *(middle - top_left_corner[0]), top_left_corner[1])

        if middle_right_corner is None and middle_left_corner is None:
            print("I can't calculate the bottom pockets")
        elif middle_right_corner is None or middle_left_corner is None:
            if bottom_left_corner is None or bottom_right_corner is None:
                print("I can't calculate the middle pockets")
            else:
                middle = (bottom_left_corner[0] + bottom_right_corner[0]) // 2
                if middle_left_corner is None:
                    middle_left_corner = (middle_right_corner[0] - 2 *(middle_right_corner[0] - middle), middle_right_corner[1])
                else:
                    middle_right_corner = (middle_left_corner[0] + 2 *(middle - middle_left_corner[0]), middle_left_corner[1])


    if show_results:
        img = cv.circle(img, bottom_left_corner, 5, 255, 2) 
        img = cv.circle(img, bottom_right_corner, 5, 255, 2) 
        img = cv.circle(img, top_right_corner, 5, 255, 2) 
        img = cv.circle(img, top_left_corner, 5, 255, 2) 
        img = cv.circle(img, middle_left_corner, 5, 255, 2) 
        img = cv.circle(img, middle_right_corner, 5, 255, 2) 

        cv2_imshow(cv.resize(img, (0, 0), fx=0.5, fy=0.5))

    return bottom_left_corner, bottom_right_corner, top_right_corner, top_left_corner, middle_left_corner, middle_right_corner

#2
for i in range(1, 51):
    print(i)
    imgColor = cv.imread(folder_in + 'Task1/'+ str(i) +'.jpg')
    img = cv.cvtColor(imgColor, cv.COLOR_BGR2GRAY)
    
    _, _, _, _, _, _ = get_pockets(img.copy(), True)

#3
def get_areas_of_interet(frame):
    imgColor = frame.copy()
    img = cv.cvtColor(imgColor, cv.COLOR_BGR2GRAY)
            
    bottom_left_corner, bottom_right_corner, top_right_corner, top_left_corner, middle_left_corner, middle_right_corner = get_pockets(img.copy(), False)
    if bottom_left_corner is None or bottom_right_corner is None or top_right_corner is None or top_left_corner is None or middle_right_corner is None or middle_left_corner is None:
        return None, None, None, None, None, None

    mask = np.zeros(imgColor.shape, dtype=np.uint8)
            
    corners = np.array([[bottom_left_corner, bottom_right_corner, top_right_corner, top_left_corner]], dtype=np.int32)
    channel_count = imgColor.shape[2]  # i.e. 3 or 4 depending on your image
    ignore_mask_color = (255,)*channel_count
    cv.fillPoly(mask, corners, ignore_mask_color)
    # apply the mask
    masked_image = cv.bitwise_and(imgColor, mask)
    masked_image = masked_image[top_right_corner[1]:bottom_left_corner[1], bottom_left_corner[0]:bottom_right_corner[0]]
    padidngSize = 100
    masked_image = cv.copyMakeBorder(masked_image, padidngSize, padidngSize, padidngSize, padidngSize, cv.BORDER_CONSTANT)

    x_cut = bottom_left_corner[0] - padidngSize
    y_cut = top_right_corner[1] - padidngSize

    top_left_corner = (top_left_corner[0] - x_cut, top_left_corner[1] - y_cut) # 1
    top_right_corner = (top_right_corner[0] - x_cut, top_right_corner[1] - y_cut) # 2
    bottom_left_corner = (bottom_left_corner[0] - x_cut, bottom_left_corner[1] - y_cut) # 3
    bottom_right_corner = (bottom_right_corner[0] - x_cut, bottom_right_corner[1] - y_cut) # 4
    middle_left_corner = (middle_left_corner[0] - x_cut, middle_left_corner[1] - y_cut) # 5
    middle_right_corner = (middle_right_corner[0] - x_cut, middle_right_corner[1] - y_cut) # 6
    ###################

    interestRaze = 90

    area_of_interest_1 = masked_image[top_left_corner[1] - interestRaze : top_left_corner[1] + interestRaze, top_left_corner[0] - interestRaze : top_left_corner[0] + interestRaze]
    area_of_interest_2 = masked_image[top_right_corner[1] - interestRaze : top_right_corner[1] + interestRaze, top_right_corner[0] - interestRaze : top_right_corner[0] + interestRaze]
    area_of_interest_3 = masked_image[bottom_left_corner[1] - interestRaze : bottom_left_corner[1] + interestRaze, bottom_left_corner[0] - interestRaze : bottom_left_corner[0] + interestRaze]
    area_of_interest_4 = masked_image[bottom_right_corner[1] - interestRaze : bottom_right_corner[1] + interestRaze, bottom_right_corner[0] - interestRaze : bottom_right_corner[0] + interestRaze]
    area_of_interest_5 = masked_image[middle_left_corner[1] - interestRaze : middle_left_corner[1] + interestRaze, middle_left_corner[0] - interestRaze : middle_left_corner[0] + interestRaze]
    area_of_interest_6 = masked_image[middle_right_corner[1] - interestRaze : middle_right_corner[1] + interestRaze, middle_right_corner[0] - interestRaze : middle_right_corner[0] + interestRaze]

    return area_of_interest_1, area_of_interest_2, area_of_interest_3, area_of_interest_4, area_of_interest_5, area_of_interest_6

#4
from collections import defaultdict

cap = cv.VideoCapture(folder_in + 'Task2/23.mp4')  
current_frame = 0

ret, frame = cap.read() # Read the frame
area_of_interest_1, area_of_interest_2, area_of_interest_3, area_of_interest_4, area_of_interest_5, area_of_interest_6 = get_areas_of_interet(frame)
while area_of_interest_1 is None:
    current_frame = current_frame + 1
    ret, frame = cap.read() # Read the frame
    area_of_interest_1, area_of_interest_2, area_of_interest_3, area_of_interest_4, area_of_interest_5, area_of_interest_6 = get_areas_of_interet(frame)

area_of_interest_1_gray = cv.cvtColor(area_of_interest_1,cv.COLOR_BGR2GRAY)
old_area_of_interest_1_gray = area_of_interest_1_gray
area_of_interest_2_gray = cv.cvtColor(area_of_interest_2,cv.COLOR_BGR2GRAY)
old_area_of_interest_2_gray = area_of_interest_2_gray
area_of_interest_3_gray = cv.cvtColor(area_of_interest_3,cv.COLOR_BGR2GRAY)
old_area_of_interest_3_gray = area_of_interest_3_gray
area_of_interest_4_gray = cv.cvtColor(area_of_interest_4,cv.COLOR_BGR2GRAY)
old_area_of_interest_4_gray = area_of_interest_4_gray
area_of_interest_5_gray = cv.cvtColor(area_of_interest_5,cv.COLOR_BGR2GRAY)
old_area_of_interest_5_gray = area_of_interest_5_gray
area_of_interest_6_gray = cv.cvtColor(area_of_interest_6,cv.COLOR_BGR2GRAY)
old_area_of_interest_6_gray = area_of_interest_6_gray

interestCorners = defaultdict(list)

while(cap.isOpened()): 
    ret, frame = cap.read() # Read the frame
    if ret is True: 
        current_frame = current_frame + 1 

        if current_frame % 2:
            continue
        else:
            area_of_interest_1, area_of_interest_2, area_of_interest_3, area_of_interest_4, area_of_interest_5, area_of_interest_6 = get_areas_of_interet(frame)
            if area_of_interest_1 is None:
                continue
            area_of_interest_1_gray = cv.cvtColor(area_of_interest_1,cv.COLOR_BGR2GRAY)
            area_of_interest_2_gray = cv.cvtColor(area_of_interest_2,cv.COLOR_BGR2GRAY)
            area_of_interest_3_gray = cv.cvtColor(area_of_interest_3,cv.COLOR_BGR2GRAY)
            area_of_interest_4_gray = cv.cvtColor(area_of_interest_4,cv.COLOR_BGR2GRAY)
            area_of_interest_5_gray = cv.cvtColor(area_of_interest_5,cv.COLOR_BGR2GRAY)
            area_of_interest_6_gray = cv.cvtColor(area_of_interest_6,cv.COLOR_BGR2GRAY)
                    
            # compute frame diff            
            temp_1 = np.abs(np.float16(area_of_interest_1_gray) - np.float16(old_area_of_interest_1_gray))
            diff_1 = np.uint8(temp_1)
            temp_2 = np.abs(np.float16(area_of_interest_2_gray) - np.float16(old_area_of_interest_2_gray))
            diff_2 = np.uint8(temp_2)
            temp_3 = np.abs(np.float16(area_of_interest_3_gray) - np.float16(old_area_of_interest_3_gray))
            diff_3 = np.uint8(temp_3)
            temp_4 = np.abs(np.float16(area_of_interest_4_gray) - np.float16(old_area_of_interest_4_gray))
            diff_4 = np.uint8(temp_4)
            temp_5 = np.abs(np.float16(area_of_interest_5_gray) - np.float16(old_area_of_interest_5_gray))
            diff_5 = np.uint8(temp_5)
            temp_6 = np.abs(np.float16(area_of_interest_6_gray) - np.float16(old_area_of_interest_6_gray))
            diff_6 = np.uint8(temp_6)

            #print(current_frame)
            #cv2_imshow(np.concatenate((diff_1, diff_2, diff_3, diff_4, diff_5, diff_6), axis=1))

            diffs = [diff_1, diff_2, diff_3, diff_4, diff_5, diff_6]

            for i in range(len(diffs)):
                if np.sum(diffs[i]) // 255 > 50:
                    interestCorners[i+1].append((diffs[i], current_frame))

            old_area_of_interest_1_gray = area_of_interest_1_gray
            old_area_of_interest_2_gray = area_of_interest_2_gray
            old_area_of_interest_3_gray = area_of_interest_3_gray
            old_area_of_interest_4_gray = area_of_interest_4_gray
            old_area_of_interest_5_gray = area_of_interest_5_gray
            old_area_of_interest_6_gray = area_of_interest_6_gray

    else:
        break

# after playing the video, release the video capture    
cap.release()
# close all the frames
cv.destroyAllWindows()

#5
for i in range(1, 7):
    all = np.zeros((100,0))
    sum = 0
    print(i)
    for j in interestCorners[i]:
        print(j[1])
        cv2_imshow(j[0])
    #sum = np.sum(all) // 255
    #if not np.array_equal(all, np.zeros((100,0))):
        #print(i)
        #cv2_imshow(all)
        #print(sum)

#6
def get_suspicious_pockets(path, suspicious_color):

    cap = cv.VideoCapture(path)  
    current_frame = 0

    ret, frame = cap.read() # Read the frame
    area_of_interest_1, area_of_interest_2, area_of_interest_3, area_of_interest_4, area_of_interest_5, area_of_interest_6 = get_areas_of_interet(frame)
    while area_of_interest_1 is None:
        current_frame = current_frame + 1
        ret, frame = cap.read() # Read the frame
        area_of_interest_1, area_of_interest_2, area_of_interest_3, area_of_interest_4, area_of_interest_5, area_of_interest_6 = get_areas_of_interet(frame)

    area_of_interest_1_gray = cv.cvtColor(area_of_interest_1,cv.COLOR_BGR2GRAY)
    old_area_of_interest_1_gray = area_of_interest_1_gray
    area_of_interest_2_gray = cv.cvtColor(area_of_interest_2,cv.COLOR_BGR2GRAY)
    old_area_of_interest_2_gray = area_of_interest_2_gray
    area_of_interest_3_gray = cv.cvtColor(area_of_interest_3,cv.COLOR_BGR2GRAY)
    old_area_of_interest_3_gray = area_of_interest_3_gray
    area_of_interest_4_gray = cv.cvtColor(area_of_interest_4,cv.COLOR_BGR2GRAY)
    old_area_of_interest_4_gray = area_of_interest_4_gray
    area_of_interest_5_gray = cv.cvtColor(area_of_interest_5,cv.COLOR_BGR2GRAY)
    old_area_of_interest_5_gray = area_of_interest_5_gray
    area_of_interest_6_gray = cv.cvtColor(area_of_interest_6,cv.COLOR_BGR2GRAY)
    old_area_of_interest_6_gray = area_of_interest_6_gray

    interestCorners = defaultdict(list)

    while(cap.isOpened()): 
        ret, frame = cap.read() # Read the frame
        if ret is True: 
            current_frame = current_frame + 1 

            if current_frame % 2:
                continue
            else:
                area_of_interest_1, area_of_interest_2, area_of_interest_3, area_of_interest_4, area_of_interest_5, area_of_interest_6 = get_areas_of_interet(frame.copy())
                if area_of_interest_1 is None:
                    continue
                area_of_interest_1_gray = cv.cvtColor(area_of_interest_1,cv.COLOR_BGR2GRAY)
                area_of_interest_2_gray = cv.cvtColor(area_of_interest_2,cv.COLOR_BGR2GRAY)
                area_of_interest_3_gray = cv.cvtColor(area_of_interest_3,cv.COLOR_BGR2GRAY)
                area_of_interest_4_gray = cv.cvtColor(area_of_interest_4,cv.COLOR_BGR2GRAY)
                area_of_interest_5_gray = cv.cvtColor(area_of_interest_5,cv.COLOR_BGR2GRAY)
                area_of_interest_6_gray = cv.cvtColor(area_of_interest_6,cv.COLOR_BGR2GRAY)
                        
                # compute frame diff            
                temp_1 = np.abs(np.float16(area_of_interest_1_gray) - np.float16(old_area_of_interest_1_gray))
                diff_1 = np.uint8(temp_1)
                temp_2 = np.abs(np.float16(area_of_interest_2_gray) - np.float16(old_area_of_interest_2_gray))
                diff_2 = np.uint8(temp_2)
                temp_3 = np.abs(np.float16(area_of_interest_3_gray) - np.float16(old_area_of_interest_3_gray))
                diff_3 = np.uint8(temp_3)
                temp_4 = np.abs(np.float16(area_of_interest_4_gray) - np.float16(old_area_of_interest_4_gray))
                diff_4 = np.uint8(temp_4)
                temp_5 = np.abs(np.float16(area_of_interest_5_gray) - np.float16(old_area_of_interest_5_gray))
                diff_5 = np.uint8(temp_5)
                temp_6 = np.abs(np.float16(area_of_interest_6_gray) - np.float16(old_area_of_interest_6_gray))
                diff_6 = np.uint8(temp_6)

                #print(current_frame)
                #cv2_imshow(np.concatenate((diff_1, diff_2, diff_3, diff_4, diff_5, diff_6), axis=1))

                if np.sum(diff_1) // 255 > 50:
                    interestCorners[1].append((area_of_interest_1, current_frame))
                if np.sum(diff_2) // 255 > 50:
                    interestCorners[2].append((area_of_interest_2, current_frame))
                if np.sum(diff_3) // 255 > 50:
                    interestCorners[3].append((area_of_interest_3, current_frame))
                if np.sum(diff_4) // 255 > 50:
                    interestCorners[4].append((area_of_interest_4, current_frame))
                if np.sum(diff_5) // 255 > 50:
                    interestCorners[5].append((area_of_interest_5, current_frame))
                if np.sum(diff_6) // 255 > 50:
                    interestCorners[6].append((area_of_interest_6, current_frame))

                old_area_of_interest_1_gray = area_of_interest_1_gray
                old_area_of_interest_2_gray = area_of_interest_2_gray
                old_area_of_interest_3_gray = area_of_interest_3_gray
                old_area_of_interest_4_gray = area_of_interest_4_gray
                old_area_of_interest_5_gray = area_of_interest_5_gray
                old_area_of_interest_6_gray = area_of_interest_6_gray

        else:
            break

    # after playing the video, release the video capture    
    cap.release()
    # close all the frames
    cv.destroyAllWindows()

    low_white = (80, 140, 65)
    high_white = (255, 255, 255)
 
    low_blue = (91, 0, 0)
    high_blue = (180, 135, 55)

    low_dark_green = (37, 60, 0)
    high_dark_green = (60, 255, 30)

    low_yellow = (0, 107, 68)
    high_yellow = (60, 255, 255)

    low_red = (0, 0, 50)
    high_red = (30, 50, 255)

    low_pink = (45, 0, 166)
    high_pink = (255, 162, 255)

    low_brown = (0, 52, 37)
    high_brown = (45, 105, 166)

    low_black = (0, 0, 0)
    high_black = (45, 38, 55)


    if suspicious_color == "white":
        low_color = low_white
        high_color = high_white
    if suspicious_color == "red":
        low_color = low_red
        high_color = high_red
    if suspicious_color == "blue":
        low_color = low_blue
        high_color = high_blue
    if suspicious_color == "pink":
        low_color = low_pink
        high_color = high_pink
    if suspicious_color == "green":
        low_color = low_dark_green
        high_color = high_dark_green
    if suspicious_color == "brown":
        low_color = low_brown
        high_color = high_brown
    if suspicious_color == "yellow":
        low_color = low_yellow
        high_color = high_yellow
    if suspicious_color == "black":
        low_color = low_black
        high_color = high_black

    corners_scores = np.zeros(7)

    for i in range(1, 7):
        sum = 0
        for j in interestCorners[i]:
            score = np.sum(cv.inRange(j[0], low_color, high_color))
            sum += score
        if len(interestCorners[i]) > 0:
            sum = sum / len(interestCorners[i])
        corners_scores[i] = sum
    
    return np.argmax(corners_scores)

#7
def get_numbers_and_colors(imgColor, show_results = True):
    img = cv.cvtColor(imgColor, cv.COLOR_BGR2GRAY)
    
    bottom_left_corner, bottom_right_corner, top_right_corner, top_left_corner = get_corners(img.copy(), False)

    #points_template = np.float32([(2000, 2000), (0, 2000), (0, 0), (2000, 0)]) # the points of the 4 corners (table) in the template image
    #points_query = np.float32([bottom_left_corner, bottom_right_corner, top_right_corner, top_left_corner]) # the points of the 4 corners (table) in the test image
    #homography = cv.getPerspectiveTransform(points_query, points_template)
    #height, width = 2000, 2000 # the shape with respect to the template image
    #aligned_image2 = cv.warpPerspective(imgColor, homography, (width, height), flags=cv.INTER_NEAREST)

    mask = np.zeros(imgColor.shape, dtype=np.uint8)
    
    corners = np.array([[bottom_left_corner, bottom_right_corner, top_right_corner, top_left_corner]], dtype=np.int32)
    channel_count = imgColor.shape[2]  # i.e. 3 or 4 depending on your image
    ignore_mask_color = (255,)*channel_count
    cv.fillPoly(mask, corners, ignore_mask_color)
    # apply the mask
    masked_image = cv.bitwise_and(imgColor, mask)

    masked_image = masked_image[top_right_corner[1]:bottom_left_corner[1], bottom_left_corner[0]:bottom_right_corner[0]]

    ####################

    imgColor = masked_image
    img = cv.cvtColor(imgColor, cv.COLOR_BGR2GRAY)

    w, h = img.shape[::-1]

    thresholds = [0.7, 0.7, 0.7, 0.8, 0.8, 0.8, 0.7]
    locations = []

    for template_numer in range(1, 8):
        template = cv.imread(data_dir_drive + '/data/templates/ball' + str(template_numer) + '.jpg', 0)
        w_T, h_T = template.shape[::-1]

        threshold = thresholds[template_numer-1]
        loc = get_balls_template_matching(img, template, threshold)
        for pt in zip(*loc):
            locations.append((pt[0], pt[1], pt[0] + w_T, pt[1] + h_T))

    boxes = non_max_suppression_fast(np.array(locations), 0.3)

    copyImgColor = imgColor.copy()
    colors = []

    for box in boxes:
        if not false_pozitive(imgColor, box, 0.5):
            color = get_box_color(imgColor, box, show_results)
            colors.append(color)
            if show_results:
                cv.rectangle(copyImgColor, (box[0], box[1]), (box[2], box[3]), (0,0,255), 2) 

    balls, num_balls = process_colors(colors)
    balls["white"] = 1

    if show_results:
        cv2_imshow(copyImgColor)
    
    return num_balls, balls

#8
def cmp_dict(d1, d2):
    if d1["white"] > d2["white"]:
        return "white"
    if d1["red"] > d2["red"]:
        return "red"
    if d1["blue"] > d2["blue"]:
        return "blue"
    if d1["dark_green"] > d2["dark_green"]:
        return "green"
    if d1["brown"] > d2["brown"]:
        return "brown"
    if d1["yellow"] > d2["yellow"]:
        return "yellow"
    if d1["black"] > d2["black"]:
        return "black"
    return True

#9
acc = 0

for i in range(1, 26):
    print(i)
    cap = cv.VideoCapture(folder_in + 'Task2/'+ str(i) +'.mp4')
    length = int(cap.get(cv.CAP_PROP_FRAME_COUNT))
    current_frame = 0

    ret, frame = cap.read() # Read the frame
    num_balls_init, balls_init = get_numbers_and_colors(frame)

    for _ in range(length-10):
        ret, frame = cap.read() # Read the frame
    ret, frame = cap.read() # Read the frame
    num_balls_final, balls_final = get_numbers_and_colors(frame)

    print(balls_init)
    print(balls_final)

    my_ans = ""
    my_color = ""
    my_pocket = 0

    if cmp_dict(balls_init, balls_final) == True:
        my_ans = "NO"
    else:
        my_ans = "YES"
        my_color = cmp_dict(balls_init, balls_final)
        my_pocket = get_suspicious_pockets(folder_in + 'Task2/'+ str(i) +'.mp4', my_color)

    # load ground truth
    ans_truth  = ""
    pocket_truth = 0
    color_truth = ""
    with open(folder_in + 'Task2/ground-truth/' + str(i) + '_gt.txt') as f:
        ans_truth = f.readline().replace('\n', '')
        if ans_truth == "YES":
            pocket_truth = int(f.readline())
            color_truth = f.readline().replace('\n', '')

    print(my_ans, ans_truth, my_ans == ans_truth)
    print(my_pocket, pocket_truth, my_pocket == pocket_truth)
    print(my_color, color_truth, my_color == color_truth)

    if my_ans == ans_truth:
        if ans_truth == "NO":
            acc += 1
        else:
            acc += 0.5
            if my_color == color_truth:
                acc += 0.25
            if my_pocket == pocket_truth:
                acc += 0.25
        
print('acc = ' + str(acc * 4) + '%')

#10
def task2(path):

    cap = cv.VideoCapture(path)
    length = int(cap.get(cv.CAP_PROP_FRAME_COUNT))
    current_frame = 0

    ret, frame = cap.read() # Read the frame
    num_balls_init, balls_init = get_numbers_and_colors(frame, False)

    for _ in range(length-10):
        ret, frame = cap.read() # Read the frame
    ret, frame = cap.read() # Read the frame
    num_balls_final, balls_final = get_numbers_and_colors(frame, False)

    #print(balls_init)
    #print(balls_final)

    my_ans = ""
    my_color = ""
    my_pocket = 0

    if cmp_dict(balls_init, balls_final) == True:
        my_ans = "NO"
    else:
        my_ans = "YES"
        my_color = cmp_dict(balls_init, balls_final)
        my_pocket = get_suspicious_pockets(path, my_color)

    return my_ans, my_color, my_pocket

#11
my_ans, my_color, my_pocket = task2(folder_in + 'Task2/1.mp4')
print(my_ans, my_color, my_pocket)

"""# **Task 3**
You receive a test set of 25 videos containing a player making a shot. In each of the 25 videos the snooker table is seen from above. The task is to track the cue ball (the white ball) and another specified ball. The initial bounding boxes of the two balls to be tracked are provided for the first frame (they follow the format [xmin ymin xmax ymax], where (xmin,ymin) is the top left corner and (xmax,ymax) is the bottom right corner of the initial bounding-box). By correctly solving the tracking problem for a video you will earn 0.04 points (0.02 points for each correctly tracked ball). In each video we will consider that your algorithm correctly tracks a ball if in more (greater or equal) than 80% of the frames your algorithm correctly localizes the ball to be tracked. We consider that your algorithm correctly localizes the ball to be tracked in a specific frame if the value of the IOU (intersection over union) beetween the window provided by your algorithm and the ground-truth window is more than 20%.The first frame for which we provide the bounding boxes initialization has index 0, the last frame of a video with N frames has index N − 1.
"""

#1
def bb_intersection_over_union(boxA, boxB):
    # determine the (x, y)-coordinates of the intersection rectangle
    xA = max(boxA[0], boxB[0])
    yA = max(boxA[1], boxB[1])
    xB = min(boxA[2], boxB[2])
    yB = min(boxA[3], boxB[3])
    # compute the area of intersection rectangle
    interArea = max(0, xB - xA + 1) * max(0, yB - yA + 1)
    # compute the area of both the prediction and ground-truth
    # rectangles
    boxAArea = (boxA[2] - boxA[0] + 1) * (boxA[3] - boxA[1] + 1)
    boxBArea = (boxB[2] - boxB[0] + 1) * (boxB[3] - boxB[1] + 1)
    # compute the intersection over union by taking the intersection
    # area and dividing it by the sum of prediction + ground-truth
    # areas - the intersection area
    if float(boxAArea + boxBArea - interArea) > 0:
        iou = interArea / float(boxAArea + boxBArea - interArea)
    else: 
        iou = 1
    # return the intersection over union value
    return iou

#2
def compute_percentage_tracking(gt_bboxes, predicted_bboxes, num_frames):
    num_frames = int(num_frames)
    
    tp = 0
    fp = 0
    
    gt_dict = {}
    for gt_box in gt_bboxes:
        gt_dict[gt_box[0]] = gt_box[1:]
    
    pred_dict = {}
    for pred_bbox in predicted_bboxes:
        pred_dict[pred_bbox[0]] = pred_bbox[1:]
        
    for i in range(num_frames):
        if gt_dict.get(i, None) is None and pred_dict.get(i, None) is None: # the ball is not on the table
            tp += 1 
        
        elif gt_dict.get(i, None) is not None and pred_dict.get(i, None) is None: # the ball is not detected
            fp += 1
            
        elif gt_dict.get(i, None) is None and pred_dict.get(i, None) is not None: # the ball is not on the table, but it is 'detected'
            fp += 1
            
        elif gt_dict.get(i, None) is not None and pred_dict.get(i, None) is not None: # the ball is on the table and it is detected
            
            iou = bb_intersection_over_union(gt_dict[i], pred_dict[i])
            if iou >= 0.2:
                tp += 1
            else:
                fp += 1 
             
            
    print(tp, fp)
    assert tp + fp == num_frames
    perc = tp / (tp + fp)
    
    return perc

#3
def track_ball_using_hist_of_colors(video_path, first_box, max_number_of_frame_to_run, show_results = True):
    
    bboxes = []
    
    cap = cv.VideoCapture(video_path)
    ret, first_frame = cap.read() # Read the first frame 
    
    (x, y, X, Y) = first_box
    w = X-x
    h = Y-y
    track_window = (x, y, w, h)
    
    roi = first_frame[y: y + h, x: x + w]
    annotated_frame = cv.rectangle(first_frame, (x, y), (x+w,y+h), 255, 2)
 
    if show_results:
        cv2_imshow(annotated_frame)
    
    roi_hist = cv.calcHist([roi], [0 ,1, 2], None, [4, 4, 4], [0, 256, 0, 256, 0, 256]) 
    roi_hist_norm = roi_hist / roi_hist.sum()

    roi_gray = cv.cvtColor(roi, cv.COLOR_BGR2GRAY)
    
    frame_idx = 0
    while cap.isOpened():
        frame_idx += 1
        ret, frame = cap.read()

        if ret is True: 
            frame_gray = cv.cvtColor(frame, cv.COLOR_BGR2GRAY)
            mask1 = np.int8(np.zeros(frame_gray.shape))
            center = (y + h//2, x + h//2)

            y_min = np.max((0, center[0] - (3*h)))
            y_max = np.min((frame.shape[0], center[0] + (3*h)))
            x_min = np.max((0, center[1] - (3*w)))
            x_max = np.min((frame.shape[1], center[1] + (3*w)))
            
            mask1[y_min: y_max, x_min: x_max] = 255

            frame_gray_mask = cv.bitwise_and(frame_gray,frame_gray,mask=mask1)

            res = cv.matchTemplate(frame_gray_mask, roi_gray, cv.TM_CCOEFF_NORMED)        
            min_val, max_val, min_loc, max_loc = cv.minMaxLoc(res)

            y = max_loc[1]
            x = max_loc[0]
            bboxes.append([frame_idx, x, y, x + w, y + h])
            img2 = cv.rectangle(frame, (x, y), (x + w, y + h), 255, 2)
            if show_results:
                cv2_imshow(img2)
            
            if current_frame > max_number_of_frame_to_run:
                break

        else:
            break
    # after playing the video, release the video capture    
    cap.release()
    # close all the frames
    cv.destroyAllWindows()
    return bboxes

#4
def track_ball_using_CSRT_tracker(video_path, first_box, max_number_of_frame_to_run, show_results = True):
    
    bboxes = []
    
    cap = cv.VideoCapture(video_path)
    ret, first_frame = cap.read() # Read the first frame 
    
    (x, y, X, Y) = first_box
    w = X-x
    h = Y-y
    track_window = (x-10, y-10, w+20, h+20)
    
    roi = first_frame[y: y + h, x: x + w]
    annotated_frame = cv.rectangle(first_frame, (x, y), (x+w,y+h), 255, 2)
 
    if show_results:
        cv2_imshow(annotated_frame)
    
    tracker = cv.TrackerCSRT_create()
    tracker.init(first_frame, track_window)
    frame_idx = 0

    while cap.isOpened():
        frame_idx += 1
        ret, frame = cap.read()

        if ret is True:             
            (success, box) = tracker.update(frame)

            if success:
                (x, y, w, h) = [int(p) for p in box]
                x = x + 10
                y = y + 10
                w = w - 20
                h = h - 20
                cv.rectangle(frame, (x, y), (x + w, y + h),
                    (0, 255, 0), 2)

            bboxes.append([frame_idx, x, y, x + w, y + h])
            img2 = cv.rectangle(frame, (x, y), (x + w, y + h), 255, 2)
            if show_results:
                cv2_imshow(img2)

            if frame_idx > max_number_of_frame_to_run:
                break
                
        else:
            break
    # after playing the video, release the video capture    
    cap.release()
    # close all the frames
    cv.destroyAllWindows()
    return bboxes

#5
idx = 3
video_name = folder_in + "Task3/" + str(idx) + ".mp4"

first_frame1 = None
first_frame2 = None
max_number_of_frame_to_run1 = 0
max_number_of_frame_to_run2 = 0

with open(folder_in + "Task3/" + str(idx) + "_ball_1.txt") as f:
    max_number_of_frame_to_run1 = int(f.readline().split(' ')[0])
    ans_truth = f.readline().replace('\n', '').split(' ')
    first_frame1 = (int(ans_truth[1]), int(ans_truth[2]), int(ans_truth[3]), int(ans_truth[4]))

with open(folder_in + "Task3/" + str(idx) + "_ball_2.txt") as f:
    max_number_of_frame_to_run2 = int(f.readline().split(' ')[0])
    ans_truth = f.readline().replace('\n', '').split(' ')
    first_frame2 = (int(ans_truth[1]), int(ans_truth[2]), int(ans_truth[3]), int(ans_truth[4]))

bboxes1 = track_ball_using_CSRT_tracker(video_name, first_frame1, max_number_of_frame_to_run1, False)
bboxes2 = track_ball_using_CSRT_tracker(video_name, first_frame2, max_number_of_frame_to_run2, False)

#bboxes3 = track_ball_using_hist_of_colors(video_name, first_frame1, max_number_of_frame_to_run1, False)
#bboxes4 = track_ball_using_hist_of_colors(video_name, first_frame2, max_number_of_frame_to_run2, False)

#6
def get_coordinates_of_interet(frame):
    imgColor = frame.copy()
    img = cv.cvtColor(imgColor, cv.COLOR_BGR2GRAY)
            
    bottom_left_corner, bottom_right_corner, top_right_corner, top_left_corner, middle_left_corner, middle_right_corner = get_pockets(img.copy(), False)
    if bottom_left_corner is None or bottom_right_corner is None or top_right_corner is None or top_left_corner is None or middle_right_corner is None or middle_left_corner is None:
        return None, None, None, None, None, None

    interestRaze = 75

    area_of_interest_1 = [top_left_corner[1] - interestRaze, top_left_corner[1] + interestRaze, top_left_corner[0] - interestRaze, top_left_corner[0] + interestRaze]
    area_of_interest_2 = [top_right_corner[1] - interestRaze, top_right_corner[1] + interestRaze, top_right_corner[0] - interestRaze, top_right_corner[0] + interestRaze]
    area_of_interest_3 = [bottom_left_corner[1] - interestRaze, bottom_left_corner[1] + interestRaze, bottom_left_corner[0] - interestRaze, bottom_left_corner[0] + interestRaze]
    area_of_interest_4 = [bottom_right_corner[1] - interestRaze, bottom_right_corner[1] + interestRaze, bottom_right_corner[0] - interestRaze, bottom_right_corner[0] + interestRaze]
    area_of_interest_5 = [middle_left_corner[1] - interestRaze, middle_left_corner[1] + interestRaze, middle_left_corner[0] - interestRaze, middle_left_corner[0] + interestRaze]
    area_of_interest_6 = [middle_right_corner[1] - interestRaze, middle_right_corner[1] + interestRaze, middle_right_corner[0] - interestRaze, middle_right_corner[0] + interestRaze]

    return area_of_interest_1, area_of_interest_2, area_of_interest_3, area_of_interest_4, area_of_interest_5, area_of_interest_6

#7
def update_boxes(video_path, boxes, show_results = True):
    cap = cv.VideoCapture(video_path)  
    current_frame = 0

    ret, frame = cap.read() # Read the frame
    area_of_interest_1, area_of_interest_2, area_of_interest_3, area_of_interest_4, area_of_interest_5, area_of_interest_6 = get_coordinates_of_interet(frame)
    while area_of_interest_1 is None:
        current_frame = current_frame + 1
        ret, frame = cap.read() # Read the frame
        area_of_interest_1, area_of_interest_2, area_of_interest_3, area_of_interest_4, area_of_interest_5, area_of_interest_6 = get_coordinates_of_interet(frame)

    is_moving = False
    frame_start_moving = 0
    for i in range(5, len(boxes), 5):
        if bb_intersection_over_union(boxes[i-5][1:], boxes[i][1:]) < 0.5:
            is_moving = True
            frame_start_moving = i -5
            break

    if show_results:
        print("frame_start_moving = ", frame_start_moving)

    if is_moving:
        frame_end_moving = len(boxes) - 1
        for i in range(frame_start_moving + 5, len(boxes), 5):
            if bb_intersection_over_union(boxes[i-5][1:], boxes[i][1:]) > 0.9:
                frame_end_moving = i
                break

        if show_results:
            print("frame_end_moving = ", frame_end_moving)
            print(area_of_interest_3)
            print(boxes[frame_end_moving][1:])

            img2 = cv.rectangle(frame, (area_of_interest_3[2], area_of_interest_3[0]), (area_of_interest_3[3], area_of_interest_3[1]), (255,0,0), 2)
            img2 = cv.rectangle(frame, (boxes[frame_start_moving][1], boxes[frame_start_moving][2]), (boxes[frame_start_moving][3], boxes[frame_start_moving][4]), (0,255,0), 2)
            img2 = cv.rectangle(frame, (boxes[frame_end_moving][1], boxes[frame_end_moving][2]), (boxes[frame_end_moving][3], boxes[frame_end_moving][4]), (0,0,255), 2)
            cv2_imshow(img2)
        
        result = boxes
        if area_of_interest_1[0] <= boxes[frame_end_moving][2] <= area_of_interest_1[1] and area_of_interest_1[0] <= boxes[frame_end_moving][4] <= area_of_interest_1[1] and area_of_interest_1[2] <= boxes[frame_end_moving][1] <= area_of_interest_1[3] and area_of_interest_1[2] <= boxes[frame_end_moving][3] <= area_of_interest_1[3]:
            result = boxes[:frame_end_moving]
            print("case1")
        elif area_of_interest_2[0] <= boxes[frame_end_moving][2] <= area_of_interest_2[1] and area_of_interest_2[0] <= boxes[frame_end_moving][4] <= area_of_interest_2[1] and area_of_interest_2[2] <= boxes[frame_end_moving][1] <= area_of_interest_2[3] and area_of_interest_2[2] <= boxes[frame_end_moving][3] <= area_of_interest_2[3]:
            result = boxes[:frame_end_moving]
            print("case2")
        elif area_of_interest_3[0] <= boxes[frame_end_moving][2] <= area_of_interest_3[1] and area_of_interest_3[0] <= boxes[frame_end_moving][4] <= area_of_interest_3[1] and area_of_interest_3[2] <= boxes[frame_end_moving][1] <= area_of_interest_3[3] and area_of_interest_3[2] <= boxes[frame_end_moving][3] <= area_of_interest_3[3]:
            result = boxes[:frame_end_moving]
            print("case3")
        elif area_of_interest_4[0] <= boxes[frame_end_moving][2] <= area_of_interest_4[1] and area_of_interest_4[0] <= boxes[frame_end_moving][4] <= area_of_interest_4[1] and area_of_interest_4[2] <= boxes[frame_end_moving][1] <= area_of_interest_4[3] and area_of_interest_4[2] <= boxes[frame_end_moving][3] <= area_of_interest_4[3]:
            result = boxes[:frame_end_moving]
            print("case4")
        elif area_of_interest_5[0] <= boxes[frame_end_moving][2] <= area_of_interest_5[1] and area_of_interest_5[0] <= boxes[frame_end_moving][4] <= area_of_interest_5[1] and area_of_interest_1[2] <= boxes[frame_end_moving][1] <= area_of_interest_5[3] and area_of_interest_5[2] <= boxes[frame_end_moving][3] <= area_of_interest_5[3]:
            result = boxes[:frame_end_moving]
            print("case5")
        elif area_of_interest_6[0] <= boxes[frame_end_moving][2] <= area_of_interest_6[1] and area_of_interest_6[0] <= boxes[frame_end_moving][4] <= area_of_interest_6[1] and area_of_interest_1[2] <= boxes[frame_end_moving][1] <= area_of_interest_6[3] and area_of_interest_6[2] <= boxes[frame_end_moving][3] <= area_of_interest_6[3]:
            result = boxes[:frame_end_moving]
            print("case6")

        return result
    else:
        return boxes

#8
bboxes1_new = update_boxes(video_name, bboxes1, False)
bboxes2_new = update_boxes(video_name, bboxes2, False)

#9
# load the ground-truth file
ball1_gt = np.loadtxt(folder_in + "Task3/ground-truth/" + str(idx) + "_ball_1_gt.txt")
ball2_gt = np.loadtxt(folder_in + "Task3/ground-truth/" + str(idx) + "_ball_2_gt.txt")

#10
compute_percentage_tracking(ball1_gt[1:], bboxes1_new, ball1_gt[0][0])

#11
compute_percentage_tracking(ball2_gt[1:], bboxes2_new, ball2_gt[0][0])

#12
acc = 0
for idx in range (1, 26):

    video_name = folder_in + "Task3/" + str(idx) + ".mp4"

    first_frame1 = None
    first_frame2 = None
    max_number_of_frame_to_run1 = 0
    max_number_of_frame_to_run2 = 0

    with open(folder_in + "Task3/" + str(idx) + "_ball_1.txt") as f:
        max_number_of_frame_to_run1 = int(f.readline().split(' ')[0])
        ans_truth = f.readline().replace('\n', '').split(' ')
        first_frame1 = (int(ans_truth[1]), int(ans_truth[2]), int(ans_truth[3]), int(ans_truth[4]))

    with open(folder_in + "Task3/" + str(idx) + "_ball_2.txt") as f:
        max_number_of_frame_to_run2 = int(f.readline().split(' ')[0])
        ans_truth = f.readline().replace('\n', '').split(' ')
        first_frame2 = (int(ans_truth[1]), int(ans_truth[2]), int(ans_truth[3]), int(ans_truth[4]))

    bboxes1 = track_ball_using_CSRT_tracker(video_name, first_frame1, max_number_of_frame_to_run1, False)
    bboxes2 = track_ball_using_CSRT_tracker(video_name, first_frame2, max_number_of_frame_to_run2, False)

    #bboxes3 = track_ball_using_hist_of_colors(video_name, first_frame1, max_number_of_frame_to_run1, False)
    #bboxes4 = track_ball_using_hist_of_colors(video_name, first_frame2, max_number_of_frame_to_run2, False)

    bboxes1_new = update_boxes(video_name, bboxes1, False)
    bboxes2_new = update_boxes(video_name, bboxes2, False)

    # load the ground-truth file
    ball1_gt = np.loadtxt(folder_in + "Task3/ground-truth/" + str(idx) + "_ball_1_gt.txt")
    ball2_gt = np.loadtxt(folder_in + "Task3/ground-truth/" + str(idx) + "_ball_2_gt.txt")

    acc1 = compute_percentage_tracking(ball1_gt[1:], bboxes1_new, ball1_gt[0][0])
    acc2 = compute_percentage_tracking(ball2_gt[1:], bboxes2_new, ball2_gt[0][0])
    print(idx, acc1, acc2)
    if acc1 > 0.8:
        acc += 0.5
    if acc2 > 0.8:
        acc += 0.5

print("acc=", acc * 4, "%")

#13
def task3(video_name, init_ball1, init_ball2):

    first_frame1 = None
    first_frame2 = None
    max_number_of_frame_to_run1 = 0
    max_number_of_frame_to_run2 = 0

    with open(init_ball1) as f:
        max_number_of_frame_to_run1 = int(f.readline().split(' ')[0])
        ans_truth = f.readline().replace('\n', '').split(' ')
        first_frame1 = (int(ans_truth[1]), int(ans_truth[2]), int(ans_truth[3]), int(ans_truth[4]))

    with open(init_ball2) as f:
        max_number_of_frame_to_run2 = int(f.readline().split(' ')[0])
        ans_truth = f.readline().replace('\n', '').split(' ')
        first_frame2 = (int(ans_truth[1]), int(ans_truth[2]), int(ans_truth[3]), int(ans_truth[4]))

    bboxes1 = track_ball_using_CSRT_tracker(video_name, first_frame1, max_number_of_frame_to_run1, False)
    bboxes2 = track_ball_using_CSRT_tracker(video_name, first_frame2, max_number_of_frame_to_run2, False)

    #bboxes3 = track_ball_using_hist_of_colors(video_name, first_frame1, max_number_of_frame_to_run1, False)
    #bboxes4 = track_ball_using_hist_of_colors(video_name, first_frame2, max_number_of_frame_to_run2, False)

    bboxes1_new = update_boxes(video_name, bboxes1, False)
    bboxes2_new = update_boxes(video_name, bboxes2, False)

    return bboxes1_new, bboxes2_new

#14
bboxes1, bboxes2 = task3(folder_in + 'Task3/1.mp4', folder_in + "Task3/1_ball_1.txt", folder_in + "Task3/1_ball_2.txt")
print(bboxes1, bboxes2)

"""# **(bonus) Task 4**
You receive a test set of 25 videos containing a player making a shot. Different than the previous tasks, in this videos the snooker table can be filmed from different viewpoints (not only from above). The task is to track the cue ball. The task here is much harder as you have diferent scales of the cue ball, changes in camera viewpoint, the initial bounding box of the white ball is not provided. The initial bounding box of the cue ball to be tracked is provided for the first frame (it follows the format [xmin ymin xmax ymax], where (xmin,ymin) is the top left corner and (xmax,ymax) is the bottom right corner of the initial bounding-box). By correctly solving the tracking problem for a video you will earn 0.04 points. The rules described at Task 3 apply here, meaning that your algorithm should correctly localizes the cue ball in more than 80% of the frames, at each frame correctly localization means that the window provided by your algorithm should have IOU more than 20% with respect to the ground-truth window. The first frame for which we provide the bounding boxes initialization has index 0, the last frame of a video with N frames has index N − 1.
"""

#1

"""# **Running the entire code**"""

#1
def run_task_1(have_ground_truth = False):
    acc_colors = 0
    acc_numbers = 0

    for i in range(1, 51):
        # change to test folder
        imgColor = cv.imread(folder_in + 'Task1/'+ str(i) +'.jpg')
        
        balls, num_balls = task1(imgColor)

        with open(folder_out + 'Task1/' + str(i) + '.txt', 'w') as f:
            f.write(str(num_balls) + '\n')
            for color, nr in balls.items(): 
                f.write(str(nr) + " " + color + '\n')
        
        
        if have_ground_truth:
            # load ground truth
            ground_truth = {}
            num_balls_truth  = 0
            with open(folder_in + 'Task1/ground-truth/' + str(i) + '_gt.txt') as f:
                num_balls_truth = int(f.readline())
                lines = f.readlines()
                for line in lines:
                    info = line.replace('\n', '').split(' ')
                    nr = int(info[0])
                    color = info[1]

                    ground_truth[color] = nr

            ok = True
            nr = 0
            for color, number in balls.items(): 
                nr += number
                if color == 'dark_green':
                    if ground_truth['green'] != number:
                        ok = False
                else:
                    if ground_truth[color] != number:
                        ok = False
                        print("greseala - ground_truth:", ground_truth[color], "me: ", number, color)

            if ok:
                acc_colors += 1
            else:
                print(i)
                print(num_balls_truth, num_balls)     
                print(ground_truth)
                print(balls)

            if num_balls_truth == nr:
                acc_numbers += 1

                ####################
            #cv.imwrite(data_dir_drive + '/data/aux/' + str(i) + '.jpg', masked_image)

    if have_ground_truth:
        print('acc numbers = ' + str(acc_numbers * 2) + '%')
        print('acc colors = ' + str(acc_colors * 2) + '%')

#2
run_task_1(True)

#3
def run_task_2(have_ground_truth = False):

    acc = 0

    for i in range(1, 26):
        print(i)
        
        my_ans, my_color, my_pocket = task2(folder_in + 'Task2/'+ str(i) +'.mp4')

        with open(folder_out + 'Task2/' + str(i) + '.txt', 'w') as f:
            f.write(my_ans + '\n')
            if my_ans == 'YES':
                f.write(str(my_pocket) + '\n' + my_color)

        if have_ground_truth:
            # load ground truth
            ans_truth  = ""
            pocket_truth = 0
            color_truth = ""
            with open(folder_in + 'Task2/ground-truth/' + str(i) + '_gt.txt') as f:
                ans_truth = f.readline().replace('\n', '')
                if ans_truth == "YES":
                    pocket_truth = int(f.readline())
                    color_truth = f.readline().replace('\n', '')

            print(my_ans, ans_truth, my_ans == ans_truth)
            print(my_pocket, pocket_truth, my_pocket == pocket_truth)
            print(my_color, color_truth, my_color == color_truth)

            if my_ans == ans_truth:
                if ans_truth == "NO":
                    acc += 1
                else:
                    acc += 0.5
                    if my_color == color_truth:
                        acc += 0.25
                    if my_pocket == pocket_truth:
                        acc += 0.25

    if have_ground_truth:        
        print('acc = ' + str(acc * 4) + '%')

#4
run_task_2(True)

#5
def run_task_3(have_ground_truth = False):
    acc = 0
    for idx in range (1, 26):

        video_name = folder_in + "Task3/" + str(idx) + ".mp4"
        txt_ball_1 = folder_in + "Task3/" + str(idx) + "_ball_1.txt"
        txt_ball_2 = folder_in + "Task3/" + str(idx) + "_ball_2.txt"

        bboxes1, bboxes2 = task3(video_name, txt_ball_1, txt_ball_2)

        info_ball_1 = ''
        with open(txt_ball_1) as f:
            info_ball_1 = f.readlines()
        with open(folder_out + 'Task3/' + str(idx) + '_ball_1.txt', 'w') as f:
            f.write(info_ball_1[0] + info_ball_1[1] + '\n')
            for box in bboxes1:
                f.write(str(box[0]) + ' ' + str(box[1]) + ' ' + str(box[2]) + ' ' + str(box[3]) + ' ' + str(box[4]) + '\n')

        info_ball_2 = ''
        with open(txt_ball_2) as f:
            info_ball_2 = f.readlines()
        with open(folder_out + 'Task3/' + str(idx) + '_ball_2.txt', 'w') as f:
            f.write(info_ball_2[0] + info_ball_2[1] + '\n')
            for box in bboxes2:
                f.write(str(box[0]) + ' ' + str(box[1]) + ' ' + str(box[2]) + ' ' + str(box[3]) + ' ' + str(box[4]) + '\n')

        if have_ground_truth:
            # load the ground-truth file
            ball1_gt = np.loadtxt(folder_in + "Task3/ground-truth/" + str(idx) + "_ball_1_gt.txt")
            ball2_gt = np.loadtxt(folder_in + "Task3/ground-truth/" + str(idx) + "_ball_2_gt.txt")

            acc1 = compute_percentage_tracking(ball1_gt[1:], bboxes1, ball1_gt[0][0])
            acc2 = compute_percentage_tracking(ball2_gt[1:], bboxes2, ball2_gt[0][0])
            print(idx, acc1, acc2)
            if acc1 > 0.8:
                acc += 0.5
            if acc2 > 0.8:
                acc += 0.5

    if have_ground_truth:
        print("acc=", acc * 4, "%")

#6
run_task_3(True)