import cv2
import pytesseract
import numpy as np
​
class ResultBoxDimensions:
    def __init__(self, playerCount, anchor, width, padding, margin):
        self.playerCount = playerCount
        self.anchor  = anchor
        self.width = width
        self.padding = padding
        self.margin = margin

dimensions = []
dimensions.append(ResultBoxDimensions(2, (279, 532), 572, 17, 184))
dimensions.append(ResultBoxDimensions(3, ( 46, 532), 572, 17,  22))
dimensions.append(ResultBoxDimensions(4, ( 17, 533), 446, 17,   0))
dimensions.append(ResultBoxDimensions(5, (126, 538), 264, 13,  60))
dimensions.append(ResultBoxDimensions(6, ( 38, 538), 264, 13,  26))
dimensions.append(ResultBoxDimensions(7, ( 34, 538), 220, 10,  32))
dimensions.append(ResultBoxDimensions(8, ( 10, 538), 220, 10,   0))
dimensions.sort(key=lambda x: x.anchor[0])
​
#Methods
def RotateImage(image, angle):
    image_center = tuple(np.array(image.shape[1::-1]) / 2)
    rot_mat = cv2.getRotationMatrix2D(image_center, angle, 1.0)
    result = cv2.warpAffine(image, rot_mat, image.shape[1::-1], flags=cv2.INTER_LINEAR)
    return result
​
def GetPlayersFromResultScreen(image,
                   threshold = 3,
                   pointChecks = 3,
                   topIdealColor = 28,
                   bottomIdealColor = 72,
                   drawCheckPoints = True):
    for d in dimensions:
        for p in range(d.playerCount):
            isResultScreen = True
            for c in range(pointChecks):
                koTopColor = image[d.anchor[1], d.anchor[0] + (d.width/pointChecks)*c + (d.width + 2*d.padding + d.margin)*p]
                koBottomColor = image[d.anchor[1] + 42, d.anchor[0] + (d.width/pointChecks)*c + (d.width + 2*d.padding + d.margin)*p]
                fallsTopColor = image[d.anchor[1] + 112 , d.anchor[0] + (d.width/pointChecks)*c + (d.width + 2*d.padding + d.margin)*p]
                fallsBottomColor = image[d.anchor[1] + 112 + 42, d.anchor[0] + (d.width/pointChecks)*c + (d.width + 2*d.padding + d.margin)*p]
                if ( abs(topIdealColor - koTopColor) <= threshold and
                     abs(topIdealColor - fallsTopColor) <= threshold and
                     abs(bottomIdealColor - koBottomColor) <= threshold and
                     abs(bottomIdealColor - fallsBottomColor) <= threshold ):
                    continue
                isResultScreen = False
                break
            if(isResultScreen):
                for count in range(pointChecks):
                    radius = 5
                    color = (0, 0, 255)
                    thickness = 1
                    cv2.circle(image, ((d.anchor[0] + ((d.width-2)/pointChecks)*count + (d.width + 2*d.padding + d.margin)*p), d.anchor[1]), radius, color, thickness)
                    cv2.circle(image, ((d.anchor[0] + ((d.width-2)/pointChecks)*count + (d.width + 2*d.padding + d.margin)*p), d.anchor[1] + 42), radius, color, thickness)
                    cv2.circle(image, ((d.anchor[0] + ((d.width-2)/pointChecks)*count + (d.width + 2*d.padding + d.margin)*p), d.anchor[1] + 112), radius, color, thickness)
                    cv2.circle(image, ((d.anchor[0] + ((d.width-2)/pointChecks)*count + (d.width + 2*d.padding + d.margin)*p), d.anchor[1] + 112 + 42), radius, color, thickness)
                return (p+1)
    return 0
​
pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files (x86)\Tesseract-OCR\tesseract'
cap = cv2.VideoCapture(1) #cv2.VideoCapture("C:/tools/SmashParser/SmashVideo_Trim.mp4")
​
i=0
while(True):
    ret, frame = cap.read()
    if ret == False or (not cap.isOpened()):
        break
​
    i+=1
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

    # add frame count
    cv2.putText(gray, str(i), (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 0), 2, cv2.LINE_AA)
​
    # export frame if it's result screen
    playerCount = GetPlayersFromResultScreen(gray)
    if (playerCount > 0):
        print("RESULT SCREEN: " + str(playerCount))
        #cv2.imwrite(('C:/tools/SmashParser/Images/'+str(i)+'.png'),gray)
​
    cv2.imshow('frame', gray)
​
    # get player names and print them to stdout
    if cv2.waitKey(1) & 0xFF == ord('s'):
        cv2.imwrite(('C:/tools/SmashParser/Images/'+str(i)+'_'+str(0)+'.png'),gray)
        #for x in range(8):
            #crop_img = gray[0:250, (240*x):(240*x)+240]
            #ret, crop_img = cv2.threshold(crop_img,210,255,cv2.THRESH_BINARY)
            #crop_img = RotateImage(crop_img, -5.1)
            #print(pytesseract.image_to_string(crop_img))
            #cv2.imwrite(('C:/tools/SmashParser/Images/'+str(i)+'_'+str(x)+'.png'),crop_img)
​
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break
cap.release()
cv2.destroyAllWindows()