refactor backend for deployment on aws

.dockerignore

@@ -2,6 +2,7 @@ venv
 .env
 __pycache__
 yolov8m-pose.pt
+# data/
 *.mp4
 tmp/
 .DS_Store

Dockerfile

@@ -7,4 +7,8 @@ COPY . .
 
 RUN pip install -r requirements.txt
 
+RUN apt-get update && apt-get install -y \
+  libgl1-mesa-glx \
+  && rm -rf /var/lib/apt/lists/*
+
 CMD uvicorn src.api.backend:app --reload --host 0.0.0.0 --port 8000

src/api/backend.py

@@ -4,23 +4,23 @@
 import time
 import io
 from fastapi import FastAPI, File, UploadFile
-from fastapi.responses import StreamingResponse
+from fastapi.responses import StreamingResponse, FileResponse
 import pandas as pd
 import boto3
 import os
 import subprocess
 from dotenv import load_dotenv
+import shutil
 
 # Amazon S3 Connection
-load_dotenv()
-access_key = os.getenv("AWS_ACCESS_KEY_ID")
-secret_key = os.getenv("AWS_SECRET_ACCESS_KEY")
+# load_dotenv()
+# access_key = os.getenv("AWS_ACCESS_KEY_ID")
+# secret_key = os.getenv("AWS_SECRET_ACCESS_KEY")
 
-s3 = boto3.client("s3", aws_access_key_id=access_key, aws_secret_access_key=secret_key)
+s3 = boto3.client("s3")
 
 app = FastAPI()
 
-
 # Root
 @app.get("/")
 async def root():
@@ -34,7 +34,7 @@ async def upload_file(video_file: UploadFile = File(...)):
     """
     file_name = time.strftime("%Y%m%d-%H%M%S")
     try:
-        s3.upload_fileobj(video_file.file, "ball-uploads", file_name)
+        s3.upload_fileobj(video_file.file, "hooptracker-uploads", file_name + ".mp4")
         return {"message": file_name, "status": "success"}
     except Exception as ex:
         return {"error": str(ex)}
@@ -46,26 +46,61 @@ async def process_file(file_name: str):
     runs main to process file
     TODO change from local to cloud
     """
+
     try:
-        command = ["python", "src/main.py", "--source", file_name]
+
+        # create a data directory if it does not exist already
+        os.makedirs("data", exist_ok=True)
+
+        # create a tmp directory if it does not exist already
+        os.makedirs("tmp", exist_ok=True)
+
+        new_path = "data/" + file_name + ".mp4"
+
+        s3.download_file("hooptracker-uploads", file_name + ".mp4", new_path)
+        command = ["python", "src/main.py", "--source", new_path]
         subprocess.run(command)
+        results_path = "results-" + file_name + ".txt"
+        court_reenc_path = "court_video_reenc-" + file_name + ".mp4"
+        s3.upload_file("tmp/results.txt", "hooptracker-uploads", results_path)
+        s3.upload_file("tmp/court_video_reenc.mp4", "hooptracker-uploads", court_reenc_path)
         return {"message": f"successfully processed {file_name}", "status": "success"}
     except Exception as ex:
         return {"error": str(ex)}
 
 
-# TODO Not being used RN
 @app.get("/download/{file_name}")
-async def download_file(file_name: str, download_path: str):
+async def download_file(file_name: str):
     """
     Download video with file_name to download_path
     """
     try:
-        s3.download_file("ball-uploads", file_name, download_path)
-        return {
-            "message": f"successfully downloaded {file_name} to {download_path}",
-            "status": "success",
-        }
+        print(file_name)
+        os.makedirs("tmp", exist_ok=True)
+        download_path_video = "tmp/court_video_reenc-" + file_name + ".mp4"
+        download_path_txt = "tmp/results-" + file_name + ".txt"
+
+        s3.download_file("hooptracker-uploads", "court_video_reenc-" + file_name + ".mp4", download_path_video)
+        s3.download_file("hooptracker-uploads", "results-" + file_name + ".txt", download_path_txt)
+
+        temp_dir = "temporary"
+        os.makedirs(temp_dir, exist_ok=True)
+
+        print("downloaded files")
+
+        # # copy all the files to the temp directory
+        shutil.copy(download_path_video, temp_dir)
+        shutil.copy(download_path_txt, temp_dir)
+
+        # # zip the files
+        zip_path = "files.zip"
+        shutil.make_archive("files", "zip", temp_dir)
+
+        # clean up temporary directory
+        shutil.rmtree(temp_dir)
+
+        return FileResponse(zip_path, media_type="application/zip", filename="files.zip")
+        # return {"message": "successfully downloaded", "status": "success"}
     except Exception as ex:
         return {"error": str(ex)}
 

src/view/app.py

@@ -6,6 +6,8 @@
 import hydralit_components as hc
 import pandas as pd
 import requests
+import zipfile
+import shutil
 
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
 from main import main
@@ -28,7 +30,7 @@
     st.session_state.user_file = "tmp/user_upload.mp4"
 
 # Backend Connection
-SERVER_URL = "http://127.0.0.1:8000/"
+SERVER_URL = "http://35.171.133.54:8000/"
 
 
 def process_video(video_file):
@@ -37,6 +39,10 @@ def process_video(video_file):
     the processed video name into session state
     Temporarily: stores the processed video into tmp/user_upload.mp4
     """
+
+    # delete any unzipped files
+    shutil.rmtree("unzipped_files", ignore_errors=True)
+
     user_video: str = st.session_state.user_file
     # UPLOAD VIDEO
     if video_file is None:
@@ -48,26 +54,29 @@ def process_video(video_file):
         print("Successfully uploaded file")
         data = r.json()
         st.session_state.upload_name = data.get("message")
-        with open(user_video, "wb") as f:  # TODO is local write; temp fix
-            f.write(video_file.getvalue())
+        # with open(user_video, "wb") as f:  # TODO is local write; temp fix
+        #     f.write(video_file.getvalue())
     else:
         print("Error uploading file")  # TODO make an error handler in frontend
         return False
     st.session_state.is_downloaded = False
 
     # PROCESS VIDEO
     print("User Video", user_video)
+    print("Upload Name", st.session_state.upload_name)
+
     # ASSUME process updates results locally for now TODO
-    r = requests.post(SERVER_URL + "process", params={"file_name": user_video})
+    r = requests.post(SERVER_URL + "process", params={"file_name": st.session_state.upload_name})
     if r.status_code == 200:
         print(r.json().get("message"))
-        with open("tmp/results.txt", "r") as file:
-            st.session_state.result_string = file.read()
-        st.session_state.processed_video = "tmp/court_video_reenc.mp4"
+        # with open("tmp/results.txt", "r") as file:
+        #     st.session_state.result_string = file.read()
+        # st.session_state.processed_video = "tmp/court_video_reenc.mp4"
     else:
         print(f"Error processing file: {r.text}")
         return False
-    return True
+
+    return download_results(st.session_state.upload_name)
 
 def upload(video_file):
     user_video: str = st.session_state.user_file
@@ -102,6 +111,37 @@ def health_check():
         return False
     st.session_state.is_downloaded = False
 
+def download_results(upload_name):
+    r = requests.get(f"{SERVER_URL}download/{upload_name}")
+    if r.status_code == 200:
+        zip_file_path = "downloaded_files.zip"
+        with open(zip_file_path, "wb") as file:
+            file.write(r.content)
+
+        # unzip the files
+        unzip_dir = "unzipped_files"
+        os.makedirs(unzip_dir, exist_ok=True)
+
+        # unzip the files
+        with zipfile.ZipFile(zip_file_path, 'r') as zip_ref:
+            zip_ref.extractall(unzip_dir)
+
+        processed_video_path = os.path.join(unzip_dir, f"court_video_reenc-{upload_name}.mp4")
+        result_string_path = os.path.join(unzip_dir, f"results-{upload_name}.txt")
+
+        with open(result_string_path, "r") as file:
+            st.session_state.result_string = file.read()
+        st.session_state.processed_video = processed_video_path
+
+        # clean up temporary directory
+        if os.path.exists(zip_file_path):
+            os.remove(zip_file_path)
+
+        return True
+    else:
+        print(f"Error downloading file: {r.text}")
+        return False
+
 # Pages
 def main_page():
     """
@@ -139,7 +179,7 @@ def loading_page():
         "",
         hc.Loaders.pulse_bars,
     ):
-        finished = upload(video_file=st.session_state.video_file)
+        finished = process_video(st.session_state.video_file)
         if finished:
             state = 2
         else: