CPU support, standalone image detection and save results features

Makefile

@@ -10,8 +10,11 @@ build: requirements.txt
 format:
  @$(PYTHON) -m black .
 
-run:
- @$(PYTHON) src/detect.py
+run_img:
+ @$(PYTHON) src/detect.py --image docs/yolo/inference.jpg
+
+run_video:
+ @$(PYTHON) src/detect.py --video docs/yolo/autocross.mp4
 
 lint:
  @$(PYTHON) -m black --diff --check $(FILES)

build_requirements.txt

@@ -1,2 +1,3 @@
 opencv-python
-ultralytics
+ultralytics
+py-cpuinfo

src/camera_perception/common.py

@@ -1,5 +1,6 @@
 import re
-from platform import platform
+import typing
+import platform
 
 import torch
 
@@ -11,14 +12,26 @@ def get_current_os() -> str:
 
 
 def is_linux_os() -> bool:
- return platform().find("Linux") != -1
+ return platform.platform().find("Linux") != -1
 
 
-def processing_unit() -> str:
- return "CUDA" if torch.cuda.is_available() else "CPU"
+def pytorch_info(version) -> typing.Tuple[str, str]:
+ regex = r"([0-9]*.*)\+([a-z]*)"
+ match = re.match(regex, version)
+ version, engine = match.group(1), match.group(2)
+ return version, engine
 
 
-def get_resolution(resolution_env):
+def processing_unit() -> typing.Tuple[str, str]:
+ version, engine = pytorch_info(torch.__version__)
+ if engine == "cu":
+ return "CUDA", version
+ if engine == "cpu":
+ return "CPU", version
+ return engine.upper(), version
+
+
+def get_resolution(resolution_env) -> typing.Tuple[int, int]:
  regex = r"\d+"
  resolution = re.findall(regex, resolution_env)
  return int(resolution[0]), int(resolution[1])
@@ -32,14 +45,15 @@ class Environment:
  def __init__(self, env):
  self.os = get_current_os()
  self.conf = float(env.get("CONFIDENCE", 0.7))
- self.processing_unit = processing_unit()
+ self.processing_unit, self.version = processing_unit()
  self.resolution = get_resolution(env.get("RESOLUTION", "1280x720"))
+ self.img_graphics_format = "jpg"
 
  @staticmethod
  def from_env(env):
  return Environment(env)
 
- def to_info_string(self):
+ def to_info_string(self) -> str:
  return "os: {}, processing unit: {}, detections confidence: {}%, resolution: {}x{}".format(
  self.os,
  self.processing_unit,
@@ -48,10 +62,16 @@ def to_info_string(self):
  self.resolution[1],
  )
 
- def cuda_to_info_string(self):
+ def cuda_to_info_string(self) -> str:
  return (
  f"GPU: {torch.cuda.get_device_name(0)} "
  f"Memory usage -> "
  f"allocated: {round(torch.cuda.memory_allocated(0) / 1024 ** 3, 1)} GB, "
  f"cached: {round(torch.cuda.memory_reserved(0) / 1024 ** 3, 1)} GB."
  )
+
+ def cpu_to_info_string(self) -> str:
+ import cpuinfo
+
+ cpu_info = cpuinfo.get_cpu_info()
+ return f"CPU: {cpu_info['brand_raw']}, Arch: {cpu_info['arch_string_raw']}, Cores: {cpu_info['count']}"

src/camera_perception/main.py

@@ -6,7 +6,7 @@
 import typing
 import cv2
 
-from camera_perception import logs, common
+from camera_perception import logs, common, utils
 from camera_perception.camera import Camera, CameraError
 from cones_detection.cones_detector import ConesDetector
 from cones_detection.landmark import ConeLandmark
@@ -17,9 +17,7 @@ def get_help_epilog():
 Exit codes:
  0 - successful execution
  1 - missing input file
- 2 - not implemented option
  3 - unable to open camera device
