From 80ebba3a5a9aef453da3fbfb27fa56c05e125679 Mon Sep 17 00:00:00 2001 From: Kenneth Chiem Date: Sun, 3 Dec 2023 12:37:46 -0500 Subject: [PATCH] refactor backend for deployment on aws --- .dockerignore | 1 + Dockerfile | 4 +++ src/api/backend.py | 65 +++++++++++++++++++++++++++++++++++----------- src/view/app.py | 58 ++++++++++++++++++++++++++++++++++------- 4 files changed, 104 insertions(+), 24 deletions(-) diff --git a/.dockerignore b/.dockerignore index 0754c9c6..edb0778b 100644 --- a/.dockerignore +++ b/.dockerignore @@ -2,6 +2,7 @@ venv .env __pycache__ yolov8m-pose.pt +# data/ *.mp4 tmp/ .DS_Store diff --git a/Dockerfile b/Dockerfile index 1b1f07b3..94762fa7 100644 --- a/Dockerfile +++ b/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 \ No newline at end of file diff --git a/src/api/backend.py b/src/api/backend.py index 78cf79fb..b4ea94bf 100644 --- a/src/api/backend.py +++ b/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)} diff --git a/src/view/app.py b/src/view/app.py index 1aaa1c8d..0bceeb28 100644 --- a/src/view/app.py +++ b/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,8 +54,8 @@ 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 @@ -57,17 +63,20 @@ def process_video(video_file): # 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: