stock (low accuracy)

aoirint · Jul 26, 2019 · fc81d07 · fc81d07
1 parent b46322b
commit fc81d07
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/main_camera.py b/main_camera.py
@@ -5,14 +5,17 @@
 from SSBUDigitClassifier import SSBUDigitClassifier
 from SSBUNameRecognizer import SSBUNameRecognizer
 from SSBUCharaClassifier import SSBUCharaClassifier
+from SSBUStockClassifier import SSBUStockClassifier
 from SSBUFrameAnalyzer import SSBUFrameAnalyzer
 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('-f', '--fighter_num', type=int, default=2) # TODO: predict
+parser.add_argument('-p', '--capture_index', type=int, default=2)
 args = parser.parse_args()
 
 os.makedirs(args.outdir, exist_ok=True)
@@ -30,12 +33,16 @@
 chara_classifier = SSBUCharaClassifier(feature_json=args.chara_dictionary)
 print('loaded chara classifier')
 
-analyzer = SSBUFrameAnalyzer(digit_classifier=digit_classifier, name_recognizer=name_recognizer, chara_classifier=chara_classifier)
+print('loading stock classifier...')
+stock_classifier = SSBUStockClassifier(feature_json=args.stock_dictionary)
+print('loaded stock classifier')
+
+analyzer = SSBUFrameAnalyzer(digit_classifier=digit_classifier, name_recognizer=name_recognizer, chara_classifier=chara_classifier, stock_classifier=stock_classifier)
 
 
 fps = 30
 
-cap = cv2.VideoCapture(2)
+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)