setup automated tests

.github/workflows/py.yml

@@ -0,0 +1,23 @@
+name: py
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: 3.x
+          cache: "pip"
+      - name: Install Dependencies
+        run: |
+          pip install -e .
+          pip install -r requirements.txt
+      - run: sh download_ml_models.sh
+      - run: python -m pytest tests --isort --black

main.py

@@ -14,7 +14,7 @@ async def fake_video_streamer():
     vid = cv2.VideoCapture(0)
 
     try:
-        with open("models/face.svg", "r") as model_file:
+        with open("tests/models/face.svg", "r") as model_file:
             tracking = Tracking()
             model = Model(model_file.read())
             tracking.set_model(model)

preview_image.py

@@ -4,34 +4,35 @@
 
 from src.svg.model import Model
 from src.tracking.tracking import Tracking
-from src.utils import overlay_webcam
+from tests.utils import overlay_webcam
 
 if __name__ == "__main__":
-    with open("models/face.svg", "r") as model_file:
-        tracking = Tracking()
+    with open("tests/models/face.svg", "r") as model_file:
         model = Model(model_file.read())
-        tracking.set_model(model)
 
-        frame = cv2.imread("tests/test_normal.jpeg")
-        frame = cv2.imread("tests/test_irises.jpeg")
-        # frame = cv2.imread("tests/test_eyes_mouth.jpeg")
-        # frame = cv2.imread("tests/test_head_tilted_2.jpeg")
-        # frame = cv2.imread("tests/test_lowlight.jpeg")
-        # frame = cv2.imread("tests/test_lowlight_2.jpeg")
+    tracking = Tracking()
+    tracking.set_model(model)
 
-        tracking.process(frame)
+    frame = cv2.imread("tests/images/test_normal.jpeg")
+    frame = cv2.imread("tests/images/test_irises.jpeg")
+    # frame = cv2.imread("tests/images/test_eyes_mouth.jpeg")
+    # frame = cv2.imread("tests/images/test_head_tilted_2.jpeg")
+    # frame = cv2.imread("tests/images/test_lowlight.jpeg")
+    # frame = cv2.imread("tests/images/test_lowlight_2.jpeg")
 
-        while tracking.last_frame_ms is None:
-            pass
+    tracking.process(frame)
 
-        image_bytes = resvg_python.svg_to_png(model.tostring())
-        image_decoded = cv2.imdecode(np.array(bytearray(image_bytes)), cv2.IMREAD_COLOR)
+    while tracking.last_frame_ms is None:
+        pass
 
-        final_image = overlay_webcam(frame, image_decoded)
+    image_bytes = resvg_python.svg_to_png(model.tostring())
+    image_decoded = cv2.imdecode(np.array(bytearray(image_bytes)), cv2.IMREAD_COLOR)
 
-        while True:
-            cv2.imshow("rein", final_image)
+    final_image = overlay_webcam(frame, image_decoded)
 
-            if cv2.waitKey(1) & 0xFF == ord("q"):
-                cv2.destroyAllWindows()
-                break
+    while True:
+        cv2.imshow("rein", final_image)
+
+        if cv2.waitKey(1) & 0xFF == ord("q"):
+            cv2.destroyAllWindows()
+            break

preview_webcam.py

@@ -4,32 +4,31 @@
 
 from src.svg.model import Model
 from src.tracking.tracking import Tracking
-from src.utils import overlay_webcam
+from tests.utils import overlay_webcam
 
 if __name__ == "__main__":
-    with open("models/face.svg", "r") as model_file:
-        tracking = Tracking()
-        vid = cv2.VideoCapture(0)
+    with open("tests/models/face.svg", "r") as model_file:
         model = Model(model_file.read())
-        tracking.set_model(model)
 
-        while True:
-            ret, frame = vid.read()
+    tracking = Tracking()
+    vid = cv2.VideoCapture(0)
+    tracking.set_model(model)
 
-            if not ret:
-                break
+    while True:
+        ret, frame = vid.read()
 
-            tracking.process(frame)
-            image_bytes = resvg_python.svg_to_png(model.tostring())
-            image_decoded = cv2.imdecode(
-                np.array(bytearray(image_bytes)), cv2.IMREAD_COLOR
-            )
+        if not ret:
+            break
 
-            final_image = overlay_webcam(frame, image_decoded)
+        tracking.process(frame)
+        image_bytes = resvg_python.svg_to_png(model.tostring())
+        image_decoded = cv2.imdecode(np.array(bytearray(image_bytes)), cv2.IMREAD_COLOR)
 
-            cv2.imshow("rein", final_image)
+        final_image = overlay_webcam(frame, image_decoded)
 
