Merge pull request #17 from CornellDataScience/fix-madeshot

Revamped backend

.github/workflows/ci.yml

@@ -27,4 +27,4 @@ jobs:
       run: echo $DATA >> /home/runner/work/Ball-101/Ball-101/.env
     - name: Run CI tests
       run: |
-        python src/processing/shot_detect.py
+        python src/main.py data/training_data.mp4

.gitignore

@@ -1,3 +1,6 @@
 venv
 .env
 __pycache__
+.DS_Store
+.vscode/
+tmp/

README.md

@@ -39,8 +39,7 @@ Enable AWS connection by pasting the .env file into the repo.
 
 Start the server backend by running
 ```
-cd src/api
-uvicorn backend:app --reload
+uvicorn src.api/.backend:app --reload
 ```
 
 Open a new bash terminal and start the frontend by running

requirements.txt

@@ -44,6 +44,9 @@ yapf
 isort==4.3.21
 imageio
 
+# Processing
+scikit-learn
+
 # View
 streamlit>=1.18.1
 hydralit_components>= 1.0.10

src/main.py

@@ -4,12 +4,14 @@
 import yaml
 from modelrunner import ModelRunner
 from processrunner import ProcessRunner
+
 # before main is called:
 # frontend boots up, awaits user to upload a video
 # upload button triggers backend call to upload video to s3
 # fetch the video from cloud and download to tmp/uploaded_video.mp4
 # calls main
 
+
 # load in configs from config.yaml
 # initialise modelrunner and processrunner
 # feed video into the yolo model, pass into modelrunner
@@ -20,7 +22,7 @@ def load_config(path):
     """
     TODO Loads the config yaml file to read in parameters and settings.
     """
-    with open(path, 'r') as file:
+    with open(path, "r") as file:
         config = yaml.safe_load(file)
     return config
 
@@ -31,21 +33,31 @@ def main(video_path):
     Input: Path of the user uploaded video.
     Returns: Results of the processing, in string format TODO change to csv/json?
     """
-    config = load_config('config.yaml')
-    model_vars = config['model_vars']
+    config = load_config("config.yaml")
+    model_vars = config["model_vars"]
 
     modelrunner = ModelRunner(video_path, model_vars)
     modelrunner.run()
     people_output, ball_output = modelrunner.fetch_output()
-    output_video_path = 'tmp/court_video.mp4'
-    output_video_path_reenc = 'tmp/court_video_reenc.mp4'
+    output_video_path = "tmp/court_video.mp4"
+    output_video_path_reenc = "tmp/court_video_reenc.mp4"
 
-    processrunner = ProcessRunner(video_path, people_output, ball_output, output_video_path, 
-                                  output_video_path_reenc)
+    processrunner = ProcessRunner(
+        video_path,
+        people_output,
+        ball_output,
+        output_video_path,
+        output_video_path_reenc,
+    )
     processrunner.run()
     results = processrunner.get_results()
     return results
 
 
-if __name__ == '__main__':
-    main('data/training_data.mp4')
+if __name__ == "__main__":
+    import sys
+
+    if len(sys.argv) <= 1:
+        main("tmp/training_data.mp4")
+    else:
+        main(sys.argv[1])  # Pass the first command-line argument to the main function