stock & HERO

aoirint · Aug 5, 2019 · 078fbca · 078fbca
1 parent fc81d07
commit 078fbca
Show file tree

Hide file tree

Showing 7 changed files with 259 additions and 22 deletions.
diff --git a/README.md b/README.md
@@ -9,7 +9,7 @@ python3 SSBUFrameAnalyzer.py DIGIT_DICTIONARY_FILE CHARA_DICTIONARY_FILE SCREENS
 
 1280x720のカラー画像のみ受け付けます（mainの場合内部で自動リサイズします）。
 
-`DIGIT_DICTIONALY_FILE`、`CHARA_DICTIONARY_FILE`はそれぞれ、ダメージ値の数値画像から計算したHOG特徴量、対戦中のキャラ顔画像から計算したHOG特徴量を格納したjsonです（Releasesに添付）。あんまり検証してません。
+`DIGIT_DICTIONALY_FILE`、`CHARA_DICTIONARY_FILE`、`STOCK_DICTIONARY_FILE`はそれぞれ、ダメージ値の数値画像から計算したHOG特徴量、対戦中のキャラ顔画像から計算したHOG特徴量、対戦中のキャラストック画像から計算したHOG特徴量を格納したjsonです（Releasesに添付）。（精度は）あんまり検証してません。
 
 いまのところ、2人、3人、4人対戦のみです。
 
@@ -23,3 +23,6 @@ python3 SSBUFrameAnalyzer.py DIGIT_DICTIONARY_FILE CHARA_DICTIONARY_FILE SCREENS
     - numpy
     - scikit-image
 
