Initial code commit

xJakub · xJakub · commit 249d45dddf5c · 2017-11-12T15:16:08.000+01:00
diff --git a/parse.py b/parse.py
@@ -0,0 +1,360 @@
+import sys, itertools, requests, json, os.path, hashlib, time, cv2, numpy as np
+
+api_cache_folder = 'cache'
+api_url = 'http://schoolido.lu/api/cards/?page_size=100'
+round_card_images_folder = 'round_card_images'
+all_patterns_file = 'cache/all_patterns.png'
+
+circle_size = 128
+
+rarity_top = 2
+rarity_height = 22
+rarity_left = 2
+rarity_width = 22
+
+pattern_top = 16
+pattern_height = 80
+pattern_left = 16
+pattern_width = 80
+pattern_ratio = 0.5
+
+match_top = 12
+match_height = 88
+match_left = 12
+match_width = 88
+
+
+def api_get(url):
+    """
+    Make a request against the schoolido.lu api. The results are cached
+    :param url: url
+    :return: result data
+    """
+    cache_file = "%s/%s.json" % (api_cache_folder, hashlib.md5(url).hexdigest())
+    if not os.path.isfile(cache_file):
+        if not os.path.isdir(api_cache_folder):
+            os.mkdir(api_cache_folder)
+        api_result = []
+        while url is not None:
+            sys.stderr.write('Fetching %s\n' % url)
+            r = requests.get(url)
+            if r.status_code != 200:
+                raise Exception(r.status_code)
+            json_data = r.json()
+            url = json_data['next']
+            api_result += json_data['results']
+        with open(cache_file, 'w') as f:
+            json.dump(api_result, f, indent=4, sort_keys=True)
+
+        if os.path.isfile(all_patterns_file):
+            os.unlink(all_patterns_file)
+
+    with open(cache_file) as f:
+        result = json.load(f)
+    return result
+
+
+def get_card_images(card):
+    """
+    Returns a list of a card's round images including the remote and local (if cached) paths, and whether it
+    belongs to the idolized version of the card
+    :param card: card data from schoolido.lu api
+    :return: list of (url, local_path, idolized) tuples 
+    """
+    urls = [(card['round_card_image'], False), (card['round_card_idolized_image'], True)]
+    result = []
+    for url, idolized in urls:
+        if url is None:
+            continue
+        if url[0:2] == '//':
+            url = 'http:%s' % url
+        local_path = "%s/%s" % (round_card_images_folder, os.path.basename(url).split('?')[0])
+        result.append((url, local_path, idolized))
+    return result
+
+
+def fetch_round_card(card):
+    """
+    Downloads a card's round images (if they don't exist).
+    There is a 100ms sleep time after each download to give the host some rest.
+    :param card: card data from schoolido.lu api
+    """
+    if not os.path.isdir(round_card_images_folder):
+        os.mkdir(round_card_images_folder)
+
+    for url, dest, idolized in get_card_images(card):
+        if not os.path.isfile(dest):
+            sys.stderr.write('Fetching %s\n' % url)
+            r = requests.get(url)
+            if r.status_code != 200:
+                raise Exception(r.status_code)
+            with open(dest, 'wb') as f:
+                f.write(r.content)
+            time.sleep(0.1)
+
+
+def get_card_group(card, idolized):
+    """
+    Returns an identifier for a card's group, to help separate them by rarities
+    :param card: card data from schoolido.lu api
+    :param idolized: True if idolized version
+    :return: card group name
+    """
+    return "%s-%s%d" % (card['rarity'], card['attribute'], 1 if idolized else 0)
+
+
+def make_rarity_patterns(cards):
+    """
+    Make and return patterns for each card rarity
+    :param cards: list of card data from schoolido.lu ap
+    :return: list of (group identifier, pattern, idolized) tuples
+    """
+    did = {}
+    patterns = []
+
+    for card in cards:
+        images = get_card_images(card)
+        for index, (url, local_path, idolized) in enumerate(images):
+            group = get_card_group(card, idolized)
+            if group in did:
+                continue
+            did[group] = 1
+            im = cv2.imread(local_path)
+            im_cropped = im[rarity_top:rarity_top + rarity_height, rarity_left:rarity_left + rarity_width]
+            patterns.append((group, im_cropped, idolized))
+    return patterns
+
+
+def get_pattern_coordinates(index, idolized):
+    """
+    Returns coordinates for extracting the card's pattern from the all_patterns image  
+    :param index: card index inside the cards obtained from the api
+    :param idolized: True if idolized version
+    :return: coordinates to be used as im_patterns[coordinates]
+    """
+    real_index = (index * 2 + 1) if idolized else (index * 2)
+    y1 = int(pattern_height * pattern_ratio * real_index)
+    y2 = int(pattern_height * pattern_ratio * (real_index + 1))
+    x1 = 0
+    x2 = int(pattern_width * pattern_ratio)
+    return np.index_exp[y1:y2, x1:x2]
+
+
+def generate_all_patterns(cards):
+    """
+    Generates an all_patterns image
+    :param cards: list of card data from schoolido.lu ap
+    """
+    if not os.path.isdir(os.path.dirname(all_patterns_file)):
+        os.mkdir(os.path.dirname(all_patterns_file))
+    coordinates = get_pattern_coordinates(len(cards), True)
+    im_patterns_height, im_patterns_width = [c.stop for c in coordinates]
+    im_patterns = np.ones((im_patterns_height, im_patterns_width, 3), np.uint8)
+
+    for index, card in enumerate(cards):
+        for url, local_path, idolized in get_card_images(card):
+            if not os.path.exists(local_path):
+                fetch_round_card(card)
+            im = cv2.imread(local_path)
+            im_cropped = im[pattern_top:pattern_top + pattern_height, pattern_left:pattern_left + pattern_width]
+            im_cropped = cv2.resize(im_cropped, (0, 0), fx=pattern_ratio, fy=pattern_ratio)
+            coordinates = get_pattern_coordinates(index, idolized)
+            im_patterns[coordinates] = im_cropped
+    cv2.imwrite(all_patterns_file, im_patterns)
+
+
+def make_card_patterns(cards):
+    """
+    Return patterns for each round card image
+    :param cards: list of card data from schoolido.lu ap
+    :return: a map, with group identifiers as keys and a list of (pattern, card data, idolized) as values
+    """
+    if not os.path.isfile(all_patterns_file):
+        generate_all_patterns(cards)
+
+    im_patterns = cv2.imread(all_patterns_file)
+    result = {}
+    for index, card in enumerate(cards):
+        for url, local_path, idolized in get_card_images(card):
+            coordinates = get_pattern_coordinates(index, idolized)
+            im_cropped = im_patterns[coordinates]
+
+            groups = [get_card_group(card, idolized)]
+            # Include non-idolized patterns in both groups, since an idolized card can show the unidolized version
+            if not idolized:
+                groups.append(get_card_group(card, True))
+            for group in groups:
+                if group not in result:
+                    result[group] = []
+                result[group].append((im_cropped, card, idolized))
+    return result
+
+
+def vertical_split(im):
+    """
+    Split an image vertically. This works for member list as well as (well, most times) scouting screenshots 
+    :param im: screenshot data
+    :return: the ratio for resizing the screenshot, a list of (start, stop) chunks to use when looping
+    """
+    im_gray = cv2.cvtColor(im, cv2.COLOR_RGB2GRAY)
+    thresh = 223
+
+    im_bw = cv2.threshold(im_gray, thresh, 255, cv2.THRESH_BINARY)[1]
+    rows = [y for y in xrange(len(im_bw)) if
+            max([len(list(g)) for e, g in itertools.groupby(im_bw[y]) if e == 255] or [0]) > len(im[0]) / 2]
+
+    row_groups = [(min(arr), max(arr)) for arr in np.split(rows, np.where(np.diff(rows) != 1)[0] + 1)]
+    row_sizes = [(b[0] - a[1] - 1) for a, b in zip(row_groups, row_groups[1:]) if (b[0] - a[1] - 1) > 100]
+    row_positions = [(a[1], b[0]) for a, b in zip(row_groups, row_groups[1:]) if (b[0] - a[1] - 1) > 100]
+
+    ratio = 1.0 * circle_size / np.median(row_sizes)
+    return ratio, row_positions
+
+
+def horizontal_split(im):
+    """
+    Horizontally split a row into smaller chunks which could contain rounded card images
+    :param im: a screenshot's row obtained after using vertical_split's result
+    :return: a list of (start, stop) chunks to use when looping
+    """
+    im_gray = cv2.cvtColor(im, cv2.COLOR_RGB2GRAY)
+    thresh = 223
+
+    im_bw = cv2.threshold(im_gray, thresh, 255, cv2.THRESH_BINARY)[1]
+    im_width = len(im[0])
+
+    cols = [y for y in xrange(len(im_bw[0])) if 0 not in np.transpose(im_bw)[y]]
+
+    col_groups = [(min(arr), max(arr)) for arr in np.split(cols, np.where(np.diff(cols) != 1)[0] + 1)]
+    col_positions = [(b, b + circle_size) for a, b in col_groups] + [(a - circle_size, a) for a, b in col_groups]
+    unique_col_positions = []
+    for col_position in col_positions:
+        if col_position[0] < 0 or col_position[1] >= im_width:
+            continue
+        if any([abs(col_position[0] - col_position2[0]) <= 2 and abs(col_position[1] - col_position2[1]) <= 2 for
+                col_position2 in unique_col_positions]):
+            break
+        unique_col_positions.append(col_position)
+
+    return unique_col_positions
+
+
+def get_matching_cards(im_match, possible_cards):
+    """
+    Tries to match a region against a list of possible card patterns
+    :param im_match: region data
+    :param possible_cards: list of (card_pattern, card data, idolized) tuples
+    :return: the matching cards, as (card data, idolized) tuples
+    """
+    im_match = cv2.resize(im_match, (0, 0), fx=pattern_ratio, fy=pattern_ratio)
+    matches = []
+    threshold = 0.8
+    for card_pattern, card, idolized in possible_cards:
+        if len(im_match) < len(card_pattern) or len(im_match[0]) < len(card_pattern[0]):
+            break
+        res = cv2.matchTemplate(im_match, card_pattern, cv2.TM_CCOEFF_NORMED)
+        loc = np.where(res >= threshold)
+        card_pattern_matches = zip(*loc[::-1])
+        if card_pattern_matches:
+            matches.append((card, idolized))
+    return matches
+
+
+def search_row(im, rarity_patterns, card_patterns):
+    """
+    Searches a row for round card pattern matches
+    :param im: the image's row data
+    :param rarity_patterns: result from make_rarity_patterns
+    :param card_patterns: result from make_make_card_patterns 
+    :return: list of match information
+    """
+    found_cards = []
+
+    im_width = len(im[0])
+    im_height = len(im)
+
+    col_positions = horizontal_split(im)
+
+    for group_min_x, group_max_x in col_positions:
+        im4 = im[:, group_min_x:group_max_x]
+
+        for pattern_group, pattern, pattern_idolized in rarity_patterns:
+            im3 = im4[0:rarity_top + rarity_height + 2, 0:rarity_left + rarity_width + 2]
+            res = cv2.matchTemplate(im3, pattern, cv2.TM_CCOEFF_NORMED)
+
+            threshold = 0.7
+            loc = np.where(res >= threshold)
+            rarity_pattern_matches = zip(*loc[::-1])
+
+            for index, pt in enumerate(rarity_pattern_matches):
+                distances = [abs(pt[0] - pt2[0]) + abs(pt[1] - pt2[1]) for pt2 in rarity_pattern_matches[:index]]
+                if distances and min(distances) < 100:
+                    continue
+
+                match_x = pt[0] - rarity_left
+                match_y = pt[1] - rarity_top
+                relative_x = 1.0 * (match_x + group_min_x) / im_width
+                relative_y = 1.0 * match_y / im_height
+
+                im_match = im4[pt[1] + match_top:pt[1] + match_top + match_height,
+                           pt[0] + match_left:pt[0] + match_left + match_width]
+                for matching_card, matching_idolized in get_matching_cards(im_match, card_patterns[pattern_group]):
+                    found_cards.append({
+                        'card': matching_card,
+                        'idolized': matching_idolized,
+                        'relative_x': relative_x,
+                        'relative_y': relative_y,
+                    })
+    return found_cards
+
+
+def main():
+    """
+    Reads a screenshot and tries to find round card image matches, using data and images from schoolido.lu
+    The screenshot's path is read from argv[1]
+    The result is written as json to stdout
+    Extra info is written to stderr
+    Note:   Cache is stored inside cache/ and round_card_images/
+            Delete the cache/ folder to force downloading of card data
+    """
+    os.chdir(os.path.dirname(os.path.realpath(__file__)))
+    start = time.time()
+    cards = api_get(api_url)
+    sys.stderr.write("%f elapsed after api_get\n" % (time.time() - start))
+
+    card_patterns = make_card_patterns(cards)
+    sys.stderr.write("%f elapsed after card_patterns\n" % (time.time() - start))
+
+    rarity_patterns = make_rarity_patterns(cards)
+    sys.stderr.write("%f elapsed after rarity_patterns\n" % (time.time() - start))
+
+    im_original = cv2.imread(sys.argv[1])
+    sys.stderr.write("%f elapsed after imread\n" % (time.time() - start))
+
+    ratio, row_positions = vertical_split(im_original)
+    sys.stderr.write("%f elapsed after vertical_split\n" % (time.time() - start))
+
+    original_width = len(im_original[0])
+    original_height = len(im_original)
+    resized_width = original_width * ratio
+    resized_height = original_height * ratio
+
+    found_cards = []
+    for group_min_y, group_max_y in row_positions:
+        im = cv2.resize(im_original[group_min_y:group_max_y, :], (0, 0), fx=ratio, fy=ratio)
+        group_min_relative_y = 1.0 * group_min_y / original_height
+        group_relative_h = 1.0 * (group_max_y - group_min_y) / original_height
+        for found_card in search_row(im, rarity_patterns, card_patterns):
+            found_card['relative_y'] = group_min_relative_y + group_relative_h * found_card['relative_y']
+            found_cards.append(found_card)
+
+    for found_card in found_cards:
+        found_card['relative_w'] = 1.0 * circle_size / resized_width
+        found_card['relative_h'] = 1.0 * circle_size / resized_height
+
+    sys.stderr.write("%f elapsed before dumping\n" % (time.time() - start))
+    sys.stdout.write(json.dumps(found_cards, indent=4, sort_keys=True))
+
+
+if __name__ == '__main__':
+    main()