-            if cv2.waitKey(1) & 0xFF == ord("q"):
-                vid.release()
-                cv2.destroyAllWindows()
-                break
+        cv2.imshow("rein", final_image)
+
+        if cv2.waitKey(1) & 0xFF == ord("q"):
+            vid.release()
+            cv2.destroyAllWindows()
+            break

requirements.txt

@@ -5,4 +5,7 @@ isort
 black
 fastapi
 uvicorn
-resvg_python
+resvg_python
+pytest
+pytest-black
+pytest-isort

src/tracking/tracking.py

@@ -5,6 +5,7 @@
 
 from src.svg.model import Model
 from src.tracking.filters import OneEuroFilter
+from src.tracking.utils import optimize_image
 
 BaseOptions = mp.tasks.BaseOptions
 FaceLandmarker = mp.tasks.vision.FaceLandmarker
@@ -68,7 +69,9 @@ def __filter__(self, key, n):
 
     def process(self, image):
         frame_timestamp_ms = int(time.time() * 1000)
-        mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=image)
+        mp_image = mp.Image(
+            image_format=mp.ImageFormat.SRGB, data=optimize_image(image)
+        )
         self.__face_landmarker.detect_async(mp_image, frame_timestamp_ms)
 
     def __process_callback__(

src/tracking/utils.py

@@ -0,0 +1,15 @@
+import cv2
+
+
+def optimize_image(image, width=480, brightness=50):
+    if image.shape[1] == width:
+        return image
+
+    aspect_ratio = float(image.shape[0]) / float(image.shape[1])
+
+    height = round(width * aspect_ratio)
+
+    resized_image = cv2.resize(image, (width, height))
+    brightened_image = cv2.convertScaleAbs(resized_image, alpha=1, beta=brightness)
+
+    return brightened_image

tests/test_apply.py

@@ -0,0 +1,61 @@
+import os
+from glob import glob
+
+import cv2
+import numpy as np
+import pytest
+import resvg_python
+
+from src.svg.model import Model
+from src.tracking.tracking import Tracking
+from tests.utils import compare_img, overlay_webcam
+
+with open("tests/models/face.svg") as file:
+    face_model_file = file.read()
+
+with open("tests/models/normal.svg") as file:
+    normal_model_file = file.read()
+
+
+@pytest.mark.parametrize("file", glob("tests/images/test_*"))
+def test_basic_model(file):
+    model = Model(face_model_file)
+
+    tracking = Tracking()
+    tracking.set_model(model)
+
+    frame = cv2.imread(file)
+
+    tracking.process(frame)
+
+    while tracking.last_frame_ms is None:
+        pass
+
+    image_bytes = resvg_python.svg_to_png(model.tostring())
+    image_decoded = cv2.imdecode(np.array(bytearray(image_bytes)), cv2.IMREAD_COLOR)
+
+    final_image = overlay_webcam(frame, image_decoded)
+
+    assert compare_img(os.path.basename(file), "basic_model", final_image) < 4.5
+
+
+@pytest.mark.parametrize("file", glob("tests/images/test_*"))
+def test_normal_model(file):
+    model = Model(normal_model_file)
+
+    tracking = Tracking()
+    tracking.set_model(model)
+
+    frame = cv2.imread(file)
+
+    tracking.process(frame)
+
+    while tracking.last_frame_ms is None:
+        pass
+
+    image_bytes = resvg_python.svg_to_png(model.tostring())
+    image_decoded = cv2.imdecode(np.array(bytearray(image_bytes)), cv2.IMREAD_COLOR)
+
+    final_image = overlay_webcam(frame, image_decoded)
+
+    assert compare_img(os.path.basename(file), "normal_model", final_image) < 4.5

src/utils.py → tests/utils.py

@@ -1,3 +1,6 @@
+import os
+from warnings import warn
+
 import cv2
 import numpy as np
 
@@ -31,3 +34,28 @@ def overlay_webcam(webcam_image, output_image, x=5, y=5):
     ] = webcam_image
 
     return output_image
+
+
+def compare_img(filename, suite, b):
+    filepath = f"tests/snapshots/{suite}/{filename}"
+
+    os.makedirs(os.path.dirname(filepath), exist_ok=True)
+
+    if not os.path.exists(filepath):
+        cv2.imwrite(filepath, b)
+    else:
+        a = cv2.imread(filepath)
+
+        a = cv2.cvtColor(a, cv2.COLOR_RGBA2GRAY)
+        b = cv2.cvtColor(b, cv2.COLOR_RGBA2GRAY)
+
+        if a.shape != b.shape:
+            raise ValueError([a.shape, b.shape])
+
+        # 'Mean Squared Error' between the two images
+        err = np.sum((a.astype("float") - b.astype("float")) ** 2)
+        err /= float(a.shape[0] * b.shape[1])
+
+        return err
+
+    return 0