ultralytics 8.1.38__py3-none-any.whl → 8.1.40__py3-none-any.whl

ultralytics/solutions/object_counter.py

@@ -43,16 +43,19 @@ class ObjectCounter:
         # Object counting Information
         self.in_counts = 0
         self.out_counts = 0
-        self.counting_dict = {}
+        self.count_ids = []
+        self.class_wise_count = {}
         self.count_txt_thickness = 0
-        self.count_txt_color = (0, 0, 0)
-        self.count_color = (255, 255, 255)
+        self.count_txt_color = (255, 255, 255)
+        self.line_color = (255, 255, 255)
+        self.cls_txtdisplay_gap = 50
+        self.fontsize = 0.6
 
         # Tracks info
         self.track_history = defaultdict(list)
         self.track_thickness = 2
         self.draw_tracks = False
-        self.track_color = (0, 255, 0)
+        self.track_color = None
 
         # Check if environment support imshow
         self.env_check = check_imshow(warn=True)
@@ -68,12 +71,14 @@ class ObjectCounter:
         view_in_counts=True,
         view_out_counts=True,
         draw_tracks=False,
-        count_txt_thickness=2,
-        count_txt_color=(0, 0, 0),
-        count_color=(255, 255, 255),
-        track_color=(0, 255, 0),
+        count_txt_thickness=3,
+        count_txt_color=(255, 255, 255),
+        fontsize=0.8,
+        line_color=(255, 255, 255),
+        track_color=None,
         region_thickness=5,
         line_dist_thresh=15,
+        cls_txtdisplay_gap=50,
     ):
         """
         Configures the Counter's image, bounding box line thickness, and counting region points.
@@ -89,11 +94,13 @@ class ObjectCounter:
             draw_tracks (Bool): draw tracks
             count_txt_thickness (int): Text thickness for object counting display
             count_txt_color (RGB color): count text color value
-            count_color (RGB color): count text background color value
+            fontsize (float): Text display font size
+            line_color (RGB color): count highlighter line color
             count_reg_color (RGB color): Color of object counting region
             track_color (RGB color): color for tracks
             region_thickness (int): Object counting Region thickness
             line_dist_thresh (int): Euclidean Distance threshold for line counter
+            cls_txtdisplay_gap (int): Display gap between each class count
         """
         self.tf = line_thickness
         self.view_img = view_img
@@ -108,7 +115,7 @@ class ObjectCounter:
             self.reg_pts = reg_pts
             self.counting_region = LineString(self.reg_pts)
         elif len(reg_pts) >= 3:
-            print("Region Counter Initiated.")
+            print("Polygon Counter Initiated.")
             self.reg_pts = reg_pts
             self.counting_region = Polygon(self.reg_pts)
         else:
@@ -120,10 +127,12 @@ class ObjectCounter:
         self.track_color = track_color
         self.count_txt_thickness = count_txt_thickness
         self.count_txt_color = count_txt_color
-        self.count_color = count_color
+        self.fontsize = fontsize
+        self.line_color = line_color
         self.region_color = count_reg_color
         self.region_thickness = region_thickness
         self.line_dist_thresh = line_dist_thresh
