fix docstrings

OneMagicKey · May 28, 2024 · d081ff0 · d081ff0
1 parent 2754038
commit d081ff0
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/README.md b/README.md
@@ -106,7 +106,7 @@ follow these steps:
 1) Export a model to ONNX format:
 
    ```bash
-   !pip install ultralytics
+   pip install ultralytics
    ```
 
    ```python

diff --git a/src/utils.py b/src/utils.py
@@ -10,7 +10,13 @@
 @dataclass
 class ModelInfo:
     """
-    Dataclass for storing information about a model before initialisation.
+    A dataclass to represent information about a yolo model.
+
+    Attributes:
+        type (str): The type of the model ('detection' or 'segmentation').
+        name (str): The name of the model (e.g., 'yolov8s', 'yolov5s-seg').
+        input_size (tuple[int, int]): The size of the input data as a tuple (height, width).
+        version (Literal[5, 8]): The version of the model, restricted to either 5 or 8.
     """
 
     type: str