SpeedProg
2020-01-02 8468f19b741475059bbd0bc4db310ee05ea7e6fb
opencv_dnn.py
old mode 100644 new mode 100755
@@ -1,3 +1,4 @@
import argparse
import ast
import collections
import cv2
@@ -8,9 +9,9 @@
import pandas as pd
from PIL import Image
import time
from multiprocessing import Pool
from config import Config
import fetch_data
import transform_data
"""
@@ -21,20 +22,14 @@
https://github.com/hj3yoo/mtg_card_detector/tree/dea64611730c84a59c711c61f7f80948f82bcd31 
"""
def calc_image_hashes(card_pool, save_to=None, hash_size=32, highfreq_factor=4):
    """
    Calculate perceptual hash (pHash) value for each cards in the database, then store them if needed
    :param card_pool: pandas dataframe containing all card information
    :param save_to: path for the pickle file to be saved
    :param hash_size: param for pHash algorithm
    :param highfreq_factor: param for pHash algorithm
    :return: pandas dataframe
    """
    # Since some double-faced cards may result in two different cards, create a new dataframe to store the result
def do_calc(args):
    card_pool = args[0]
    hash_size = args[1]
    new_pool = pd.DataFrame(columns=list(card_pool.columns.values))
    new_pool['card_hash'] = np.NaN
    #new_pool['art_hash'] = np.NaN
    for hs in hash_size:
        new_pool['card_hash_%d' % hs] = np.NaN
        new_pool['set_hash_%d' % 64] = np.NaN
        #new_pool['art_hash_%d' % hs] = np.NaN
    for ind, card_info in card_pool.iterrows():
        if ind % 100 == 0:
            print('Calculating hashes: %dth card' % ind)
@@ -54,34 +49,82 @@
        for card_name in card_names:
            # Fetch the image - name can be found based on the card's information
            card_info['name'] = card_name
            img_name = '%s/card_img/png/%s/%s_%s.png' % (transform_data.data_dir, card_info['set'],
            cname = card_name
            if cname == 'con':
                cname == 'con__'
            img_name = '%s/card_img/png/%s/%s_%s.png' % (Config.data_dir, card_info['set'],
                                                         card_info['collector_number'],
                                                         fetch_data.get_valid_filename(card_info['name']))
                                                         fetch_data.get_valid_filename(cname))
            card_img = cv2.imread(img_name)
            # If the image doesn't exist, download it from the URL
            if card_img is None:
                set_name = card_info['set']
                if set_name == 'con':
                    set_name = 'con__'
                fetch_data.fetch_card_image(card_info,
                                            out_dir='%s/card_img/png/%s' % (transform_data.data_dir, card_info['set']))
                                            out_dir='%s/card_img/png/%s' % (Config.data_dir, set_name))
                card_img = cv2.imread(img_name)
            if card_img is None:
                print('WARNING: card %s is not found!' % img_name)
                continue
            """
            img_cc = cv2.cvtColor(card_img, cv2.COLOR_BGR2GRAY)
            img_thresh = cv2.adaptiveThreshold(img_cc, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY_INV, 11, 5)
            # Dilute the image, then erode them to remove minor noises
            kernel = np.ones((3, 3), np.uint8)
            img_dilate = cv2.dilate(img_thresh, kernel, iterations=1)
            img_erode = cv2.erode(img_dilate, kernel, iterations=1)
            cnts, hier = cv2.findContours(img_erode, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
            cnts2 = sorted(cnts, key=cv2.contourArea, reverse=True)
            cnts2 = cnts2[:10]
            if True:
                cv2.drawContours(img_cc, cnts2, -1, (0, 255, 0), 3)
                #cv2.imshow('Contours', card_img)
                #cv2.waitKey(10000)
            """
            set_img = card_img[595:635, 600:690]
            #cv2.imshow(card_info['name'], set_img)
            # Compute value of the card's perceptual hash, then store it to the database
            '''
            img_art = Image.fromarray(card_img[121:580, 63:685])  # For 745*1040 size card image
            art_hash = ih.phash(img_art, hash_size=hash_size, highfreq_factor=highfreq_factor)
            card_info['art_hash'] = art_hash
            '''
            #img_art = Image.fromarray(card_img[121:580, 63:685])  # For 745*1040 size card image
            img_card = Image.fromarray(card_img)
            card_hash = ih.phash(img_card, hash_size=hash_size, highfreq_factor=highfreq_factor)
            card_info['card_hash'] = card_hash
            img_set = Image.fromarray(set_img)
            #cv2.imshow('Set' + card_names[0], set_img)
            for hs in hash_size:
                card_hash = ih.phash(img_card, hash_size=hs)
                set_hash = ih.whash(img_set, hash_size=64)
                card_info['card_hash_%d' % hs] = card_hash
                card_info['set_hash_%d' % 64] = set_hash
                #print('Setting set_hash_%d' % hs)
                #art_hash = ih.phash(img_art, hash_size=hs)
                #card_info['art_hash_%d' % hs] = art_hash
            new_pool.loc[0 if new_pool.empty else new_pool.index.max() + 1] = card_info
    return new_pool
    # Remove uselesss fields, then pickle it if needed
    new_pool = new_pool[['artist', 'border_color', 'collector_number', 'color_identity', 'colors', 'flavor_text',
                         'image_uris', 'mana_cost', 'legalities', 'name', 'oracle_text', 'rarity', 'type_line',
                         'set', 'set_name', 'power', 'toughness', 'art_hash', 'card_hash']]
def calc_image_hashes(card_pool, save_to=None, hash_size=None):
    """
    Calculate perceptual hash (pHash) value for each cards in the database, then store them if needed
    :param card_pool: pandas dataframe containing all card information
    :param save_to: path for the pickle file to be saved
    :param hash_size: param for pHash algorithm
    :return: pandas dataframe
    """
    if hash_size is None:
        hash_size = [16, 32]
    elif isinstance(hash_size, int):
        hash_size = [hash_size]
    num_cores = 16
    num_partitions = round(card_pool.shape[0]/1000)
    if num_partitions < min(num_cores, card_pool.shape[0]):
        num_partitions = min(num_cores, card_pool.shape[0])
    pool = Pool(num_cores)
    df_split = np.array_split(card_pool, num_partitions)
    new_pool = pd.concat(pool.map(do_calc, [(split, hash_size) for split in df_split]))
    pool.close()
    pool.join()
    # Since some double-faced cards may result in two different cards, create a new dataframe to store the result
    if save_to is not None:
        new_pool.to_pickle(save_to)
    return new_pool
@@ -166,72 +209,6 @@
    return warped
'''
# The following functions are only used in conjunction with YOLO, and is deprecated:
# - get_outputs_names()
# - post_process()
# - draw_pred()
# Get the names of the output layers
def get_outputs_names(net):
    # Get the names of all the layers in the network
    layers_names = net.getLayerNames()
    # Get the names of the output layers, i.e. the layers with unconnected outputs
    return [layers_names[i[0] - 1] for i in net.getUnconnectedOutLayers()]
# Remove the bounding boxes with low confidence using non-maxima suppression
# https://www.learnopencv.com/deep-learning-based-object-detection-using-yolov3-with-opencv-python-c/
def post_process(frame, outs, thresh_conf, thresh_nms):
    frame_height = frame.shape[0]
    frame_width = frame.shape[1]
    # Scan through all the bounding boxes output from the network and keep only the
    # ones with high confidence scores. Assign the box's class label as the class with the highest score.
    class_ids = []
    confidences = []
    boxes = []
    for out in outs:
        for detection in out:
            scores = detection[5:]
            class_id = np.argmax(scores)
            confidence = scores[class_id]
            if confidence > thresh_conf:
                center_x = int(detection[0] * frame_width)
                center_y = int(detection[1] * frame_height)
                width = int(detection[2] * frame_width)
                height = int(detection[3] * frame_height)
                left = int(center_x - width / 2)
                top = int(center_y - height / 2)
                class_ids.append(class_id)
                confidences.append(float(confidence))
                boxes.append([left, top, width, height])
    # Perform non maximum suppression to eliminate redundant overlapping boxes with lower confidences.
    indices = [ind[0] for ind in cv2.dnn.NMSBoxes(boxes, confidences, thresh_conf, thresh_nms)]
    ret = [[class_ids[i], confidences[i], boxes[i]] for i in indices]
    return ret
# Draw the predicted bounding box
def draw_pred(frame, class_id, classes, conf, left, top, right, bottom):
    # Draw a bounding box.
    cv2.rectangle(frame, (left, top), (right, bottom), (0, 0, 255))
    label = '%.2f' % conf
    # Get the label for the class name and its confidence
    if classes:
        assert (class_id < len(classes))
        label = '%s:%s' % (classes[class_id], label)
    # Display the label at the top of the bounding box
    label_size, base_line = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
    top = max(top, label_size[1])
    cv2.putText(frame, label, (left, top), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255))
'''
def remove_glare(img):
    """
    Reduce the effect of glaring in the image
@@ -262,7 +239,7 @@
    return corrected
def find_card(img, thresh_c=5, kernel_size=(3, 3), size_thresh=10000):
def find_card(img, thresh_c=5, kernel_size=(3, 3), size_thresh=10000, debug=False):
    """
    Find contours of all cards in the image
    :param img: source image
