ultralytics 8.2.18__tar.gz → 8.2.20__tar.gz

{ultralytics-8.2.18 → ultralytics-8.2.20}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: ultralytics
-Version: 8.2.18
+Version: 8.2.20
 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

{ultralytics-8.2.18 → ultralytics-8.2.20}/tests/test_cli.py

@@ -7,7 +7,7 @@ import pytest
 from ultralytics.cfg import TASK2DATA, TASK2MODEL, TASKS
 from ultralytics.utils import ASSETS, WEIGHTS_DIR, checks
 
-from . import CUDA_DEVICE_COUNT, CUDA_IS_AVAILABLE
+from tests import CUDA_DEVICE_COUNT, CUDA_IS_AVAILABLE
 
 # Constants
 TASK_MODEL_DATA = [(task, WEIGHTS_DIR / TASK2MODEL[task], TASK2DATA[task]) for task in TASKS]

{ultralytics-8.2.18 → ultralytics-8.2.20}/tests/test_cuda.py

@@ -10,7 +10,7 @@ from ultralytics import YOLO
 from ultralytics.utils import ASSETS, WEIGHTS_DIR
 from ultralytics.cfg import TASK2DATA, TASK2MODEL, TASKS
 
-from . import CUDA_DEVICE_COUNT, CUDA_IS_AVAILABLE, MODEL, SOURCE
+from tests import CUDA_DEVICE_COUNT, CUDA_IS_AVAILABLE, MODEL, SOURCE
 
 
 def test_checks():

{ultralytics-8.2.18 → ultralytics-8.2.20}/tests/test_engine.py

@@ -9,7 +9,7 @@ from ultralytics.engine.exporter import Exporter
 from ultralytics.models.yolo import classify, detect, segment
 from ultralytics.utils import ASSETS, DEFAULT_CFG, WEIGHTS_DIR
 
-from . import MODEL
+from tests import MODEL
 
 
 def test_func(*args):  # noqa

{ultralytics-8.2.18 → ultralytics-8.2.20}/tests/test_exports.py

@@ -18,8 +18,7 @@ from ultralytics.utils import (
     checks,
 )
 from ultralytics.utils.torch_utils import TORCH_1_9, TORCH_1_13
-
-from . import MODEL, SOURCE
+from tests import MODEL, SOURCE
 
 
 def test_export_torchscript():

{ultralytics-8.2.18 → ultralytics-8.2.20}/tests/test_integrations.py

@@ -12,7 +12,7 @@ from ultralytics import YOLO, download
 from ultralytics.utils import DATASETS_DIR, SETTINGS
 from ultralytics.utils.checks import check_requirements
 
-from . import MODEL, SOURCE, TMP
+from tests import MODEL, SOURCE, TMP
 
 
 @pytest.mark.skipif(not check_requirements("ray", install=False), reason="ray[tune] not installed")

{ultralytics-8.2.18 → ultralytics-8.2.20}/tests/test_python.py

@@ -25,10 +25,9 @@ from ultralytics.utils import (
     Retry,
     checks,
 )
-from ultralytics.utils.downloads import download
+from ultralytics.utils.downloads import download, is_url
 from ultralytics.utils.torch_utils import TORCH_1_9
-
-from . import CFG, IS_TMP_WRITEABLE, MODEL, SOURCE, TMP
+from tests import CFG, IS_TMP_WRITEABLE, MODEL, SOURCE, TMP
 
 
 def test_model_forward():
@@ -102,7 +101,7 @@ def test_predict_img(model_name):
 def test_predict_visualize(model):
     """Test model predict methods with 'visualize=True' arguments."""
     YOLO(WEIGHTS_DIR / model)(SOURCE, imgsz=32, visualize=True)
-    
+
 
 def test_predict_grey_and_4ch():
     """Test YOLO prediction on SOURCE converted to greyscale and 4-channel images."""
@@ -131,6 +130,7 @@ def test_predict_grey_and_4ch():
 
 @pytest.mark.slow
 @pytest.mark.skipif(not ONLINE, reason="environment is offline")
+@pytest.mark.skipif(not is_url("https://youtu.be/G17sBkb38XQ"), reason="YouTube URL issue")
 @Retry(times=3, delay=10)
 def test_youtube():
     """

{ultralytics-8.2.18 → ultralytics-8.2.20}/ultralytics/__init__.py

@@ -1,6 +1,6 @@
 # Ultralytics YOLO 🚀, AGPL-3.0 license
 
-__version__ = "8.2.18"
+__version__ = "8.2.20"
 
 from ultralytics.data.explorer.explorer import Explorer
 from ultralytics.models import RTDETR, SAM, YOLO, YOLOWorld

ultralytics-8.2.20/ultralytics/cfg/datasets/signature.yaml

@@ -0,0 +1,20 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# Signature dataset by Ultralytics
+# Documentation: https://docs.ultralytics.com/datasets/detect/signature/
+# Example usage: yolo train data=signature.yaml
+# parent
+# ├── ultralytics
+# └── datasets
+#     └── signature  ← downloads here (11.2 MB)
+
+# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
+path: ../datasets/signature # dataset root dir
+train: train/images # train images (relative to 'path') 143 images
+val: valid/images # val images (relative to 'path') 35 images
+
+# Classes
+names:
+  0: signature
+
+# Download script/URL (optional)
+download: https://ultralytics.com/assets/signature.zip

ultralytics-8.2.20/ultralytics/cfg/models/v9/yolov9c-seg.yaml

@@ -0,0 +1,38 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# YOLOv9c-seg
+# 654 layers, 27897120 parameters, 159.4 GFLOPs
+
+# parameters
+nc: 80 # number of classes
+
+# gelan backbone
+backbone:
+  - [-1, 1, Conv, [64, 3, 2]] # 0-P1/2
+  - [-1, 1, Conv, [128, 3, 2]] # 1-P2/4
+  - [-1, 1, RepNCSPELAN4, [256, 128, 64, 1]] # 2
+  - [-1, 1, ADown, [256]] # 3-P3/8
+  - [-1, 1, RepNCSPELAN4, [512, 256, 128, 1]] # 4
+  - [-1, 1, ADown, [512]] # 5-P4/16
+  - [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]] # 6
+  - [-1, 1, ADown, [512]] # 7-P5/32
+  - [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]] # 8
+  - [-1, 1, SPPELAN, [512, 256]] # 9
+
+head:
+  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
+  - [[-1, 6], 1, Concat, [1]] # cat backbone P4
+  - [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]] # 12
+
+  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
+  - [[-1, 4], 1, Concat, [1]] # cat backbone P3
+  - [-1, 1, RepNCSPELAN4, [256, 256, 128, 1]] # 15 (P3/8-small)
+
+  - [-1, 1, ADown, [256]]
+  - [[-1, 12], 1, Concat, [1]] # cat head P4
+  - [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]] # 18 (P4/16-medium)
+
+  - [-1, 1, ADown, [512]]
+  - [[-1, 9], 1, Concat, [1]] # cat head P5
+  - [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]] # 21 (P5/32-large)
+
+  - [[15, 18, 21], 1, Segment, [nc, 32, 256]] # Segment(P3, P4, P5)

ultralytics-8.2.20/ultralytics/cfg/models/v9/yolov9c.yaml

@@ -0,0 +1,38 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# YOLOv9c
+# 618 layers, 25590912 parameters, 104.0 GFLOPs
+
+# parameters
+nc: 80 # number of classes
+
+# gelan backbone
+backbone:
+  - [-1, 1, Conv, [64, 3, 2]] # 0-P1/2
+  - [-1, 1, Conv, [128, 3, 2]] # 1-P2/4
+  - [-1, 1, RepNCSPELAN4, [256, 128, 64, 1]] # 2
+  - [-1, 1, ADown, [256]] # 3-P3/8
+  - [-1, 1, RepNCSPELAN4, [512, 256, 128, 1]] # 4
+  - [-1, 1, ADown, [512]] # 5-P4/16
+  - [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]] # 6
+  - [-1, 1, ADown, [512]] # 7-P5/32
+  - [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]] # 8
+  - [-1, 1, SPPELAN, [512, 256]] # 9
+
+head:
+  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
+  - [[-1, 6], 1, Concat, [1]] # cat backbone P4
+  - [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]] # 12
+
+  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
+  - [[-1, 4], 1, Concat, [1]] # cat backbone P3
+  - [-1, 1, RepNCSPELAN4, [256, 256, 128, 1]] # 15 (P3/8-small)
+
+  - [-1, 1, ADown, [256]]
+  - [[-1, 12], 1, Concat, [1]] # cat head P4
+  - [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]] # 18 (P4/16-medium)
+
+  - [-1, 1, ADown, [512]]
+  - [[-1, 9], 1, Concat, [1]] # cat head P5
+  - [-1, 1, RepNCSPELAN4, [512, 512, 256, 1]] # 21 (P5/32-large)
+
+  - [[15, 18, 21], 1, Detect, [nc]] # Detect(P3, P4, P5)

ultralytics-8.2.20/ultralytics/cfg/models/v9/yolov9e-seg.yaml

@@ -0,0 +1,61 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# YOLOv9c-seg
+# 1261 layers, 60512800 parameters, 248.4 GFLOPs
+
+# parameters
+nc: 80 # number of classes
+
+# gelan backbone
+backbone:
+  - [-1, 1, Silence, []]
+  - [-1, 1, Conv, [64, 3, 2]] # 1-P1/2
+  - [-1, 1, Conv, [128, 3, 2]] # 2-P2/4
+  - [-1, 1, RepNCSPELAN4, [256, 128, 64, 2]] # 3
+  - [-1, 1, ADown, [256]] # 4-P3/8
+  - [-1, 1, RepNCSPELAN4, [512, 256, 128, 2]] # 5
+  - [-1, 1, ADown, [512]] # 6-P4/16
+  - [-1, 1, RepNCSPELAN4, [1024, 512, 256, 2]] # 7
+  - [-1, 1, ADown, [1024]] # 8-P5/32
+  - [-1, 1, RepNCSPELAN4, [1024, 512, 256, 2]] # 9
+
+  - [1, 1, CBLinear, [[64]]] # 10
+  - [3, 1, CBLinear, [[64, 128]]] # 11
+  - [5, 1, CBLinear, [[64, 128, 256]]] # 12
+  - [7, 1, CBLinear, [[64, 128, 256, 512]]] # 13
+  - [9, 1, CBLinear, [[64, 128, 256, 512, 1024]]] # 14
+
+  - [0, 1, Conv, [64, 3, 2]] # 15-P1/2
+  - [[10, 11, 12, 13, 14, -1], 1, CBFuse, [[0, 0, 0, 0, 0]]] # 16
+  - [-1, 1, Conv, [128, 3, 2]] # 17-P2/4
+  - [[11, 12, 13, 14, -1], 1, CBFuse, [[1, 1, 1, 1]]] # 18
+  - [-1, 1, RepNCSPELAN4, [256, 128, 64, 2]] # 19
+  - [-1, 1, ADown, [256]] # 20-P3/8
+  - [[12, 13, 14, -1], 1, CBFuse, [[2, 2, 2]]] # 21
+  - [-1, 1, RepNCSPELAN4, [512, 256, 128, 2]] # 22
+  - [-1, 1, ADown, [512]] # 23-P4/16
+  - [[13, 14, -1], 1, CBFuse, [[3, 3]]] # 24
+  - [-1, 1, RepNCSPELAN4, [1024, 512, 256, 2]] # 25
+  - [-1, 1, ADown, [1024]] # 26-P5/32
+  - [[14, -1], 1, CBFuse, [[4]]] # 27
+  - [-1, 1, RepNCSPELAN4, [1024, 512, 256, 2]] # 28
+  - [-1, 1, SPPELAN, [512, 256]] # 29
+
+# gelan head
+head:
+  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
+  - [[-1, 25], 1, Concat, [1]] # cat backbone P4
+  - [-1, 1, RepNCSPELAN4, [512, 512, 256, 2]] # 32
+
+  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
+  - [[-1, 22], 1, Concat, [1]] # cat backbone P3
+  - [-1, 1, RepNCSPELAN4, [256, 256, 128, 2]] # 35 (P3/8-small)
+
+  - [-1, 1, ADown, [256]]
+  - [[-1, 32], 1, Concat, [1]] # cat head P4
+  - [-1, 1, RepNCSPELAN4, [512, 512, 256, 2]] # 38 (P4/16-medium)
+
+  - [-1, 1, ADown, [512]]
+  - [[-1, 29], 1, Concat, [1]] # cat head P5
+  - [-1, 1, RepNCSPELAN4, [512, 1024, 512, 2]] # 41 (P5/32-large)
+
+  - [[35, 38, 41], 1, Segment, [nc, 32, 256]] # Segment (P3, P4, P5)

ultralytics-8.2.20/ultralytics/cfg/models/v9/yolov9e.yaml

@@ -0,0 +1,61 @@
+# Ultralytics YOLO 🚀, AGPL-3.0 license
+# YOLOv9e
+# 1225 layers, 58206592 parameters, 193.0 GFLOPs
+
+# parameters
+nc: 80 # number of classes
+
+# gelan backbone
+backbone:
+  - [-1, 1, Silence, []]
+  - [-1, 1, Conv, [64, 3, 2]] # 1-P1/2
+  - [-1, 1, Conv, [128, 3, 2]] # 2-P2/4
+  - [-1, 1, RepNCSPELAN4, [256, 128, 64, 2]] # 3
+  - [-1, 1, ADown, [256]] # 4-P3/8
+  - [-1, 1, RepNCSPELAN4, [512, 256, 128, 2]] # 5
+  - [-1, 1, ADown, [512]] # 6-P4/16
+  - [-1, 1, RepNCSPELAN4, [1024, 512, 256, 2]] # 7
+  - [-1, 1, ADown, [1024]] # 8-P5/32
+  - [-1, 1, RepNCSPELAN4, [1024, 512, 256, 2]] # 9
+
+  - [1, 1, CBLinear, [[64]]] # 10
+  - [3, 1, CBLinear, [[64, 128]]] # 11
+  - [5, 1, CBLinear, [[64, 128, 256]]] # 12
+  - [7, 1, CBLinear, [[64, 128, 256, 512]]] # 13
+  - [9, 1, CBLinear, [[64, 128, 256, 512, 1024]]] # 14
+
+  - [0, 1, Conv, [64, 3, 2]] # 15-P1/2
+  - [[10, 11, 12, 13, 14, -1], 1, CBFuse, [[0, 0, 0, 0, 0]]] # 16
+  - [-1, 1, Conv, [128, 3, 2]] # 17-P2/4
+  - [[11, 12, 13, 14, -1], 1, CBFuse, [[1, 1, 1, 1]]] # 18
+  - [-1, 1, RepNCSPELAN4, [256, 128, 64, 2]] # 19
+  - [-1, 1, ADown, [256]] # 20-P3/8
+  - [[12, 13, 14, -1], 1, CBFuse, [[2, 2, 2]]] # 21
+  - [-1, 1, RepNCSPELAN4, [512, 256, 128, 2]] # 22
+  - [-1, 1, ADown, [512]] # 23-P4/16
+  - [[13, 14, -1], 1, CBFuse, [[3, 3]]] # 24
+  - [-1, 1, RepNCSPELAN4, [1024, 512, 256, 2]] # 25
+  - [-1, 1, ADown, [1024]] # 26-P5/32
+  - [[14, -1], 1, CBFuse, [[4]]] # 27
+  - [-1, 1, RepNCSPELAN4, [1024, 512, 256, 2]] # 28
+  - [-1, 1, SPPELAN, [512, 256]] # 29
+
+# gelan head
+head:
+  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
+  - [[-1, 25], 1, Concat, [1]] # cat backbone P4
+  - [-1, 1, RepNCSPELAN4, [512, 512, 256, 2]] # 32
+
+  - [-1, 1, nn.Upsample, [None, 2, "nearest"]]
+  - [[-1, 22], 1, Concat, [1]] # cat backbone P3
+  - [-1, 1, RepNCSPELAN4, [256, 256, 128, 2]] # 35 (P3/8-small)
+
+  - [-1, 1, ADown, [256]]
+  - [[-1, 32], 1, Concat, [1]] # cat head P4
+  - [-1, 1, RepNCSPELAN4, [512, 512, 256, 2]] # 38 (P4/16-medium)
+
+  - [-1, 1, ADown, [512]]
+  - [[-1, 29], 1, Concat, [1]] # cat head P5
+  - [-1, 1, RepNCSPELAN4, [512, 1024, 512, 2]] # 41 (P5/32-large)
+
+  - [[35, 38, 41], 1, Detect, [nc]] # Detect(P3, P4, P5)

{ultralytics-8.2.18 → ultralytics-8.2.20}/ultralytics/data/build.py

@@ -128,7 +128,7 @@ def build_dataloader(dataset, batch, workers, shuffle=True, rank=-1):
     """Return an InfiniteDataLoader or DataLoader for training or validation set."""
     batch = min(batch, len(dataset))
     nd = torch.cuda.device_count()  # number of CUDA devices
-    nw = min([os.cpu_count() // max(nd, 1), workers])  # number of workers
+    nw = min(os.cpu_count() // max(nd, 1), workers)  # number of workers
     sampler = None if rank == -1 else distributed.DistributedSampler(dataset, shuffle=shuffle)
     generator = torch.Generator()
     generator.manual_seed(6148914691236517205 + RANK)

{ultralytics-8.2.18 → ultralytics-8.2.20}/ultralytics/engine/exporter.py

@@ -195,6 +195,9 @@ class Exporter:
         if not hasattr(model, "names"):
             model.names = default_class_names()
         model.names = check_class_names(model.names)
+        if self.args.half and self.args.int8:
+            LOGGER.warning("WARNING ⚠️ half=True and int8=True are mutually exclusive, setting half=False.")
+            self.args.half = False
         if self.args.half and onnx and self.device.type == "cpu":
             LOGGER.warning("WARNING ⚠️ half=True only compatible with GPU export, i.e. use device=0")
             self.args.half = False
@@ -1012,12 +1015,17 @@ class Exporter:
 
     def _add_tflite_metadata(self, file):
         """Add metadata to *.tflite models per https://www.tensorflow.org/lite/models/convert/metadata."""
-        from tflite_support import flatbuffers  # noqa
-        from tflite_support import metadata as _metadata  # noqa
-        from tflite_support import metadata_schema_py_generated as _metadata_fb  # noqa
+        import flatbuffers
+
+        if MACOS:  # TFLite Support bug https://github.com/tensorflow/tflite-support/issues/954#issuecomment-2108570845
+            from tflite_support import metadata  # noqa
+            from tflite_support import metadata_schema_py_generated as schema  # noqa
+        else:
+            from tensorflow_lite_support.metadata import metadata_schema_py_generated as schema  # noqa
+            from tensorflow_lite_support.metadata.python import metadata  # noqa
 
         # Create model info
-        model_meta = _metadata_fb.ModelMetadataT()
+        model_meta = schema.ModelMetadataT()
         model_meta.name = self.metadata["description"]
         model_meta.version = self.metadata["version"]
         model_meta.author = self.metadata["author"]
@@ -1028,41 +1036,41 @@ class Exporter:
         with open(tmp_file, "w") as f:
             f.write(str(self.metadata))
 
-        label_file = _metadata_fb.AssociatedFileT()
+        label_file = schema.AssociatedFileT()
         label_file.name = tmp_file.name
-        label_file.type = _metadata_fb.AssociatedFileType.TENSOR_AXIS_LABELS
+        label_file.type = schema.AssociatedFileType.TENSOR_AXIS_LABELS
 
         # Create input info
-        input_meta = _metadata_fb.TensorMetadataT()
+        input_meta = schema.TensorMetadataT()
         input_meta.name = "image"
         input_meta.description = "Input image to be detected."
-        input_meta.content = _metadata_fb.ContentT()
-        input_meta.content.contentProperties = _metadata_fb.ImagePropertiesT()
-        input_meta.content.contentProperties.colorSpace = _metadata_fb.ColorSpaceType.RGB
-        input_meta.content.contentPropertiesType = _metadata_fb.ContentProperties.ImageProperties
+        input_meta.content = schema.ContentT()
+        input_meta.content.contentProperties = schema.ImagePropertiesT()
+        input_meta.content.contentProperties.colorSpace = schema.ColorSpaceType.RGB
+        input_meta.content.contentPropertiesType = schema.ContentProperties.ImageProperties
 
         # Create output info
-        output1 = _metadata_fb.TensorMetadataT()
+        output1 = schema.TensorMetadataT()
         output1.name = "output"
         output1.description = "Coordinates of detected objects, class labels, and confidence score"
         output1.associatedFiles = [label_file]
         if self.model.task == "segment":
-            output2 = _metadata_fb.TensorMetadataT()
+            output2 = schema.TensorMetadataT()
             output2.name = "output"
             output2.description = "Mask protos"
             output2.associatedFiles = [label_file]
 
         # Create subgraph info
-        subgraph = _metadata_fb.SubGraphMetadataT()
+        subgraph = schema.SubGraphMetadataT()
         subgraph.inputTensorMetadata = [input_meta]
         subgraph.outputTensorMetadata = [output1, output2] if self.model.task == "segment" else [output1]
         model_meta.subgraphMetadata = [subgraph]
 
         b = flatbuffers.Builder(0)
-        b.Finish(model_meta.Pack(b), _metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER)
+        b.Finish(model_meta.Pack(b), metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER)
         metadata_buf = b.Output()
 
-        populator = _metadata.MetadataPopulator.with_model_file(str(file))
+        populator = metadata.MetadataPopulator.with_model_file(str(file))
         populator.load_metadata_buffer(metadata_buf)
         populator.load_associated_files([str(tmp_file)])
         populator.populate()

{ultralytics-8.2.18 → ultralytics-8.2.20}/ultralytics/engine/results.py

@@ -402,8 +402,8 @@ class Results(SimpleClass):
             )
             return results
 
-        data = self.boxes or self.obb
         is_obb = self.obb is not None
+        data = self.obb if is_obb else self.boxes
         h, w = self.orig_shape if normalize else (1, 1)
         for i, row in enumerate(data):  # xyxy, track_id if tracking, conf, class_id
             class_id, conf = int(row.cls), round(row.conf.item(), decimals)

{ultralytics-8.2.18 → ultralytics-8.2.20}/ultralytics/solutions/__init__.py

@@ -1,6 +1,7 @@
 # Ultralytics YOLO 🚀, AGPL-3.0 license
 
 from .ai_gym import AIGym
+from .analytics import Analytics
 from .distance_calculation import DistanceCalculation
 from .heatmap import Heatmap
 from .object_counter import ObjectCounter
@@ -16,4 +17,5 @@ __all__ = (
     "ParkingManagement",
     "QueueManager",
     "SpeedEstimator",
+    "Analytics",
 )

{ultralytics-8.2.18 → ultralytics-8.2.20}/ultralytics/solutions/ai_gym.py

@@ -73,11 +73,11 @@ class AIGym:
             self.stage = ["-" for _ in results[0]]
 
         self.keypoints = results[0].keypoints.data
-        self.annotator = Annotator(im0, line_width=2)
+        self.annotator = Annotator(im0, line_width=self.tf)
 
         for ind, k in enumerate(reversed(self.keypoints)):
             # Estimate angle and draw specific points based on pose type
-            if self.pose_type in {"pushup", "pullup", "abworkout"}:
+            if self.pose_type in {"pushup", "pullup", "abworkout", "squat"}:
                 self.angle[ind] = self.annotator.estimate_pose_angle(
                     k[int(self.kpts_to_check[0])].cpu(),
                     k[int(self.kpts_to_check[1])].cpu(),
@@ -93,7 +93,7 @@ class AIGym:
                         self.stage[ind] = "up"
                         self.count[ind] += 1
 
-                elif self.pose_type == "pushup":
+                elif self.pose_type == "pushup" or self.pose_type == "squat":
                     if self.angle[ind] > self.poseup_angle:
                         self.stage[ind] = "up"
                     if self.angle[ind] < self.posedown_angle and self.stage[ind] == "up":

ultralytics-8.2.20/ultralytics/solutions/analytics.py

@@ -0,0 +1,197 @@
+from itertools import cycle
+
+import cv2
+import matplotlib.pyplot as plt
+import numpy as np
+from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
+from matplotlib.figure import Figure
+
+
+class Analytics:
+    """A class to create and update various types of charts (line, bar, pie) for visual analytics."""
+
+    def __init__(
+        self,
+        type,
+        writer,
+        im0_shape,
+        title="ultralytics",
+        x_label="x",
+        y_label="y",
+        bg_color="white",
+        fg_color="black",
+        line_color="yellow",
+        line_width=2,
+        fontsize=13,
+        view_img=False,
+        save_img=True,
+    ):
+        """
+        Initialize the Analytics class with various chart types.
+
+        Args:
+            type (str): Type of chart to initialize ('line', 'bar', or 'pie').
+            writer: Video writer object to save the frames.
+            im0_shape (tuple): Shape of the input image (width, height).
+            title (str): Title of the chart.
+            x_label (str): Label for the x-axis.
+            y_label (str): Label for the y-axis.
+            bg_color (str): Background color of the chart.
+            fg_color (str): Foreground (text) color of the chart.
+            line_color (str): Line color for line charts.
+            line_width (int): Width of the lines in line charts.
+            fontsize (int): Font size for chart text.
+            view_img (bool): Whether to display the image.
+            save_img (bool): Whether to save the image.
+        """
+
+        self.bg_color = bg_color
+        self.fg_color = fg_color
+        self.view_img = view_img
+        self.save_img = save_img
+        self.title = title
+        self.writer = writer
+
+        # Set figure size based on image shape
+        figsize = (im0_shape[0] / 100, im0_shape[1] / 100)
+
+        if type == "line":
+            # Initialize line plot
+            fig = Figure(facecolor=self.bg_color, figsize=figsize)
+            self.canvas = FigureCanvas(fig)
+            self.ax = fig.add_subplot(111, facecolor=self.bg_color)
+            (self.line,) = self.ax.plot([], [], color=line_color, linewidth=line_width)
+
+        elif type == "bar" or type == "pie":
+            # Initialize bar or pie plot
+            self.fig, self.ax = plt.subplots(figsize=figsize, facecolor=self.bg_color)
+            self.ax.set_facecolor(self.bg_color)
+            color_palette = [
+                (31, 119, 180),
+                (255, 127, 14),
+                (44, 160, 44),
+                (214, 39, 40),
+                (148, 103, 189),
+                (140, 86, 75),
+                (227, 119, 194),
+                (127, 127, 127),
+                (188, 189, 34),
+                (23, 190, 207),
+            ]
+            self.color_palette = [(r / 255, g / 255, b / 255, 1) for r, g, b in color_palette]
+            self.color_cycle = cycle(self.color_palette)
+            self.color_mapping = {}
+
+            # Ensure pie chart is circular
+            self.ax.axis("equal") if type == "pie" else None
+
+        # Set common axis properties
+        self.ax.set_title(self.title, color=self.fg_color, fontsize=fontsize)
+        self.ax.set_xlabel(x_label, color=self.fg_color, fontsize=fontsize - 3)
+        self.ax.set_ylabel(y_label, color=self.fg_color, fontsize=fontsize - 3)
+        self.ax.tick_params(axis="both", colors=self.fg_color)
+
+    def update_line(self, frame_number, total_counts):
+        """
+        Update the line graph with new data.
+
+        Args:
+            frame_number (int): The current frame number.
+            total_counts (int): The total counts to plot.
+        """
+
+        # Update line graph data
+        x_data = self.line.get_xdata()
+        y_data = self.line.get_ydata()
+        x_data = np.append(x_data, float(frame_number))
+        y_data = np.append(y_data, float(total_counts))
+        self.line.set_data(x_data, y_data)
+        self.ax.relim()
+        self.ax.autoscale_view()
+        self.canvas.draw()
+        im0 = np.array(self.canvas.renderer.buffer_rgba())
+        im0 = cv2.cvtColor(im0[:, :, :3], cv2.COLOR_RGBA2BGR)
+
+        # Display and save the updated graph
+        cv2.imshow(self.title, im0) if self.view_img else None
+        self.writer.write(im0) if self.save_img else None
+
+    def update_bar(self, count_dict):
+        """
+        Update the bar graph with new data.
+
+        Args:
+            count_dict (dict): Dictionary containing the count data to plot.
+        """
+
+        # Update bar graph data
+        self.ax.clear()
+        self.ax.set_facecolor(self.bg_color)
+        labels = list(count_dict.keys())
+        counts = list(count_dict.values())
+
+        # Map labels to colors
+        for label in labels:
+            if label not in self.color_mapping:
+                self.color_mapping[label] = next(self.color_cycle)
+
+        colors = [self.color_mapping[label] for label in labels]
+
+        bars = self.ax.bar(labels, counts, color=colors)
+        for bar, count in zip(bars, counts):
+            self.ax.text(
+                bar.get_x() + bar.get_width() / 2,
+                bar.get_height(),
+                str(count),
+                ha="center",
+                va="bottom",
+                color=self.fg_color,
+            )
+
+        # Display and save the updated graph
+        canvas = FigureCanvas(self.fig)
+        canvas.draw()
+        buf = canvas.buffer_rgba()
+        im0 = np.asarray(buf)
+        im0 = cv2.cvtColor(im0, cv2.COLOR_RGBA2BGR)
+
+        self.writer.write(im0) if self.save_img else None
+        cv2.imshow(self.title, im0) if self.view_img else None
+
+    def update_pie(self, classes_dict):
+        """
+        Update the pie chart with new data.
+
+        Args:
+            classes_dict (dict): Dictionary containing the class data to plot.
+        """
+
+        # Update pie chart data
+        labels = list(classes_dict.keys())
+        sizes = list(classes_dict.values())
+        total = sum(sizes)
+        percentages = [size / total * 100 for size in sizes]
+        start_angle = 90
+        self.ax.clear()
+
+        # Create pie chart without labels inside the slices
+        wedges, autotexts = self.ax.pie(sizes, autopct=None, startangle=start_angle, textprops={"color": self.fg_color})
+
+        # Construct legend labels with percentages
+        legend_labels = [f"{label} ({percentage:.1f}%)" for label, percentage in zip(labels, percentages)]
+        self.ax.legend(wedges, legend_labels, title="Classes", loc="center left", bbox_to_anchor=(1, 0, 0.5, 1))
+
+        # Adjust layout to fit the legend
+        self.fig.tight_layout()
+        self.fig.subplots_adjust(left=0.1, right=0.75)
+
+        # Display and save the updated chart
+        im0 = self.fig.canvas.draw()
+        im0 = np.array(self.fig.canvas.renderer.buffer_rgba())
+        im0 = cv2.cvtColor(im0[:, :, :3], cv2.COLOR_RGBA2BGR)
+        self.writer.write(im0) if self.save_img else None
+        cv2.imshow(self.title, im0) if self.view_img else None
+
+
+if __name__ == "__main__":
+    Analytics("line", writer=None, im0_shape=None)