- 4 - CUDA not enabled
  any other code indicated unrecoverable error
 
 Environment variables:
@@ -43,25 +41,8 @@ def run_app(
  logger: logging.Logger,
  environment: common.Environment,
 ) -> int:
- def bounding_box():
- cv2.rectangle(
- frame,
- cone.bounding_box.top_left,
- cone.bounding_box.bottom_right,
- cone.color,
- 3,
- )
-
- def label():
- cv2.putText(
- frame,
- f"{cone.name} {round(cone.conf * 100)}%",
- (cone.bounding_box.top_left[0], cone.bounding_box.top_left[1] - 10),
- font,
- 0.5,
- cone.color,
- 1,
- )
+ def gui_exit(window):
+ return cv2.getWindowProperty(window, cv2.WND_PROP_VISIBLE) < 1
 
  if environment.os != "linux":
  logger.warning(
@@ -71,47 +52,115 @@ def label():
  f"cones_detector: {sys.executable} argv: {argv} {environment.to_info_string()}"
  )
 
+ if args.save and os.path.exists(args.output) and os.listdir(args.output):
+ logger.error(f"Output directory {args.output} is not empty")
+ logger.info("App finished with exit code 1")
+ return 1
+
  video = pathlib.Path(args.video) if args.video else None
  if video and not video.exists():
  logger.error(f"Video file '{str(video)}' not exist.")
  logger.info("App finished with exit code 1")
  return 1
 
- if args.image:
- logger.info(f"Not implemented.")
- return 2
+ image = pathlib.Path(args.image) if args.image else None
+ if image and not image.exists():
+ logger.error(f"Image file '{str(image)}' not exist.")
+ logger.info("App finished with exit code 1")
+ return 1
 
  if environment.processing_unit == "CUDA":
  logger.info(environment.cuda_to_info_string())
+ elif environment.processing_unit == "CPU":
+ logger.info(environment.cpu_to_info_string())
  else:
- logger.error("CUDA not enabled.")
- return 4
+ logger.warning(
+ f"Detected unknown PyTorch installation {environment.processing_unit} {environment.version}. "
+ f"Some functionalities may not work correctly"
+ )
 
+ # setup
  font = cv2.FONT_HERSHEY_DUPLEX
- camera_lens = Camera(args.video, environment.resolution)
+ title = "Cones Detection"
+ exit_key = "q"
+ cache = None
+ if args.save:
+ cache = []
+
  cones_detector = ConesDetector()
- try:
- for frame in camera_lens.get_frame():
- detections = cones_detector.detect(image=frame, conf=environment.conf)
- if len(detections) > 0:
- for detection in detections.boxes:
- cone = ConeLandmark(
- cls=int(detection.cls[0]),
- conf=detection.conf[0],
- detection=detection,
- )
- bounding_box()
- label()
-
- cv2.imshow("Cones Detection", frame)
- if cv2.waitKey(1) == ord("q"):
- break
-
- camera_lens.turn_off()
- cv2.destroyAllWindows()
- except CameraError as e:
- logger.error(str(e))
- return 3
+
+ output = pathlib.Path(args.output) if args.output else None
+ if args.save and not output.exists():
+ output.mkdir()
+
+ if video:
+ camera_lens = Camera(args.video, environment.resolution)
+ try:
+ for idx, frame in enumerate(camera_lens.get_frame()):
+ detections = cones_detector.detect(image=frame, conf=environment.conf)
+ if len(detections) > 0:
+ for detection in detections.boxes:
+ cone = ConeLandmark(
+ cls=int(detection.cls[0]),
+ conf=detection.conf[0],
+ detection=detection,
+ )
+ utils.bounding_box(frame, cone)
+ utils.label(frame, cone, font)
+
+ cv2.imshow(title, frame)
+ if args.save:
+ cache.append(frame)
+ if cv2.waitKey(1) == ord(exit_key):
+ break
+ if gui_exit(title):
+ break
+
+ camera_lens.turn_off()
+ cv2.destroyAllWindows()
+ except CameraError as e:
+ logger.error(str(e))
+ return 3
+
+ if image:
+ img_inference = cv2.imread(str(image))
+ inferences = cones_detector.detect(image=img_inference, conf=environment.conf)
+ if len(inferences) > 0:
+ for inference in inferences.boxes:
+ cone = ConeLandmark(
+ cls=int(inference.cls[0]),
+ conf=inference.conf[0],
+ detection=inference,
+ )
+ utils.bounding_box(img_inference, cone)
+ utils.label(img_inference, cone, font)
+
+ cv2.imshow(title, img_inference)
+ if args.save:
+ cv2.imwrite(
+ str(
+ output.joinpath(
+ f"{image.stem}_detection.{environment.img_graphics_format}"
+ )
+ ),
+ img_inference,
+ )
+ while True:
+ if cv2.waitKey(1) == ord(exit_key):
+ break
+ if gui_exit(title):
+ break
+ cv2.destroyAllWindows()
+
+ if args.save and video:
+ logger.info("Saving images...")
+ for idx, frame in enumerate(cache):
+ cv2.imwrite(
+ str(
+ output.joinpath(f"detection{idx}.{environment.img_graphics_format}")
+ ),
+ frame,
+ )
 
  logger.info("App finished with exit code 0")
  return 0
