Skip to content

Commit

Permalink
Fixed formatting
Browse files Browse the repository at this point in the history
  • Loading branch information
Ethan118 committed Feb 7, 2024
1 parent 435992c commit 73a28a8
Show file tree
Hide file tree
Showing 3 changed files with 18 additions and 14 deletions.
2 changes: 1 addition & 1 deletion main_2024.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
13 changes: 7 additions & 6 deletions tests/model_example/generate_expected.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down Expand Up @@ -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()
Expand All @@ -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]
Expand Down
17 changes: 10 additions & 7 deletions tests/test_detect_target.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,10 +2,10 @@
TODO: PointsAndTime.
"""
import copy
import pathlib

import cv2
import numpy as np
import pathlib
import pytest
import torch

Expand All @@ -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:
Expand All @@ -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])
Expand All @@ -59,6 +61,7 @@ def create_detections(detections_from_file: np.ndarray) -> detections_and_time.D

return detections


@pytest.fixture()
def detector():
"""
Expand Down Expand Up @@ -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()` .
Expand All @@ -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,
Expand Down

0 comments on commit 73a28a8

Please sign in to comment.