From 73a28a88c96cabcfd887b76ff8d1bf77dc5b0460 Mon Sep 17 00:00:00 2001
From: exahn <exahn@uwaterloo.ca>
Date: Wed, 7 Feb 2024 18:17:14 -0500
Subject: [PATCH] Fixed formatting

---
 main_2024.py                             |  2 +-
 tests/model_example/generate_expected.py | 13 +++++++------
 tests/test_detect_target.py              | 17 ++++++++++-------
 3 files changed, 18 insertions(+), 14 deletions(-)

diff --git a/main_2024.py b/main_2024.py
index 3ed8b52e..1b53be8a 100644
--- a/main_2024.py
+++ b/main_2024.py
@@ -171,7 +171,7 @@ def main() -> int:
             print("pitch: " + str(orientation.pitch))
             print("")
 
-        if cv2.waitKey(1) == ord("q"):
+        if cv2.waitKey(1) == ord('q'):
             print("Exiting main loop")
             break
 
diff --git a/tests/model_example/generate_expected.py b/tests/model_example/generate_expected.py
index c239bd16..5dd4fcb2 100644
--- a/tests/model_example/generate_expected.py
+++ b/tests/model_example/generate_expected.py
@@ -2,15 +2,16 @@
 Generates expected output using pretrained default model and images.
 TODO: PointsAndTime
 """
+import pathlib
 
 import cv2
 import numpy as np
 import ultralytics
-import pathlib
 
 
-# Downloaded from: https://github.com/ultralytics/assets/releases
 TEST_PATH = pathlib.Path("tests", "model_example")
+
+# Downloaded from: https://github.com/ultralytics/assets/releases
 MODEL_PATH = pathlib.Path(TEST_PATH, "yolov8s_ultralytics_pretrained_default.pt")
 
 BUS_IMAGE_PATH = pathlib.Path(TEST_PATH, "bus.jpg")
@@ -44,6 +45,10 @@
     image_bus_annotated = results_bus[0].plot(conf=True)
     image_zidane_annotated = results_zidane[0].plot(conf=True)
 
+    # Save image
+    cv2.imwrite(BUS_IMAGE_ANNOTATED_PATH, image_bus_annotated)
+    cv2.imwrite(ZIDANE_IMAGE_ANNOTATED_PATH, image_zidane_annotated)
+
     # Generate expected
     bounding_box_bus = results_bus[0].boxes.xyxy.detach().cpu().numpy()
     bounding_box_zidane = results_zidane[0].boxes.xyxy.detach().cpu().numpy()
@@ -56,10 +61,6 @@
 
     predictions_bus = np.insert(bounding_box_bus, 0, [conf_bus, labels_bus], axis=1)
     predictions_zidane = np.insert(bounding_box_zidane, 0, [conf_zidane, labels_zidane], axis=1)
-
-    # Save image
-    cv2.imwrite(BUS_IMAGE_ANNOTATED_PATH, image_bus_annotated)
-    cv2.imwrite(ZIDANE_IMAGE_ANNOTATED_PATH, image_zidane_annotated)
     
     # Save expected to text file
     # Format: [confidence, label, x1, y1, x2, y2]
diff --git a/tests/test_detect_target.py b/tests/test_detect_target.py
index b3182698..f742e933 100644
--- a/tests/test_detect_target.py
+++ b/tests/test_detect_target.py
@@ -2,10 +2,10 @@
 TODO: PointsAndTime.
 """
 import copy
+import pathlib
 
 import cv2
 import numpy as np
-import pathlib
 import pytest
 import torch
 
@@ -22,8 +22,8 @@
 IMAGE_ZIDANE_PATH =             pathlib.Path("tests", "model_example", "zidane.jpg")
 BOUNDING_BOX_ZIDANE_PATH =      pathlib.Path("tests", "model_example", "bounding_box_zidane.txt")
 
-BOUNDING_BOX_DESIRED_TOLERANCE = 7
-CONFIDENCE_DESIRED_TOLERANCE = 3
+BOUNDING_BOX_DECIMAL_TOLERANCE = 1
+CONFIDENCE_DECIMAL_TOLERANCE = 2
 
 
 def compare_detections(actual: detections_and_time.DetectionsAndTime, expected: detections_and_time.DetectionsAndTime) -> None:
@@ -44,12 +44,14 @@ def compare_detections(actual: detections_and_time.DetectionsAndTime, expected:
         np.testing.assert_almost_equal(actual_detection.x2, expected_detection.x2, decimal=BOUNDING_BOX_DESIRED_TOLERANCE)
         np.testing.assert_almost_equal(actual_detection.y2, expected_detection.y2, decimal=BOUNDING_BOX_DESIRED_TOLERANCE)
 
+
 def create_detections(detections_from_file: np.ndarray) -> detections_and_time.DetectionsAndTime:
     """
     Create DetectionsAndTime from expected.
+    Format: [confidence, label, x1, y1, x2, y2]
     """
-    detections = detections_and_time.DetectionsAndTime(0)
     assert detections_from_file.shape[1] == 6
+    detections = detections_and_time.DetectionsAndTime(0)
 
     for i in range(0, detections_from_file.shape[0]):
         result, detection = detections_and_time.Detection.create(detections_from_file[i][2:], int(detections_from_file[i][1]), detections_from_file[i][0])
@@ -59,6 +61,7 @@ def create_detections(detections_from_file: np.ndarray) -> detections_and_time.D
 
     return detections
 
+
 @pytest.fixture()
 def detector():
     """
@@ -91,24 +94,25 @@ def image_zidane():
     assert zidane_image is not None
     yield zidane_image
 
+
 @pytest.fixture()
 def expected_bus():
     """
     Load expected bus detections.
-    Format: [confidence, label, x1, y1, x2, y2]
     """
     expected_bus = np.loadtxt(BOUNDING_BOX_BUS_PATH)
     yield create_detections(expected_bus)
 
+
 @pytest.fixture()
 def expected_zidane():
     """
     Load expected Zidane detections.
-    Format: [confidence, label, x1, y1, x2, y2]
     """
     expected_zidane = np.loadtxt(BOUNDING_BOX_ZIDANE_PATH)
     yield create_detections(expected_zidane)
 
+
 class TestDetector:
     """
     Tests `DetectTarget.run()` .
@@ -128,7 +132,6 @@ def test_single_bus_image(self,
         assert actual is not None
 
         compare_detections(actual, expected_bus)
-        
 
     def test_single_zidane_image(self,
                                  detector: detect_target.DetectTarget,