+        self.cls_txtdisplay_gap = cls_txtdisplay_gap
 
     def mouse_event_for_region(self, event, x, y, flags, params):
         """
@@ -171,7 +180,13 @@ class ObjectCounter:
             # Extract tracks
             for box, track_id, cls in zip(boxes, track_ids, clss):
                 # Draw bounding box
-                self.annotator.box_label(box, label=f"{track_id}:{self.names[cls]}", color=colors(int(track_id), True))
+                self.annotator.box_label(box, label=f"{self.names[cls]}#{track_id}", color=colors(int(track_id), True))
+
+                # Store class info
+                if self.names[cls] not in self.class_wise_count:
+                    if len(self.names[cls]) > 5:
+                        self.names[cls] = self.names[cls][:5]
+                    self.class_wise_count[self.names[cls]] = {"in": 0, "out": 0}
 
                 # Draw Tracks
                 track_line = self.track_history[track_id]
@@ -182,68 +197,68 @@ class ObjectCounter:
                 # Draw track trails
                 if self.draw_tracks:
                     self.annotator.draw_centroid_and_tracks(
-                        track_line, color=self.track_color, track_thickness=self.track_thickness
+                        track_line,
+                        color=self.track_color if self.track_color else colors(int(track_id), True),
+                        track_thickness=self.track_thickness,
                     )
 
                 prev_position = self.track_history[track_id][-2] if len(self.track_history[track_id]) > 1 else None
-                centroid = Point((box[:2] + box[2:]) / 2)
-
-                # Count objects
-                if len(self.reg_pts) >= 3:  # any polygon
-                    is_inside = self.counting_region.contains(centroid)
-                    current_position = "in" if is_inside else "out"
 
-                    if prev_position is not None:
-                        if self.counting_dict[track_id] != current_position and is_inside:
-                            self.in_counts += 1
-                            self.counting_dict[track_id] = "in"
-                        elif self.counting_dict[track_id] != current_position and not is_inside:
-                            self.out_counts += 1
-                            self.counting_dict[track_id] = "out"
-                        else:
-                            self.counting_dict[track_id] = current_position
+                # Count objects in any polygon
+                if len(self.reg_pts) >= 3:
+                    is_inside = self.counting_region.contains(Point(track_line[-1]))
 
-                    else:
-                        self.counting_dict[track_id] = current_position
+                    if prev_position is not None and is_inside and track_id not in self.count_ids:
+                        self.count_ids.append(track_id)
 
-                elif len(self.reg_pts) == 2:
-                    if prev_position is not None:
-                        is_inside = (box[0] - prev_position[0]) * (
-                            self.counting_region.centroid.x - prev_position[0]
-                        ) > 0
-                        current_position = "in" if is_inside else "out"
-
-                        if self.counting_dict[track_id] != current_position and is_inside:
+                        if (box[0] - prev_position[0]) * (self.counting_region.centroid.x - prev_position[0]) > 0:
                             self.in_counts += 1
-                            self.counting_dict[track_id] = "in"
-                        elif self.counting_dict[track_id] != current_position and not is_inside:
-                            self.out_counts += 1
-                            self.counting_dict[track_id] = "out"
+                            self.class_wise_count[self.names[cls]]["in"] += 1
                         else:
-                            self.counting_dict[track_id] = current_position
-                    else:
-                        self.counting_dict[track_id] = None
-
-        incount_label = f"In Count : {self.in_counts}"
-        outcount_label = f"OutCount : {self.out_counts}"
-
-        # Display counts based on user choice
-        counts_label = None
-        if not self.view_in_counts and not self.view_out_counts:
-            counts_label = None
-        elif not self.view_in_counts:
-            counts_label = outcount_label
-        elif not self.view_out_counts:
-            counts_label = incount_label
-        else:
-            counts_label = f"{incount_label} {outcount_label}"
+                            self.out_counts += 1
+                            self.class_wise_count[self.names[cls]]["out"] += 1
 
-        if counts_label is not None:
-            self.annotator.count_labels(
-                counts=counts_label,
-                count_txt_size=self.count_txt_thickness,
+                # Count objects using line
+                elif len(self.reg_pts) == 2:
+                    is_inside = (box[0] - prev_position[0]) * (self.counting_region.centroid.x - prev_position[0]) > 0
+
+                    if prev_position is not None and is_inside and track_id not in self.count_ids:
+                        distance = Point(track_line[-1]).distance(self.counting_region)
+
+                        if distance < self.line_dist_thresh and track_id not in self.count_ids:
+                            self.count_ids.append(track_id)
+
+                            if (box[0] - prev_position[0]) * (self.counting_region.centroid.x - prev_position[0]) > 0:
+                                self.in_counts += 1
+                                self.class_wise_count[self.names[cls]]["in"] += 1
+                            else:
+                                self.out_counts += 1
+                                self.class_wise_count[self.names[cls]]["out"] += 1
+
+        label = "Ultralytics Analytics \t"
+
+        for key, value in self.class_wise_count.items():
+            if value["in"] != 0 or value["out"] != 0:
+                if not self.view_in_counts and not self.view_out_counts:
+                    label = None
+                elif not self.view_in_counts:
+                    label += f"{str.capitalize(key)}: IN {value['in']} \t"
+                elif not self.view_out_counts:
+                    label += f"{str.capitalize(key)}: OUT {value['out']} \t"
+                else:
+                    label += f"{str.capitalize(key)}: IN {value['in']} OUT {value['out']} \t"
+
+        label = label.rstrip()
+        label = label.split("\t")
+
+        if label is not None:
+            self.annotator.display_counts(
+                counts=label,
+                tf=self.count_txt_thickness,
+                fontScale=self.fontsize,
                 txt_color=self.count_txt_color,
-                color=self.count_color,
+                line_color=self.line_color,
+                classwise_txtgap=self.cls_txtdisplay_gap,
             )
 
     def display_frames(self):

ultralytics/trackers/byte_tracker.py

@@ -47,7 +47,7 @@ class STrack(BaseTrack):
         """Initialize new STrack instance."""
         super().__init__()
         # xywh+idx or xywha+idx
-        assert len(xywh) in [5, 6], f"expected 5 or 6 values but got {len(xywh)}"
+        assert len(xywh) in {5, 6}, f"expected 5 or 6 values but got {len(xywh)}"
         self._tlwh = np.asarray(xywh2ltwh(xywh[:4]), dtype=np.float32)
         self.kalman_filter = None
         self.mean, self.covariance = None, None

ultralytics/trackers/track.py

@@ -31,7 +31,7 @@ def on_predict_start(predictor: object, persist: bool = False) -> None:
     tracker = check_yaml(predictor.args.tracker)
     cfg = IterableSimpleNamespace(**yaml_load(tracker))
 
-    if cfg.tracker_type not in ["bytetrack", "botsort"]:
+    if cfg.tracker_type not in {"bytetrack", "botsort"}:
         raise AssertionError(f"Only 'bytetrack' and 'botsort' are supported for now, but got '{cfg.tracker_type}'")
 
     trackers = []

ultralytics/trackers/utils/gmc.py

@@ -94,7 +94,7 @@ class GMC:
             array([[1, 2, 3],
                    [4, 5, 6]])
         """