+## 対応キャラクター（仮）
+
+勇者までの全キャラクター（DLC: パックンフラワー、ジョーカー、勇者含む）
diff --git a/calc_chara_hog.py b/calc_chara_hog.py
@@ -0,0 +1,32 @@
+import argparse
+from skimage.feature import hog
+import numpy as np
+import os
+import cv2
+import json
+
+parser = argparse.ArgumentParser()
+parser.add_argument('in_dir', type=str)
+parser.add_argument('out_file', type=str)
+
+args = parser.parse_args()
+
+data = {}
+for cat_name in sorted(os.listdir(args.in_dir)):
+    cat_path = os.path.join(args.in_dir, cat_name)
+    if not os.path.isdir(cat_path):
+        continue
+
+    data_chara = []
+    for in_file in sorted(os.listdir(cat_path)):
+        in_path = os.path.join(cat_path, in_file)
+        img = cv2.imread(in_path, 0) # grayscale
+        # img = cv2.resize(img, (110, 110))
+        ft = hog(img)
+        data_chara.append(ft.tolist())
+
+    data[cat_name] = data_chara
+
+os.makedirs(os.path.dirname(args.out_file), exist_ok=True)
+with open(args.out_file, 'w') as fp:
+    json.dump(data, fp)
diff --git a/calc_damage_hog.py b/calc_damage_hog.py
@@ -0,0 +1,35 @@
+import argparse
+from skimage.feature import hog
+import numpy as np
+import os
+import cv2
+import json
+
+parser = argparse.ArgumentParser()
+parser.add_argument('in_dir', type=str)
+parser.add_argument('out_file', type=str)
+
+args = parser.parse_args()
+
+data = {}
+for cat_name in sorted(os.listdir(args.in_dir)):
+    cat_path = os.path.join(args.in_dir, cat_name)
+    if not os.path.isdir(cat_path):
+        continue
+
+    digit = int(cat_name)
+    print(digit)
+
+    data_digit = []
+    for in_file in sorted(os.listdir(cat_path)):
+        in_path = os.path.join(cat_path, in_file)
+        img = cv2.imread(in_path, 0) # grayscale
+        img = cv2.resize(img, (35, 55))
+        ft = hog(img)
+        data_digit.append(ft.tolist())
+
+    data[digit] = data_digit
+
+os.makedirs(os.path.dirname(args.out_file), exist_ok=True)
+with open(args.out_file, 'w') as fp:
+    json.dump(data, fp)
diff --git a/calc_stock_hog.py b/calc_stock_hog.py
@@ -0,0 +1,32 @@
+import argparse
+from skimage.feature import hog
+import numpy as np
+import os
+import cv2
+import json
+
+parser = argparse.ArgumentParser()
+parser.add_argument('in_dir', type=str)
+parser.add_argument('out_file', type=str)
+
+args = parser.parse_args()
+
+data = {}
+for cat_name in sorted(os.listdir(args.in_dir)):
+    cat_path = os.path.join(args.in_dir, cat_name)
+    if not os.path.isdir(cat_path):
+        continue
+
+    data_chara = []
+    for in_file in sorted(os.listdir(cat_path)):
+        in_path = os.path.join(cat_path, in_file)
+        img = cv2.imread(in_path, 0) # grayscale
+        img = cv2.resize(img, (30, 30))
+        ft = hog(img)
+        data_chara.append(ft.tolist())
+
+    data[cat_name] = data_chara
+
+os.makedirs(os.path.dirname(args.out_file), exist_ok=True)
+with open(args.out_file, 'w') as fp:
+    json.dump(data, fp)
diff --git a/get_chara_shots.py b/get_chara_shots.py
@@ -0,0 +1,56 @@
+import cv2
+import argparse
+from PIL import Image
+import os
+import numpy as np
+import SSBUBoundingBoxUtil
+
+parser = argparse.ArgumentParser()
+parser.add_argument('in_file', type=str)
+parser.add_argument('out_dir', type=str)
+parser.add_argument('fighter_num', type=int)
+args = parser.parse_args()
+
+
+os.makedirs(args.out_dir, exist_ok=True)
+
+frame_files = None
+cap = None
+if os.path.isdir(args.in_file):
+    frame_files = [ os.path.join(args.in_file, file) for file in os.listdir(args.in_file) ]
+else:
+    cap = cv2.VideoCapture(args.in_file)
+
+frame_idx = 0
+def next_frame():
+    global frame_idx
+
+    if frame_files is not None:
+        if len(frame_files) <= frame_idx:
+            return None
+        frame = cv2.imread(frame_files[frame_idx])
+    else:
+        ret, frame = cap.read()
+
+    frame_idx += 1
+    return frame
+
+in_file = os.path.basename(args.in_file)
+
+while True:
+    img = next_frame()
+    if img is None:
+        break
+
+    bboxes = SSBUBoundingBoxUtil.fighters_chara_bbox(args.fighter_num)
+    for fighter_idx in range(len(bboxes)):
+        # if fighter_idx != 3:
+        #     continue
+        bbox = bboxes[fighter_idx]
+        chara = img[bbox[1]:bbox[1]+bbox[3], bbox[0]:bbox[0]+bbox[2]]
+
+        file = '%s-f%d-%d.png' % (in_file, fighter_idx, frame_idx, )
+        path = os.path.join(args.out_dir, file)
+        cv2.imwrite(path, chara)
+
+        print(frame_idx, fighter_idx)
diff --git a/get_stock_shots.py b/get_stock_shots.py
@@ -0,0 +1,59 @@
+import cv2
+import argparse
+from PIL import Image
+import os
+import numpy as np
+import SSBUBoundingBoxUtil
+
+parser = argparse.ArgumentParser()
+parser.add_argument('in_file', type=str)
+parser.add_argument('out_dir', type=str)
+parser.add_argument('fighter_num', type=int)
+args = parser.parse_args()
+
+
+os.makedirs(args.out_dir, exist_ok=True)
+
+
+frame_files = None
+cap = None
+if os.path.isdir(args.in_file):
+    frame_files = [ os.path.join(args.in_file, file) for file in os.listdir(args.in_file) ]
+else:
+    cap = cv2.VideoCapture(args.in_file)
+
+frame_idx = 0
+def next_frame():
+    global frame_idx
+
+    if frame_files is not None:
+        if len(frame_files) <= frame_idx:
+            return None
+        frame = cv2.imread(frame_files[frame_idx])
+    else:
+        ret, frame = cap.read()
+
+    frame_idx += 1
+    return frame
+
+in_file = os.path.basename(args.in_file)
+
+fighters_stock_bboxes = SSBUBoundingBoxUtil.fighters_stock_bboxes(args.fighter_num, stock_num=3)
+
+while True:
+    img = next_frame()
+    if img is None:
+        break
+
+    for fighter_idx, bboxes in enumerate(fighters_stock_bboxes):
+        for idx in range(len(bboxes)):
+            # if idx != 3:
+            #     continue
+            bbox = bboxes[idx]
+            chara = img[bbox[1]:bbox[1]+bbox[3], bbox[0]:bbox[0]+bbox[2]]
+
+            file = '%s-f%d-s%d-%d.png' % (in_file, fighter_idx, idx, frame_idx, )
+            path = os.path.join(args.out_dir, file)
+            cv2.imwrite(path, chara)
+
+        print(frame_idx, fighter_idx)
diff --git a/main_camera.py b/main_camera.py
@@ -10,14 +10,19 @@
 import threading
 
 parser = argparse.ArgumentParser()
