michimani · November 23, 2020 14:59 · Aug 29, 2020 · Aug 29, 2020
diff --git a/detect_smash_bros_sp_power_v2.py b/detect_smash_bros_sp_power_v2.py
@@ -0,0 +1,118 @@
+import boto3
+import cv2
+import datetime
+import hashlib
+import json
+import os
+import re
+import sys
+import time
+
+TMP_DIR = './tmp'
+THUMBNAIL_DIR = './char_thumb'
+
+reko = boto3.client('rekognition')
+
+target_area = {
+    'A': {'X': 0.57187, 'Y': 0.1875},
+    'B': {'X': 0.9, 'Y': 0.1875},
+    'C': {'X': 0.9, 'Y': 0.43125},
+    'D': {'X': 0.57187, 'Y': 0.43125},
+}
+
+character_area = {
+    'A': {'X': 911, 'Y': 474},
+    'B': {'X': 959, 'Y': 474},
+    'C': {'X': 959, 'Y': 522},
+    'D': {'X': 911, 'Y': 522},
+}
+
+
+def get_datetime_from_image_path(image_path):
+    datetime_part = re.sub(r'.*(\d{16}).*', r'\1', image_path)[:14]
+    dt = datetime.datetime.strptime(datetime_part, '%Y%m%d%H%M%S')
+    datetime_str = '{0:%Y-%m-%d %H:%M:%S}'.format(dt)
+    return datetime_str
+
+
+def get_character_hash(image_path):
+    img = cv2.imread(image_path)
+
+    trimed = img[
+        character_area['A']['Y']: character_area['D']['Y'],
+        character_area['A']['X']: character_area['C']['X']]
+
+    trimed_bytes = trimed.tobytes()
+    trimed_hash = hashlib.sha1(trimed_bytes).hexdigest()
+
+    thumb_path = f'{THUMBNAIL_DIR}/{trimed_hash}.jpg'
+    if os.path.exists(thumb_path) is False:
+        cv2.imwrite(thumb_path, trimed)
+
+    return trimed_hash
+
+
+def get_bytes_of_image(image_path):
+    img = cv2.imread(image_path)
+    trimed = img[
+        360: 720,
+        640: 1280
+    ]
+
+    ts = str(time.time())
+    trimed_tmp_path = f'{TMP_DIR}/{ts}.jpg'
+    cv2.imwrite(trimed_tmp_path, trimed)
+
+    with open(trimed_tmp_path, mode='rb') as img:
+        image_bytes = img.read()
+
+    os.remove(trimed_tmp_path)
+
+    return image_bytes
+
+
+def detect_text(image_path):
+    image_bytes = get_bytes_of_image(image_path)
+    return reko.detect_text(
+        Image={
+            'Bytes': image_bytes
+        }
+    )
+
+
+def get_smash_power(image_path):
+    smash_power = 0
+    detect_res = detect_text(image_path)
+
+    for detected in detect_res['TextDetections']:
+        polygon = detected['Geometry']['Polygon']
+        if polygon[0]['X'] > target_area['A']['X'] and polygon[0]['Y'] > target_area['A']['Y'] \
+                and polygon[1]['X'] < target_area['B']['X'] and polygon[1]['Y'] > target_area['B']['Y'] \
+                and polygon[2]['X'] < target_area['C']['X'] and polygon[2]['Y'] < target_area['C']['Y'] \
+                and polygon[3]['X'] > target_area['D']['X'] and polygon[3]['Y'] < target_area['D']['Y']:
+            smash_power = int(detected['DetectedText'].replace(',', ''))
+
+    return smash_power
+
+
+def generate_smash_power_item(image_path):
+    # Datetime of data
+    datetime_str = get_datetime_from_image_path(image_path)
+
+    # Get Smash Power via Amazon Rekognition
+    smash_power = get_smash_power(image_path)
+
+    # Get character hash (and Generate character thumbnail)
+    character_hash = get_character_hash(image_path)
+
+    return {
+        'character_hash': character_hash,
+        'capture_datetime': datetime_str,
+        'power': smash_power
+    }
+
+
+if __name__ == '__main__':
+    image_path = sys.argv[1]
+    smash_power_item = generate_smash_power_item(image_path)
+    print(json.dumps(smash_power_item, indent=2, ensure_ascii=False))
No results found