-        if self.method in ["orb", "sift"]:
+        if self.method in {"orb", "sift"}:
             return self.applyFeatures(raw_frame, detections)
         elif self.method == "ecc":
             return self.applyEcc(raw_frame)

ultralytics/utils/__init__.py

@@ -41,7 +41,7 @@ VERBOSE = str(os.getenv("YOLO_VERBOSE", True)).lower() == "true"  # global verbo
 TQDM_BAR_FORMAT = "{l_bar}{bar:10}{r_bar}" if VERBOSE else None  # tqdm bar format
 LOGGING_NAME = "ultralytics"
 MACOS, LINUX, WINDOWS = (platform.system() == x for x in ["Darwin", "Linux", "Windows"])  # environment booleans
-ARM64 = platform.machine() in ("arm64", "aarch64")  # ARM64 booleans
+ARM64 = platform.machine() in {"arm64", "aarch64"}  # ARM64 booleans
 HELP_MSG = """
     Usage examples for running YOLOv8:
 
@@ -359,7 +359,7 @@ def yaml_load(file="data.yaml", append_filename=False):
     Returns:
         (dict): YAML data and file name.
     """
-    assert Path(file).suffix in (".yaml", ".yml"), f"Attempting to load non-YAML file {file} with yaml_load()"
+    assert Path(file).suffix in {".yaml", ".yml"}, f"Attempting to load non-YAML file {file} with yaml_load()"
     with open(file, errors="ignore", encoding="utf-8") as f:
         s = f.read()  # string
 
@@ -866,7 +866,7 @@ def set_sentry():
         """
         if "exc_info" in hint:
             exc_type, exc_value, tb = hint["exc_info"]
-            if exc_type in (KeyboardInterrupt, FileNotFoundError) or "out of memory" in str(exc_value):
+            if exc_type in {KeyboardInterrupt, FileNotFoundError} or "out of memory" in str(exc_value):
                 return None  # do not send event
 
         event["tags"] = {
@@ -879,7 +879,7 @@ def set_sentry():
 
     if (
         SETTINGS["sync"]
-        and RANK in (-1, 0)
+        and RANK in {-1, 0}
         and Path(ARGV[0]).name == "yolo"
         and not TESTS_RUNNING
         and ONLINE

ultralytics/utils/benchmarks.py

@@ -115,7 +115,7 @@ def benchmark(
 
             # Predict
             assert model.task != "pose" or i != 7, "GraphDef Pose inference is not supported"
-            assert i not in (9, 10), "inference not supported"  # Edge TPU and TF.js are unsupported
+            assert i not in {9, 10}, "inference not supported"  # Edge TPU and TF.js are unsupported
             assert i != 5 or platform.system() == "Darwin", "inference only supported on macOS>=10.13"  # CoreML
             exported_model.predict(ASSETS / "bus.jpg", imgsz=imgsz, device=device, half=half)
 
@@ -220,7 +220,7 @@ class ProfileModels:
         output = []
         for file in files:
             engine_file = file.with_suffix(".engine")
-            if file.suffix in (".pt", ".yaml", ".yml"):
+            if file.suffix in {".pt", ".yaml", ".yml"}:
                 model = YOLO(str(file))
                 model.fuse()  # to report correct params and GFLOPs in model.info()
                 model_info = model.info()

ultralytics/utils/callbacks/comet.py

@@ -71,7 +71,7 @@ def _get_experiment_type(mode, project_name):
 
 def _create_experiment(args):
     """Ensures that the experiment object is only created in a single process during distributed training."""
-    if RANK not in (-1, 0):
+    if RANK not in {-1, 0}:
         return
     try:
         comet_mode = _get_comet_mode()

ultralytics/utils/callbacks/mlflow.py

@@ -108,7 +108,7 @@ def on_train_end(trainer):
         for f in trainer.save_dir.glob("*"):  # log all other files in save_dir
             if f.suffix in {".png", ".jpg", ".csv", ".pt", ".yaml"}:
                 mlflow.log_artifact(str(f))
-        keep_run_active = os.environ.get("MLFLOW_KEEP_RUN_ACTIVE", "False").lower() in ("true")
+        keep_run_active = os.environ.get("MLFLOW_KEEP_RUN_ACTIVE", "False").lower() == "true"
         if keep_run_active:
             LOGGER.info(f"{PREFIX}mlflow run still alive, remember to close it using mlflow.end_run()")
         else:

ultralytics/utils/checks.py

@@ -237,7 +237,7 @@ def check_version(
             result = False
         elif op == "!=" and c == v:
             result = False
-        elif op in (">=", "") and not (c >= v):  # if no constraint passed assume '>=required'
+        elif op in {">=", ""} and not (c >= v):  # if no constraint passed assume '>=required'
             result = False
         elif op == "<=" and not (c <= v):
             result = False
@@ -500,7 +500,7 @@ def check_file(file, suffix="", download=True, hard=True):
             raise FileNotFoundError(f"'{file}' does not exist")
         elif len(files) > 1 and hard:
             raise FileNotFoundError(f"Multiple files match '{file}', specify exact path: {files}")
-        return files[0] if len(files) else [] if hard else file  # return file
+        return files[0] if len(files) else []  # return file
 
 
 def check_yaml(file, suffix=(".yaml", ".yml"), hard=True):
@@ -632,7 +632,7 @@ def check_amp(model):
         (bool): Returns True if the AMP functionality works correctly with YOLOv8 model, else False.
     """
     device = next(model.parameters()).device  # get model device
-    if device.type in ("cpu", "mps"):
+    if device.type in {"cpu", "mps"}:
         return False  # AMP only used on CUDA devices
 
     def amp_allclose(m, im):

ultralytics/utils/downloads.py

@@ -356,13 +356,13 @@ def safe_download(
                     raise ConnectionError(emojis(f"❌  Download failure for {url}. Retry limit reached.")) from e
                 LOGGER.warning(f"⚠️ Download failure, retrying {i + 1}/{retry} {url}...")
 
-    if unzip and f.exists() and f.suffix in ("", ".zip", ".tar", ".gz"):
+    if unzip and f.exists() and f.suffix in {"", ".zip", ".tar", ".gz"}:
         from zipfile import is_zipfile
 
         unzip_dir = (dir or f.parent).resolve()  # unzip to dir if provided else unzip in place
         if is_zipfile(f):
             unzip_dir = unzip_file(file=f, path=unzip_dir, exist_ok=exist_ok, progress=progress)  # unzip
-        elif f.suffix in (".tar", ".gz"):
+        elif f.suffix in {".tar", ".gz"}:
             LOGGER.info(f"Unzipping {f} to {unzip_dir}...")
             subprocess.run(["tar", "xf" if f.suffix == ".tar" else "xfz", f, "--directory", unzip_dir], check=True)
         if delete:

ultralytics/utils/loss.py

@@ -157,7 +157,7 @@ class v8DetectionLoss:
         self.hyp = h
         self.stride = m.stride  # model strides
         self.nc = m.nc  # number of classes
-        self.no = m.no
+        self.no = m.nc + m.reg_max * 4
         self.reg_max = m.reg_max
         self.device = device
 

ultralytics/utils/metrics.py

@@ -298,7 +298,7 @@ class ConfusionMatrix:
         self.task = task
         self.matrix = np.zeros((nc + 1, nc + 1)) if self.task == "detect" else np.zeros((nc, nc))
         self.nc = nc  # number of classes
-        self.conf = 0.25 if conf in (None, 0.001) else conf  # apply 0.25 if default val conf is passed
+        self.conf = 0.25 if conf in {None, 0.001} else conf  # apply 0.25 if default val conf is passed
         self.iou_thres = iou_thres
 
     def process_cls_preds(self, preds, targets):

ultralytics/utils/plotting.py

@@ -363,35 +363,49 @@ class Annotator:
         cv2.polylines(self.im, [points], isClosed=False, color=color, thickness=track_thickness)
         cv2.circle(self.im, (int(track[-1][0]), int(track[-1][1])), track_thickness * 2, color, -1)
 
-    def count_labels(self, counts=0, count_txt_size=2, color=(255, 255, 255), txt_color=(0, 0, 0)):
+    def display_counts(
+        self, counts=None, tf=2, fontScale=0.6, line_color=(0, 0, 0), txt_color=(255, 255, 255), classwise_txtgap=55
+    ):
         """
-        Plot counts for object counter.
+        Display counts on im0.
 
         Args:
-            counts (int): objects counts value
-            count_txt_size (int): text size for counts display
-            color (tuple): background color of counts display
-            txt_color (tuple): text color of counts display
+            counts (str): objects count data
+            tf (int): text thickness for display
+            fontScale (float): text fontsize for display
+            line_color (RGB Color): counts highlighter color
+            txt_color (RGB Color): counts display color
+            classwise_txtgap (int): Gap between each class count data
         """
-        self.tf = count_txt_size
-        tl = self.tf or round(0.002 * (self.im.shape[0] + self.im.shape[1]) / 2) + 1
+
+        tl = tf or round(0.002 * (self.im.shape[0] + self.im.shape[1]) / 2) + 1
         tf = max(tl - 1, 1)
 
-        # Get text size for in_count and out_count
-        t_size_in = cv2.getTextSize(str(counts), 0, fontScale=tl / 2, thickness=tf)[0]
+        t_sizes = [cv2.getTextSize(str(count), 0, fontScale=0.8, thickness=tf)[0] for count in counts]
 
-        # Calculate positions for counts label
-        text_width = t_size_in[0]
-        text_x = (self.im.shape[1] - text_width) // 2  # Center x-coordinate
-        text_y = t_size_in[1]
+        max_text_width = max([size[0] for size in t_sizes])
+        max_text_height = max([size[1] for size in t_sizes])
 
-        # Create a rounded rectangle for in_count
-        cv2.rectangle(
-            self.im, (text_x - 5, text_y - 5), (text_x + text_width + 7, text_y + t_size_in[1] + 7), color, -1
-        )
-        cv2.putText(
-            self.im, str(counts), (text_x, text_y + t_size_in[1]), 0, tl / 2, txt_color, self.tf, lineType=cv2.LINE_AA
-        )
+        text_x = self.im.shape[1] - max_text_width - 20
+        text_y = classwise_txtgap
+
+        for i, count in enumerate(counts):
+            text_x_pos = text_x
+            text_y_pos = text_y + i * classwise_txtgap
+
+            cv2.putText(
+                self.im,
+                str(count),
+                (text_x_pos, text_y_pos),
+                cv2.FONT_HERSHEY_SIMPLEX,
+                fontScale=fontScale,
+                color=txt_color,
+                thickness=tf,
+                lineType=cv2.LINE_AA,
+            )
+
+            line_y_pos = text_y_pos + max_text_height + 5
+            cv2.line(self.im, (text_x_pos, line_y_pos), (text_x_pos + max_text_width, line_y_pos), line_color, tf)
 
     @staticmethod
     def estimate_pose_angle(a, b, c):
@@ -890,7 +904,7 @@ def plot_results(file="path/to/results.csv", dir="", segment=False, pose=False,
                 ax[i].plot(x, y, marker=".", label=f.stem, linewidth=2, markersize=8)  # actual results
                 ax[i].plot(x, gaussian_filter1d(y, sigma=3), ":", label="smooth", linewidth=2)  # smoothing line
                 ax[i].set_title(s[j], fontsize=12)
-                # if j in [8, 9, 10]:  # share train and val loss y axes
+                # if j in {8, 9, 10}:  # share train and val loss y axes
                 #     ax[i].get_shared_y_axes().join(ax[i], ax[i - 5])
         except Exception as e:
             LOGGER.warning(f"WARNING: Plotting error for {f}: {e}")

ultralytics/utils/torch_utils.py

@@ -37,7 +37,7 @@ TORCHVISION_0_13 = check_version(torchvision.__version__, "0.13.0")
 def torch_distributed_zero_first(local_rank: int):
     """Decorator to make all processes in distributed training wait for each local_master to do something."""
     initialized = torch.distributed.is_available() and torch.distributed.is_initialized()
-    if initialized and local_rank not in (-1, 0):
+    if initialized and local_rank not in {-1, 0}:
         dist.barrier(device_ids=[local_rank])
     yield
     if initialized and local_rank == 0:
@@ -109,7 +109,7 @@ def select_device(device="", batch=0, newline=False, verbose=True):
     for remove in "cuda:", "none", "(", ")", "[", "]", "'", " ":
         device = device.replace(remove, "")  # to string, 'cuda:0' -> '0' and '(0, 1)' -> '0,1'
     cpu = device == "cpu"
-    mps = device in ("mps", "mps:0")  # Apple Metal Performance Shaders (MPS)
+    mps = device in {"mps", "mps:0"}  # Apple Metal Performance Shaders (MPS)
     if cpu or mps:
         os.environ["CUDA_VISIBLE_DEVICES"] = "-1"  # force torch.cuda.is_available() = False
     elif device:  # non-cpu device requested
@@ -347,7 +347,7 @@ def initialize_weights(model):
         elif t is nn.BatchNorm2d:
             m.eps = 1e-3
             m.momentum = 0.03
-        elif t in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU]:
+        elif t in {nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU}:
             m.inplace = True
 
 
@@ -505,6 +505,20 @@ def strip_optimizer(f: Union[str, Path] = "best.pt", s: str = "") -> None:
     LOGGER.info(f"Optimizer stripped from {f},{f' saved as {s},' if s else ''} {mb:.1f}MB")
 
 
+def convert_optimizer_state_dict_to_fp16(state_dict):
+    """
+    Converts the state_dict of a given optimizer to FP16, focusing on the 'state' key for tensor conversions.
+
+    This method aims to reduce storage size without altering 'param_groups' as they contain non-tensor data.
+    """
+    for state in state_dict["state"].values():
+        for k, v in state.items():
+            if isinstance(v, torch.Tensor) and v.dtype is torch.float32:
+                state[k] = v.half()
+
+    return state_dict
+
+
 def profile(input, ops, n=10, device=None):
     """
     Ultralytics speed, memory and FLOPs profiler.

{ultralytics-8.1.38.dist-info → ultralytics-8.1.40.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ultralytics
-Version: 8.1.38
+Version: 8.1.40
 Summary: Ultralytics YOLOv8 for SOTA object detection, multi-object tracking, instance segmentation, pose estimation and image classification.
 Author: Glenn Jocher, Ayush Chaurasia, Jing Qiu
 Maintainer: Glenn Jocher, Ayush Chaurasia, Jing Qiu