@@ -274,19 +251,29 @@
    # Typical pre-processing - grayscale, blurring, thresholding
    img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    img_blur = cv2.medianBlur(img_gray, 5)
    img_thresh = cv2.adaptiveThreshold(img_blur, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY_INV, 5, thresh_c)
    img_thresh = cv2.adaptiveThreshold(img_blur, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY_INV, 11, thresh_c)
    if debug:
        cv2.imshow('Thres', img_thresh)
    # Dilute the image, then erode them to remove minor noises
    kernel = np.ones(kernel_size, np.uint8)
    img_dilate = cv2.dilate(img_thresh, kernel, iterations=1)
    img_erode = cv2.erode(img_dilate, kernel, iterations=1)
    if debug:
        cv2.imshow('Eroded', img_erode)
    # Find the contour
    _, cnts, hier = cv2.findContours(img_erode, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    cnts, hier = cv2.findContours(img_erode, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    if len(cnts) == 0:
        print('no contours')
        return []
    img_cont = cv2.cvtColor(img_erode, cv2.COLOR_GRAY2BGR)
    img_cont_base = img_cont.copy()
    cnts2 = sorted(cnts, key=cv2.contourArea, reverse=True)
    cnts2 = cnts2[:10]
    for i in range(0, len(cnts2)):
        print(i, len(cnts2[i]))
    if debug:
        cv2.drawContours(img_cont, cnts2, -1, (0, 255, 0), 3)
        cv2.imshow('Contours', img_cont)
    # The hierarchy from cv2.findContours() is similar to a tree: each node has an access to the parent, the first child
    # their previous and next node
    # Using recursive search, find the uppermost contour in the hierarchy that satisfies the condition
@@ -302,8 +289,47 @@
        size = cv2.contourArea(cnt)
        peri = cv2.arcLength(cnt, True)
        approx = cv2.approxPolyDP(cnt, 0.04 * peri, True)
        print('Base Size:', size)
        print('Len Approx:', len(approx))
        if size >= size_thresh and len(approx) == 4:
            cnts_rect.append(approx)
            # lets see if we got a contour very close in size as child
            if i_child != -1:
                img_ccont = img_cont_base.copy()
                # lets collect all children
                c_list = [cnts[i_child]]
                h_info = hier[0][i_child]
                while h_info[0] != -1:
                    cld = cnts[h_info[0]]
                    c_list.append(cld)
                    h_info = hier[0][h_info[0]]
                # child with biggest area
                c_list.sort(key=cv2.contourArea, reverse=True)
                c_cnt = c_list[0]  # the biggest child
                if debug:
                    cv2.drawContours(img_ccont, c_list[:1], -1, (0, 255, 0), 1)
                    cv2.imshow('CCont', img_ccont)
                c_size = cv2.contourArea(c_cnt)
                c_approx = cv2.approxPolyDP(c_cnt, 0.04 * peri, True)
                if len(c_approx) == 4 and (c_size/size) > 0.85:
                    rect = cv2.minAreaRect(c_cnt)
                    box = cv2.boxPoints(rect)
                    box = np.intp(box)
                    print(c_cnt)
                    print(box)
                    print('CSize:', c_size, '%:', c_size/size)
                    b2 = []
                    for x in box:
                        b2.append([x])
                    cnts_rect.append(np.array(b2))
                else:
                    print('CF:', (c_size/size))
                    print('Size:', size)
                    cnts_rect.append(approx)
            else:
                #print('CF:', (c_size/size))
                print('Size:', size)
                cnts_rect.append(approx)
        else:
            if i_child != -1:
                stack.append((i_child, hier[0][i_child]))
@@ -343,14 +369,16 @@
        card_set = key[key.find('(') + 1:key.find(')')]
        confidence = sum(val) / f_len
        card_info = card_pool[(card_pool['name'] == card_name) & (card_pool['set'] == card_set)].iloc[0]
        img_name = '%s/card_img/tiny/%s/%s_%s.png' % (transform_data.data_dir, card_info['set'],
        img_name = '%s/card_img/tiny/%s/%s_%s.png' % (Config.data_dir, card_info['set'],
                                                      card_info['collector_number'],
                                                      fetch_data.get_valid_filename(card_info['name']))
        # If the card image is not found, just leave it blank
        if os.path.exists(img_name):
            card_img = cv2.imread(img_name)
        else:
            card_img = np.ones((h_card, w_card))
            card_img = np.ones((h_card, w_card, 3)) * 255
            cv2.putText(card_img, 'X', ((w_card - int(txt_scale * 25)) // 2, (h_card + int(txt_scale * 25)) // 2),
                        cv2.FONT_HERSHEY_SIMPLEX, txt_scale, (0, 0, 0), 2)
        # Insert the card image, card name, and confidence bar to the graph
        img_graph[y_anchor:y_anchor + h_card, x_anchor:x_anchor + w_card] = card_img
@@ -369,14 +397,13 @@
    return img_graph
def detect_frame(img, card_pool, hash_size=32, highfreq_factor=4, size_thresh=10000,
def detect_frame(img, card_pool, hash_size=32, size_thresh=10000,
                 out_path=None, display=True, debug=False):
    """
    Identify all cards in the input frame, display or save the frame if needed
    :param img: input frame
    :param card_pool: pandas dataframe of all card's information
    :param hash_size: param for pHash algorithm
    :param highfreq_factor: param for pHash algorithm
    :param size_thresh: threshold for size (in pixel) of the contour to be a candidate
    :param out_path: path to save the result
    :param display: flag for displaying the result
@@ -387,8 +414,10 @@
    img_result = img.copy()  # For displaying and saving
    det_cards = []
    # Detect contours of all cards in the image
    cnts = find_card(img_result, size_thresh=size_thresh)
    cnts = find_card(img_result, size_thresh=size_thresh, debug=debug)
    print('Countours:', len(cnts))
    for i in range(len(cnts)):
        print('Contour', i)
        cnt = cnts[i]
        # For the region of the image covered by the contour, transform them into a rectangular image
        pts = np.float32([p[0] for p in cnt])
@@ -402,44 +431,81 @@
        '''
        img_art = img_warp[47:249, 22:294]
        img_art = Image.fromarray(img_art.astype('uint8'), 'RGB')
        art_hash = ih.phash(img_art, hash_size=hash_size, highfreq_factor=highfreq_factor).hash.flatten()
        art_hash = ih.phash(img_art, hash_size=hash_size).hash.flatten()
        card_pool['hash_diff'] = card_pool['art_hash'].apply(lambda x: np.count_nonzero(x != art_hash))
        '''
        img_card = Image.fromarray(img_warp.astype('uint8'), 'RGB')
        img_card_size = img_warp.shape
        print(img_card_size)
        cut = [round(img_card_size[0]*0.57),round(img_card_size[0]*0.615),round(img_card_size[1]*0.81),round(img_card_size[1]*0.940)]
        print(cut)
        img_set_part = img_warp[cut[0]:cut[1], cut[2]:cut[3]]
        print(img_set_part.shape)
        img_set = Image.fromarray(img_set_part.astype('uint8'), 'RGB')
        print('img set')
        if debug:
            cv2.imshow("Set Img#%d" % i, img_set_part)
        # the stored values of hashes in the dataframe is pre-emptively flattened already to minimize computation time
        card_hash = ih.phash(img_card, hash_size=hash_size, highfreq_factor=highfreq_factor).hash.flatten()
        card_pool['hash_diff'] = card_pool['card_hash'].apply(lambda x: np.count_nonzero(x != card_hash))
        card_hash = ih.phash(img_card, hash_size=hash_size).hash.flatten()
        card_pool['hash_diff'] = card_pool['card_hash_%d' % hash_size]
        card_pool['hash_diff'] = card_pool['hash_diff'].apply(lambda x: np.count_nonzero(x != card_hash))
        min_card = card_pool[card_pool['hash_diff'] == min(card_pool['hash_diff'])].iloc[0]
        hash_diff = min_card['hash_diff']
        top_matches = sorted(card_pool['hash_diff'])
        card_one = card_pool[card_pool['hash_diff'] == top_matches[0]].iloc[0]
        card_two = card_pool[card_pool['hash_diff'] == top_matches[1]].iloc[0]
        if card_one['name'] == card_two['name'] and card_one['set'] != card_two['set']:
            set_img_hash = ih.whash(img_set, hash_size=hash_size).hash.flatten()
            cd_data = pd.DataFrame(columns=list(card_pool.columns.values))
            print(list(card_pool.columns.values))
            candidates = []
            for ix in range(0, 2):
                cd = card_pool[card_pool['hash_diff'] == top_matches[ix]].iloc[0]
                cd_data.loc[0 if cd_data.empty else cd_data.index.max()+1] = cd
                print('Idx:', ix, 'Name:', cd['name'], 'Set:', cd['set'], 'Diff:', top_matches[ix])
            cd_data['set_hash_diff'] = cd_data['set_hash_%d' % 64]
            cd_data['set_hash_diff'] = cd_data['set_hash_diff'].apply(lambda x: np.count_nonzero(x != set_img_hash))
            conf = sorted(cd_data['set_hash_diff'])
            print('Confs:', conf)
            best_match = cd_data[cd_data['set_hash_diff'] == min(cd_data['set_hash_diff'])].iloc[0]
            print('Best Match', 'Name:', best_match['name'], 'Set:', best_match['set'])
            min_card = best_match
        card_name = min_card['name']
        card_set = min_card['set']
        det_cards.append((card_name, card_set))
        hash_diff = min_card['hash_diff']
        # Render the result, and display them if needed
        cv2.drawContours(img_result, [cnt], -1, (0, 255, 0), 2)
        cv2.putText(img_result, card_name, (pts[0][0], pts[0][1]), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
        cv2.putText(img_result, card_name, (min(pts[0][0], pts[1][0]), min(pts[0][1], pts[1][1])),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
        if debug:
            # cv2.rectangle(img_warp, (22, 47), (294, 249), (0, 255, 0), 2)
            cv2.putText(img_warp, card_name + ', ' + str(hash_diff), (0, 50),
                        cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
            cv2.putText(img_warp, card_name + ':' + card_set + ', ' + str(hash_diff), (0, 20),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.4, (255, 255, 255), 1)
            cv2.imshow('card#%d' % i, img_warp)
    if display:
        cv2.imshow('Result', img_result)
        cv2.waitKey(0)
        inp = cv2.waitKey(0)
    if out_path is not None:
        print(out_path)
        cv2.imwrite(out_path, img_result.astype(np.uint8))
    return det_cards, img_result
def detect_video(capture, card_pool, hash_size=32, highfreq_factor=4, size_thresh=10000,
                 out_path=None, display=True, show_graph=True, debug=False):
def detect_video(capture, card_pool, hash_size=32, size_thresh=10000,
                 out_path=None, display=True, show_graph=True, debug=False, crop_x=0, crop_y=0):
    """
    Identify all cards in the continuous video stream, display or save the result if needed
    :param capture: input video stream
    :param card_pool: pandas dataframe of all card's information
    :param hash_size: param for pHash algorithm
    :param highfreq_factor: param for pHash algorithm
    :param size_thresh: threshold for size (in pixel) of the contour to be a candidate
    :param out_path: path to save the result
    :param display: flag for displaying the result
@@ -451,8 +517,9 @@
    # Get the dimension of the output video, and set it up
    if show_graph:
        img_graph = draw_card_graph({}, pd.DataFrame(), -1)  # Black image of the graph just to get the dimension
        width = round(capture.get(cv2.CAP_PROP_FRAME_WIDTH)) + img_graph.shape[1]
        height = max(round(capture.get(cv2.CAP_PROP_FRAME_HEIGHT)), img_graph.shape[0])
        width = round(capture.get(cv2.CAP_PROP_FRAME_WIDTH)) - 2*crop_x  + img_graph.shape[1]
        height = max(round(capture.get(cv2.CAP_PROP_FRAME_HEIGHT)) - 2*crop_y, img_graph.shape[0])
        height += 200  # some space to display last detected cards
    else:
        width = round(capture.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = round(capture.get(cv2.CAP_PROP_FRAME_HEIGHT))
@@ -461,9 +528,15 @@
    max_num_obj = 0
    f_len = 10  # number of frames to consider to check for existing cards
    exist_cards = {}
    exist_card_single = {}
    written_out_cards = set()
    found_cards = []
    try:
        while True:
            ret, frame = capture.read()
            croped_img = frame[crop_y:-crop_y, crop_x:-crop_x]
            fimg = cv2.flip(croped_img, -1)
            start_time = time.time()
            if not ret:
                # End of video
@@ -471,8 +544,8 @@
                cv2.waitKey(0)
                break
            # Detect all cards from the current frame
            det_cards, img_result = detect_frame(frame, card_pool, hash_size=hash_size, highfreq_factor=highfreq_factor,
                                                 size_thresh=size_thresh, out_path=None, display=False, debug=debug)
            det_cards, img_result = detect_frame(fimg, card_pool, hash_size=hash_size, size_thresh=size_thresh,
                                                 out_path=None, display=False, debug=debug)
            if show_graph:
                # If the card was already detected in the previous frame, append 1 to the list
                # If the card previously detected was not found in this trame, append 0 to the list
@@ -493,18 +566,50 @@
                    else:
                        exist_cards[key] = exist_cards[key][1 - f_len:] + [0]
                    if len(val) == f_len and sum(val) == 0:
                        gone.append(key)
                        gone.append(key)  # not there anymore
                det_card_map = {}
                gone_single =  []
                for card_name, card_set in det_cards:
                    skey = '%s (%s)' % (card_name, card_set)
                    det_card_map[skey] = (card_name, card_set)
                for key, val in exist_card_single.items():
                    if key in det_card_map:
                        exist_card_single[key] = exist_card_single[key][1 - f_len:] + [1]
                    else:
                        exist_card_single[key] = exist_card_single[key][1 - f_len:] + [0]
                    if len(val) == f_len and sum(val) == 0:
                        gone_single.append(key)
                        if key in written_out_cards:
                            written_out_cards.remove(key)
                    if len(val) == f_len and sum(val) == f_len:
                        if key not in written_out_cards and key in det_card_map:
                            written_out_cards.add(key)
                            found_cards.append(det_card_map[key])
                for key in det_card_map:
                    if key not in exist_card_single.keys():
                        exist_card_single[key] = [1]
                for key in gone_single:
                    exist_card_single.pop(key)
                for key in det_cards_list:
                    if key not in exist_cards.keys():
                        exist_cards[key] = [1]
                for key in gone:
                    exist_cards.pop(key)
                # Draw the graph based on the history of detected cards, then concatenate it with the result image
                img_graph = draw_card_graph(exist_cards, card_pool, f_len)
                img_save = np.zeros((height, width, 3), dtype=np.uint8)
                img_save[0:img_result.shape[0], 0:img_result.shape[1]] = img_result
                img_save[0:img_graph.shape[0], img_result.shape[1]:img_result.shape[1] + img_graph.shape[1]] = img_graph
                for c, card in enumerate(reversed(found_cards[-10:]), 1):
                    cv2.putText(img_save, f'{card[0]} ({card[1].upper()})',(0, height-200+18*c), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0))
            else:
                img_save = img_result
@@ -520,73 +625,118 @@
            print('Elapsed time: %.2f ms' % elapsed_ms)
            if out_path is not None:
                vid_writer.write(img_save.astype(np.uint8))
            cv2.waitKey(1)
            inp = cv2.waitKey(0)
            if 'q' == chr(inp & 255):
                break
    except KeyboardInterrupt:
        capture.release()
        if out_path is not None:
            vid_writer.release()
        cv2.destroyAllWindows()
        with open('detect.txt', 'w') as of:
            counter = collections.Counter(found_cards)
            for key in counter:
                of.write(f'{counter[key]} [{key[1].upper()}] {key[0]}\n')
def main():
def main(args):
    # Specify paths for all necessary files
    #test_path = os.path.abspath('test_file/test4.mp4')
    test_path = None
    out_dir = 'out'
    hash_size = 32
    highfreq_factor = 4
    pck_path = os.path.abspath('card_pool_%d_%d.pck' % (hash_size, highfreq_factor))
    hash_sizes = {16, 32}
    hash_sizes.add(args.hash_size)
    pck_path = os.path.abspath('card_pool.pck')
    if os.path.isfile(pck_path):
        card_pool = pd.read_pickle(pck_path)
    else:
        print('Warning: pickle for card database %s is not found!' % pck_path)
        # Merge database for all cards, then calculate pHash values of each, store them
        df_list = []
        for set_name in fetch_data.all_set_list:
            csv_name = '%s/csv/%s.csv' % (transform_data.data_dir, set_name)
        for set_name in Config.all_set_list:
            if set_name == 'con':
                set_name = 'con__'
            csv_name = '%s/csv/%s.csv' % (Config.data_dir, set_name)
            df = fetch_data.load_all_cards_text(csv_name)
            df_list.append(df)
        card_pool = pd.concat(df_list, sort=True)
        card_pool.reset_index(drop=True, inplace=True)
        card_pool.drop('Unnamed: 0', axis=1, inplace=True, errors='ignore')
        card_pool = calc_image_hashes(card_pool, save_to=pck_path, hash_size=hash_sizes)
    ch_key = 'card_hash_%d' % args.hash_size
    set_key = 'set_hash_%d' % 64
    if ch_key not in card_pool.columns:
        # we did not generate this hash_size yet
        print('We need to add hash_size=%d' % (args.hash_size,))
        card_pool = calc_image_hashes(card_pool, save_to=pck_path, hash_size=[args.hash_size])
        card_pool = calc_image_hashes(card_pool, save_to=pck_path, hash_size=hash_size, highfreq_factor=highfreq_factor)
    card_pool = card_pool[['name', 'set', 'collector_number', 'card_hash']]
    card_pool = card_pool[['name', 'set', 'collector_number', ch_key, set_key]]
    # Processing time is almost linear to the size of the database
    # Program can be much faster if the search scope for the card can be reduced
    card_pool = card_pool[card_pool['set'].isin(Config.set_2003_list)]
    # ImageHash is basically just one numpy.ndarray with (hash_size)^2 number of bits. pre-emptively flattening it
    # significantly increases speed for subtracting hashes in the future.
    card_pool['card_hash'] = card_pool['card_hash'].apply(lambda x: x.hash.flatten())
    card_pool[ch_key] = card_pool[ch_key].apply(lambda x: x.hash.flatten())
    card_pool[set_key] = card_pool[set_key].apply(lambda x: x.hash.flatten())
    # If the test file isn't given, use webcam to capture video
    if test_path is None:
        capture = cv2.VideoCapture(0)
        detect_video(capture, card_pool, out_path='%s/result.avi' % out_dir, display=True, show_graph=True, debug=False)
    if args.in_path is None:
        capture = cv2.VideoCapture(0, cv2.CAP_V4L)
        capture.set(cv2.CAP_PROP_FOURCC, cv2.VideoWriter_fourcc(*"MJPG"))
        capture.set(cv2.CAP_PROP_FRAME_WIDTH, 1920)
        capture.set(cv2.CAP_PROP_FRAME_HEIGHT, 1080)
        thres = int(((1920-2*500)*(1080-2*200)*0.3))
        print('Threshold:', thres)
        detect_video(capture, card_pool, hash_size=args.hash_size, out_path='%s/result.avi' % args.out_path,
                     display=args.display, show_graph=args.show_graph, debug=args.debug, crop_x=500, crop_y=200, size_thresh=thres)
        capture.release()
    else:
        # Save the detection result if out_dir is provided
        if out_dir is None or out_dir == '':
        # Save the detection result if args.out_path is provided
        if args.out_path is None:
            out_path = None
        else:
            f_name = os.path.split(test_path)[1]
            out_path = '%s/%s.avi' % (out_dir, f_name[:f_name.find('.')])
            f_name = os.path.split(args.in_path)[1]
            out_path = '%s/%s.avi' % (args.out_path, f_name[:f_name.find('.')])
        if not os.path.isfile(test_path):
            print('The test file %s doesn\'t exist!' % os.path.abspath(test_path))
        if not os.path.isfile(args.in_path):
            print('The test file %s doesn\'t exist!' % os.path.abspath(args.in_path))
            return
        # Check if test file is image or video
        test_ext = test_path[test_path.find('.') + 1:]
        test_ext = args.in_path[args.in_path.find('.') + 1:]
        if test_ext in ['jpg', 'jpeg', 'bmp', 'png', 'tiff']:
            # Test file is an image
            img = cv2.imread(test_path)
            detect_frame(img, card_pool, out_path=out_path)
            img = cv2.imread(args.in_path)
            if img is None:
                print('Could not read', args.in_path)
            detect_frame(img, card_pool, hash_size=args.hash_size, out_path=out_path, display=args.display,
                         debug=args.debug)
        else:
            # Test file is a video
            capture = cv2.VideoCapture(test_path)
            detect_video(capture, card_pool, out_path=out_path, display=True, show_graph=True, debug=False)
            capture = cv2.VideoCapture(args.in_path)
            detect_video(capture, card_pool, hash_size=args.hash_size, out_path=out_path, display=args.display,
                         show_graph=args.show_graph, debug=args.debug)
            capture.release()
    pass
if __name__ == '__main__':
    main()
    parser = argparse.ArgumentParser()
    parser.add_argument('-i', '--in', dest='in_path', help='Path of the input file. For webcam, leave it blank',
                        type=str)
    parser.add_argument('-o', '--out', dest='out_path', help='Path of the output directory to save the result',
                        type=str)
    parser.add_argument('-hs', '--hash_size', dest='hash_size',
                        help='Size of the hash for pHash algorithm', type=int, default=16)
    parser.add_argument('-dsp', '--display', dest='display', help='Display the result', action='store_true',
                        default=False)
    parser.add_argument('-dbg', '--debug', dest='debug', help='Enable debug mode', action='store_true', default=False)
    parser.add_argument('-gph', '--show_graph', dest='show_graph', help='Display the graph for video output',
                        action='store_true', default=False)
    args = parser.parse_args()
    if not args.display and args.out_path is None:
        # Then why the heck are you running this thing in the first place?
        print('The program isn\'t displaying nor saving any output file. Please change the setting and try again.')
        exit()
    main(args)