From f2ddc963342e944214dc9d459a1e1dc4bf599aa1 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Mon, 8 Apr 2024 12:25:02 +0200
Subject: [PATCH] Add `IS_JETSON` constant (#9852)

Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
---
 ultralytics/utils/__init__.py | 36 +++++++++++++++++++++++++++++------
 1 file changed, 30 insertions(+), 6 deletions(-)

diff --git a/ultralytics/utils/__init__.py b/ultralytics/utils/__init__.py
index cda2aaae..3a270624 100644
--- a/ultralytics/utils/__init__.py
+++ b/ultralytics/utils/__init__.py
@@ -406,6 +406,20 @@ DEFAULT_CFG_KEYS = DEFAULT_CFG_DICT.keys()
 DEFAULT_CFG = IterableSimpleNamespace(**DEFAULT_CFG_DICT)
 
 
+def read_device_model() -> str:
+    """
+    Reads the device model information from the system and caches it for quick access. Used by is_jetson() and
+    is_raspberrypi().
+
+    Returns:
+        (str): Model file contents if read successfully or empty string otherwise.
+    """
+    with contextlib.suppress(Exception):
+        with open("/proc/device-tree/model") as f:
+            return f.read()
+    return ""
+
+
 def is_ubuntu() -> bool:
     """
     Check if the OS is Ubuntu.
@@ -473,10 +487,18 @@ def is_raspberrypi() -> bool:
     Returns:
         (bool): True if running on a Raspberry Pi, False otherwise.
     """
-    with contextlib.suppress(Exception):
-        with open("/proc/device-tree/model") as f:
-            return "Raspberry Pi" in f.read()
-    return False
+    return "Raspberry Pi" in PROC_DEVICE_MODEL
+
+
+def is_jetson() -> bool:
+    """
+    Determines if the Python environment is running on a Jetson Nano or Jetson Orin device by checking the device model
+    information.
+
+    Returns:
+        (bool): True if running on a Jetson Nano or Jetson Orin, False otherwise.
+    """
+    return "Jetson" in PROC_DEVICE_MODEL  # i.e. "Jetson Nano" or "Jetson Orin"
 
 
 def is_online() -> bool:
@@ -658,9 +680,11 @@ def get_user_config_dir(sub_dir="Ultralytics"):
 
 
 # Define constants (required below)
+PROC_DEVICE_MODEL = read_device_model()  # is_jetson() and is_raspberrypi() depend on this constant
 ONLINE = is_online()
 IS_COLAB = is_colab()
 IS_DOCKER = is_docker()
+IS_JETSON = is_jetson()
 IS_JUPYTER = is_jupyter()
 IS_KAGGLE = is_kaggle()
 IS_PIP_PACKAGE = is_pip_package()
@@ -696,8 +720,8 @@ def colorstr(*input):
         (str): The input string wrapped with ANSI escape codes for the specified color and style.
 
     Examples:
-        >>> colorstr('blue', 'bold', 'hello world')
-        >>> '\033[34m\033[1mhello world\033[0m'
+        >>> colorstr("blue", "bold", "hello world")
+        >>> "\033[34m\033[1mhello world\033[0m"
     """
     *args, string = input if len(input) > 1 else ("blue", "bold", input[0])  # color arguments, string
     colors = {