-parser.add_argument('outdir', type=str)
-parser.add_argument('-d', '--digit_dictionary', type=str, default='../digit_dictionary.json')
-parser.add_argument('-c', '--chara_dictionary', type=str, default='../chara_dictionary.json')
-parser.add_argument('-s', '--stock_dictionary', type=str, default='../stock_dictionary.json')
+parser.add_argument('-d', '--digit_dictionary', type=str, default='../digit_dictionary_v1.json')
+parser.add_argument('-c', '--chara_dictionary', type=str, default='../chara_dictionary_v3.json')
+parser.add_argument('-s', '--stock_dictionary', type=str, default='../stock_dictionary_v3.json')
 parser.add_argument('-f', '--fighter_num', type=int, default=2) # TODO: predict
-parser.add_argument('-p', '--capture_index', type=int, default=2)
+parser.add_argument('-p', '--capture_device', type=int, default=None)
+parser.add_argument('-i', '--video', type=str, default=None)
+parser.add_argument('--dump', action='store_true')
+parser.add_argument('-o', '--outdir', type=str, default='./')
+parser.add_argument('--sync', action='store_true')
 args = parser.parse_args()
 
+assert (args.capture_device is not None) ^ (args.video is not None)
+
 os.makedirs(args.outdir, exist_ok=True)
 
 
@@ -40,18 +45,26 @@
 analyzer = SSBUFrameAnalyzer(digit_classifier=digit_classifier, name_recognizer=name_recognizer, chara_classifier=chara_classifier, stock_classifier=stock_classifier)
 
 
-fps = 30
+if args.capture_device:
+    fps = 30
+
+    cap = cv2.VideoCapture(args.capture_device)
+    cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
+    cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)
+    cap.set(cv2.CAP_PROP_FPS, fps)
+else:
+    cap = cv2.VideoCapture(args.video)
 
-cap = cv2.VideoCapture(args.capture_index)
-cap.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
-cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)
-cap.set(cv2.CAP_PROP_FPS, fps)
+if args.dump:
+    outfile = '%f.avi' % time.time()
+    outpath = os.path.join(args.outdir, outfile)
 
-outfile = '%f.avi' % time.time()
-outpath = os.path.join(args.outdir, outfile)
+    fourcc = cv2.VideoWriter_fourcc(*'XVID')
+    out = cv2.VideoWriter(outpath, fourcc, fps, (1280, 720))
 
-fourcc = cv2.VideoWriter_fourcc(*'XVID')
-out = cv2.VideoWriter(outpath, fourcc, fps, (1280, 720))
+    print('Dump mode: %s' % outpath)
+else:
+    print('No Dump mode')
 
 task = None
 def do_task():
@@ -64,26 +77,33 @@ def do_task():
 
     print('-' * 40)
     print(ret)
-    print('FPS: %f (%f s)' % (1/elapsed, elapsed, ))
+    print('Task FPS: %f (%f s)' % (1/elapsed, elapsed, ))
 
     task = None
 
 while True:
+    ts = time.time()
     ret, frame = cap.read()
 
-    out.write(frame)
-
-    # print(frame.shape)
-    cv2.imshow('img', frame)
-
     if task is None:
         task = threading.Thread(target=do_task)
         task.start()
+    if args.sync:
+        task.join()
+
+    if args.dump:
+        out.write(frame)
+
+    # print(frame.shape)
+    cv2.imshow('img', frame)
+    elapsed = time.time() - ts
+    print('Camera FPS: %f (%f s)' % (1/elapsed, elapsed, ))
 
     c = cv2.waitKey(1)
     if c & 0xFF == ord('q'):
         break
 
 cap.release()
-out.release()
+if args.dump:
+    out.release()
 cv2.destroyAllWindows()