diff --git a/pyproject.toml b/pyproject.toml index 63206e5..5adaef4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -57,3 +57,7 @@ ipykernel = "^6.25.2" requires = ["poetry-core"] build-backend = "poetry.core.masonry.api" + +[tool.ruff] +# More than the default (88) of `black` to make comments less of a headache +line-length = 120 diff --git a/set-detect-notify/__main__.py b/set-detect-notify/__main__.py index 9f166ec..defefac 100644 --- a/set-detect-notify/__main__.py +++ b/set-detect-notify/__main__.py @@ -50,11 +50,13 @@ def main(): help="The scale to run the detection at, default is 0.25", ) argparser.add_argument( - '--view-scale', - # Set it to the env VIEW_SCALE if it isn't blank, otherwise set it to 0.75 - default=os.environ['VIEW_SCALE'] if 'VIEW_SCALE' in os.environ and os.environ['VIEW_SCALE'] != '' else 0.75, # noqa: E501 - type=float, - help="The scale to view the detection at, default is 0.75", + "--view-scale", + # Set it to the env VIEW_SCALE if it isn't blank, otherwise set it to 0.75 + default=os.environ["VIEW_SCALE"] + if "VIEW_SCALE" in os.environ and os.environ["VIEW_SCALE"] != "" + else 0.75, # noqa: E501 + type=float, + help="The scale to view the detection at, default is 0.75", ) argparser.add_argument( @@ -188,8 +190,10 @@ def main(): # If it isn't, print a warning for obj in args.detect_object: if obj not in object_names: - print(f"Warning: {obj} is not in the list of objects the model can detect!") - + print( + f"Warning: {obj} is not in the list of objects the model can detect!" + ) + for box in r.boxes: # Get the name of the object class_id = r.names[box.cls[0].item()] @@ -205,13 +209,14 @@ def main(): # print("---") # Now do stuff (if conf > 0.5) - if conf < args.confidence_threshold or (class_id not in args.detect_object and args.detect_object != []): + if conf < args.confidence_threshold or ( + class_id not in args.detect_object and args.detect_object != [] + ): # If the confidence is too low # or if the object is not in the list of objects to detect and the list of objects to detect is not empty # then skip this iteration continue - # Add the object to the list of objects to plot plot_boxes.append( { @@ -303,7 +308,7 @@ def main(): run_scale=args.run_scale, view_scale=args.view_scale, ) - + # Display the resulting frame # cv2.imshow("", r) cv2.imshow(f"Video{i}", frame_to_show) diff --git a/set-detect-notify/utils/utils.py b/set-detect-notify/utils/utils.py index 5e58f31..82cac73 100644 --- a/set-detect-notify/utils/utils.py +++ b/set-detect-notify/utils/utils.py @@ -1,17 +1,19 @@ import cv2 import numpy as np + + def plot_label( - # list of dicts with each dict containing a label, x1, y1, x2, y2 - boxes: list = None, - # opencv image - full_frame: np.ndarray = None, - # run_scale is the scale of the image that was used to run the model - # So the coordinates will be scaled up to the view frame size - run_scale: float = None, - # view_scale is the scale of the image, in relation to the full frame - # So the coordinates will be scaled appropriately when coming from run_frame - view_scale: float = None, - font: int = cv2.FONT_HERSHEY_SIMPLEX, + # list of dicts with each dict containing a label, x1, y1, x2, y2 + boxes: list = None, + # opencv image + full_frame: np.ndarray = None, + # run_scale is the scale of the image that was used to run the model + # So the coordinates will be scaled up to the view frame size + run_scale: float = None, + # view_scale is the scale of the image, in relation to the full frame + # So the coordinates will be scaled appropriately when coming from run_frame + view_scale: float = None, + font: int = cv2.FONT_HERSHEY_SIMPLEX, ): view_frame = cv2.resize(full_frame, (0, 0), fx=view_scale, fy=view_scale) for thing in boxes: @@ -19,9 +21,15 @@ def plot_label( # Image view_frame, # Start point - (int(thing["x1"] * (run_scale/view_scale)), int(thing["y1"] * (run_scale/view_scale))), + ( + int(thing["x1"] * (run_scale / view_scale)), + int(thing["y1"] * (run_scale / view_scale)), + ), # End point - (int(thing["x2"] * (run_scale/view_scale)), int(thing["y2"] * (run_scale/view_scale))), + ( + int(thing["x2"] * (run_scale / view_scale)), + int(thing["y2"] * (run_scale / view_scale)), + ), # Color (0, 255, 0), # Thickness @@ -33,7 +41,10 @@ def plot_label( # Text thing["label"], # Origin - (int(thing["x1"] * (run_scale/view_scale)), int(thing["y1"] * (run_scale/view_scale))), + ( + int(thing["x1"] * (run_scale / view_scale)), + int(thing["y1"] * (run_scale / view_scale)), + ), # Font font, # Font Scale @@ -41,6 +52,6 @@ def plot_label( # Color (0, 255, 0), # Thickness - 1 + 1, ) - return view_frame \ No newline at end of file + return view_frame