@@ -139,5 +188,16 @@ def main(argv: typing.List[str], logger=None, environment=None) -> int:
  metavar="input_data",
  help="specifies image with environment for cones detection",
  )
+ parser.add_argument(
+ "--save", action="store_true", help="save results from image cones recognition"
+ )
+ parser.add_argument(
+ "-o",
+ "--output",
+ type=str,
+ metavar="output_dir",
+ default="results",
+ help="if '--save' option then specifies directory, where results should be saved. Has to be empty",
+ )
  parser.epilog = get_help_epilog()
  return run_app(parser.parse_args(argv[1:]), argv, logger, environment)

src/camera_perception/utils.py

@@ -0,0 +1,23 @@
+import cv2
+
+
+def bounding_box(image, cone):
+ cv2.rectangle(
+ image,
+ cone.bounding_box.top_left,
+ cone.bounding_box.bottom_right,
+ cone.color,
+ 3,
+ )
+
+
+def label(image, cone, font):
+ cv2.putText(
+ image,
+ f"{cone.name} {round(cone.conf * 100)}%",
+ (cone.bounding_box.top_left[0], cone.bounding_box.top_left[1] - 10),
+ font,
+ 0.5,
+ cone.color,
+ 1,
+ )

sys_check.py

@@ -1,10 +1,41 @@
+from camera_perception.common import pytorch_info
+
+
 # YOLO check
 import ultralytics
 
 ultralytics.checks()
 
-# CUDA check
+
+# CUDA/CPU check
 import torch
+import torchvision
+
+
+version, engine = pytorch_info(torch.__version__)
+version_vision, engine_vision = pytorch_info(torchvision.__version__)
+
+if engine != engine_vision:
+ print(
+ f"torch and torchvision packages are not compatible:\n"
+ f"torch: {engine} {version}\n"
+ f"torchvision: {engine_vision} {version_vision}"
+ )
+ exit(1)
+
+if engine == "cu" and not torch.cuda.is_available():
+ print(
+ f"CUDA is not available but PyTorch {engine.upper()} {version} is installed. "
+ f"Check environment configuration for conflicted packages."
+ )
+ exit(1)
 
 if not torch.cuda.is_available():
- print("CUDA is not available.")
+ print(f"CUDA is not available. PyTorch {engine.upper()} {version} will be used.")
+ exit(0)
+
+if engine == "cpu" and torch.cuda.is_available():
+ print(
+ f"CUDA is available but PyTorch {engine.upper()} {version} is used. "
+ f"Check environment configuration for better performance."
+ )