switch to streamlit

Dockerfile

@@ -6,11 +6,14 @@ RUN pip install ultralytics
 RUN pip install opencv-python==4.6.0.66
 RUN pip install Pillow==10.3.0
 RUN pip install uvicorn fastapi
+RUN pip install streamlit
 RUN apt update && apt install fonts-dejavu -y
 RUN wget https://github.com/ultralytics/assets/releases/download/v8.2.0/yolov8l-pose.pt
 
 COPY . .
 
-CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]
+# CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860"]
+
+CMD ["streamlit", "run", "app.py"]
 
 # curl -X POST "https://unfinity-yolov8-pose.hf.space/predict-image" -F "[email protected]"

app.py

@@ -1,10 +1,12 @@
 import streamlit as st
-import numpy as np
-from PIL import Image, ImageDraw, ImageFont
+from PIL import Image
 from ultralytics import YOLO
 import torch
 import utils
 
+import utils
+from drawing import draw_keypoints
+
 
 @st.cache_resource()
 def load_model():
@@ -16,49 +18,8 @@ def load_model():
 
 
 def draw_output(image_pil: Image.Image, keypoints: dict):    
-    draw = ImageDraw.Draw(image_pil)
-    line_width = 10
-    font = ImageFont.truetype("DejaVuSerif-Bold.ttf", 70)
-    
-    ear, eye = None, None
-    if keypoints["left_ear"] and keypoints["left_eye"]:
-        ear = keypoints["left_ear"]
-        eye = keypoints["left_eye"]
-    elif keypoints["right_ear"] and keypoints["right_eye"]:
-        ear = keypoints["right_ear"]
-        eye = keypoints["right_eye"]
-    
-    # draw extended left and right eye lines
-    if ear and eye:
-        left_new_point = utils.extend_line(ear, eye, 3)
-        l1 = [ear, left_new_point]
-        draw.line(l1, fill='red', width=line_width)
-        # draw a horizontal line from ear forwards
-        ear = np.array(ear)
-        l1 = np.array(l1)
-        l1_vector = l1[1] - l1[0]
-        x_s = np.sign(l1_vector)[0]
-        length_l1 = np.linalg.norm(l1_vector)
-        p2 = ear + np.array([length_l1*x_s, 0])
-        ear = tuple(ear.tolist())
-        l = [ear, tuple(p2.tolist())]
-        draw.line(l, fill='gray', width=line_width//2)
-        # draw angle
-        angle = utils.calculate_angle_to_horizontal(l1_vector)
-        draw.text(ear, f'{angle:.2f}', fill='red', font=font)
-
-    # draw elbow angles
-    left_elbow_angle, right_elbow_angle = utils.get_elbow_angles(keypoints)
-    if left_elbow_angle:
-        draw.text(keypoints['left_elbow'], f'{left_elbow_angle:.2f}', fill='red', font=font)
-        # draw polyline for left arm
-        draw.line([keypoints['left_shoulder'], keypoints['left_elbow'], keypoints['left_wrist']], fill='blue', width=line_width)
-    if right_elbow_angle:
-        draw.text(keypoints['right_elbow'], f'{right_elbow_angle:.2f}', fill='red', font=font)
-        # draw polyline for right arm
-        draw.line([keypoints['right_shoulder'], keypoints['right_elbow'], keypoints['right_wrist']], fill='blue', width=line_width)
-
-    return image_pil
+    output_image = draw_keypoints(image_pil, keypoints).convert("RGB")
+    return output_image
 
 
 st.title('Pose Estimation App')
@@ -82,16 +43,18 @@ if img is not None:
     # predict
     with st.spinner('Predicting...'):
         model = load_model()
-        pred = model(img)[0]
+        result = model(img)[0]
     st.markdown('**Results:**')
-    keypoints = utils.get_keypoints(pred)
+    keypoints = utils.get_keypoints(result)
     if keypoints is not None:
         img = draw_output(img, keypoints)
         st.image(img, caption='Predicted image', use_column_width=True)
+        # calculate angles
         lea, rea = utils.get_eye_angles(keypoints)
         lba, rba = utils.get_elbow_angles(keypoints)
+        angles  = {'left_eye_angle': lea, 'right_eye_angle': rea, 'left_elbow_angle': lba, 'right_elbow_angle': rba}
         st.write('Angles:')
-        st.json({'left_eye_angle': lea, 'right_eye_angle': rea, 'left_elbow_angle': lba, 'right_elbow_angle': rba})
+        st.json(angles)
         st.write('Raw keypoints:')
         st.json(keypoints)
     else:

start.py

@@ -1,3 +0,0 @@
-import subprocess
-
-subprocess.run("uvicorn api:app --host 0.0.0.0 --port 7860", shell=True)

test.py

@@ -1,7 +0,0 @@
-from fastapi import FastAPI
-
-app = FastAPI()
-
-@app.get("/")
-def read_root():
-    return {"Hello": "World!"}