|
| 1 | +```mermaid |
| 2 | +classDiagram |
| 3 | + class ControlModel { |
| 4 | + +str type |
| 5 | + +ControlTag control |
| 6 | + +str from_name |
| 7 | + +str to_name |
| 8 | + +str value |
| 9 | + +YOLO model |
| 10 | + +float model_score_threshold |
| 11 | + +Optional[Dict[str, str]] label_map |
| 12 | + +LabelStudioMLBase label_studio_ml_backend |
| 13 | + +get_cached_model(path: str) YOLO |
| 14 | + +create(cls, mlbackend: LabelStudioMLBase, control: ControlTag) ControlModel |
| 15 | + +predict_regions(path: str) List[Dict] |
| 16 | + +debug_plot(image) |
| 17 | + } |
| 18 | +
|
| 19 | + class RectangleLabelsModel { |
| 20 | + +predict_regions(path: str) List[Dict] |
| 21 | + +create_rectangles(results, path) List[Dict] |
| 22 | + } |
| 23 | +
|
| 24 | + class RectangleLabelsObbModel { |
| 25 | + +predict_regions(path: str) List[Dict] |
| 26 | + +create_rotated_rectangles(results, path) List[Dict] |
| 27 | + } |
| 28 | + |
| 29 | +
|
| 30 | + class PolygonLabelsModel { |
| 31 | + +predict_regions(path: str) List[Dict] |
| 32 | + +create_polygons(results, path) List[Dict] |
| 33 | + } |
| 34 | + |
| 35 | + class KeyPointLabelsModel { |
| 36 | + +predict_regions(path: str) List[Dict] |
| 37 | + +create_keypoints(results, path) List[Dict] |
| 38 | + } |
| 39 | +
|
| 40 | + class ChoicesModel { |
| 41 | + +predict_regions(path: str) List[Dict] |
| 42 | + +create_choices(results, path) List[Dict] |
| 43 | + } |
| 44 | +
|
| 45 | + class VideoRectangleModel { |
| 46 | + +predict_regions(path: str) List[Dict] |
| 47 | + +create_video_rectangles(results, path) List[Dict] |
| 48 | + +update_tracker_params(yaml_path: str, prefix: str) str | None |
| 49 | + } |
| 50 | +
|
| 51 | + ControlModel <|-- RectangleLabelsModel |
| 52 | + ControlModel <|-- RectangleLabelsObbModel |
| 53 | + ControlModel <|-- PolygonLabelsModel |
| 54 | + ControlModel <|-- ChoicesModel |
| 55 | + ControlModel <|-- KeyPointLabelsModel |
| 56 | + ControlModel <|-- VideoRectangleModel |
| 57 | + |
| 58 | +``` |
| 59 | + |
| 60 | +### 1. **Architecture Overview** |
| 61 | + |
| 62 | +The architecture of the project is modular and is primarily centered around integrating YOLO-based models with Label Studio to automate the labeling of images and videos. The system is organized into several Python modules that interact with each other to perform this task. The main components of the architecture include: |
| 63 | + |
| 64 | +1. **Main YOLO Integration Module (`model.py`)**: |
| 65 | + - This is the central module that connects Label Studio with YOLO models. It handles the overall process of detecting control tags from Label Studio’s configuration, running predictions on tasks, and returning the predictions in the format that Label Studio expects. |
| 66 | + |
| 67 | +2. **Control Models (`control_models/`)**: |
| 68 | + - The control models are specialized modules that correspond to different annotation types in Label Studio (e.g., RectangleLabels, PolygonLabels, Choices, VideoRectangle). Each control model is responsible for handling specific types of annotations by using the YOLO model to predict the necessary regions or labels. |
| 69 | + |
| 70 | +3. **Base Control Model (`control_models/base.py`)**: |
| 71 | + - This is an abstract base class that provides common functionality for all control models. It handles tasks like loading the YOLO model, caching it for efficiency, and providing a template for the predict and create methods. |
| 72 | + |
| 73 | +4. **Specific Control Models**: |
| 74 | + - **RectangleLabelsModel (`control_models/rectanglelabels.py`)**: Handles bounding boxes (both simple and oriented bounding boxes) for images. |
| 75 | + - **PolygonLabelsModel (`control_models/polygonlabels.py`)**: Deals with polygon annotations, typically used for segmentation tasks. |
| 76 | + - **ChoicesModel (`control_models/choices.py`)**: Manages classification tasks where the model predicts one or more labels for the entire image. |
| 77 | + - **KeyPointLabelsModel (`control_models/keypointlabels.py`)**: Supports keypoint annotations, where the model predicts the locations of keypoints on an image. |
| 78 | + - **VideoRectangleModel (`control_models/videorectangle.py`)**: Focuses on tracking objects across video frames, generating bounding boxes for each frame. |
| 79 | + |
| 80 | +### 2. **Module Descriptions** |
| 81 | + |
| 82 | +1. **`model.py` (Main YOLO Integration Module)**: |
| 83 | + - **Purpose**: This module serves as the entry point for integrating YOLO models with Label Studio. It is responsible for setting up the YOLO model, detecting which control models are needed based on the Label Studio configuration, running predictions on tasks, and returning the results in the required format. |
| 84 | + - **Key Functions**: |
| 85 | + - `setup()`: Initializes the YOLO model parameters. |
| 86 | + - `detect_control_models()`: Scans the Label Studio configuration to determine which control models to use. |
| 87 | + - `predict()`: Runs predictions on a batch of tasks and formats the results for Label Studio. |
| 88 | + - `fit()`: (Not implemented) Placeholder for updating the model based on new annotations. |
| 89 | + |
| 90 | +2. **`control_models/base.py` (Base Control Model)**: |
| 91 | + - **Purpose**: Provides a common interface and shared functionality for all specific control models. It includes methods for loading and caching the YOLO model, plotting results for debugging, and abstract methods that need to be implemented by subclasses. |
| 92 | + - **Key Functions**: |
| 93 | + - `get_cached_model()`: Retrieves a YOLO model from cache or loads it if not cached. |
| 94 | + - `create()`: Factory method to instantiate a control model. |
| 95 | + - `predict_regions()`: Abstract method to be implemented by subclasses to perform predictions. |
| 96 | + |
| 97 | +3. **`control_models/choices.py` (ChoicesModel)**: |
| 98 | + - **Purpose**: Handles classification tasks where the model predicts one or more labels for an image. It converts the YOLO model’s classification output into Label Studio’s choices format. |
| 99 | + - **Key Functions**: |
| 100 | + - `create_choices()`: Processes the YOLO model’s output and maps it to the Label Studio choices format. |
| 101 | + |
| 102 | +4. **`control_models/rectanglelabels.py` (RectangleLabelsModel)**: |
| 103 | + - **Purpose**: Manages the creation of bounding box annotations, both simple (axis-aligned) and oriented (rotated), from the YOLO model’s output. |
| 104 | + - **Key Functions**: |
| 105 | + - `create_rectangles()`: Converts the YOLO model’s bounding box predictions into Label Studio’s rectangle labels format. |
| 106 | + - `create_rotated_rectangles()`: Handles oriented bounding boxes (OBB) by processing rotation angles and converting them to the required format. |
| 107 | + |
| 108 | +5. **`control_models/polygonlabels.py` (PolygonLabelsModel)**: |
| 109 | + - **Purpose**: Converts segmentation masks generated by the YOLO model into polygon annotations for Label Studio. This is useful for tasks where precise boundaries around objects are required. |
| 110 | + - **Key Functions**: |
| 111 | + - `create_polygons()`: Transforms the YOLO model’s segmentation output into polygon annotations. |
| 112 | + |
| 113 | +6. **`control_models/keypointlabels.py` (KeyPointLabelsModel)**: |
| 114 | + - **Purpose**: Supports keypoint annotations by predicting the locations of keypoints on an image using the pose YOLO model. |
| 115 | + - **Key Functions**: |
| 116 | + - `create_keypoints()`: Processes the YOLO model’s keypoint predictions and converts them into Label Studio’s keypoint labels format. |
| 117 | + |
| 118 | +7. **`control_models/videorectangle.py` (VideoRectangleModel)**: |
| 119 | + - **Purpose**: Focuses on tracking objects across video frames, using YOLO’s tracking capabilities to generate bounding box annotations for each frame in a video sequence. |
| 120 | + - **Key Functions**: |
| 121 | + - `predict_regions()`: Runs YOLO’s tracking model on a video and converts the results into Label Studio’s video rectangle format. |
| 122 | + - `create_video_rectangles()`: Processes the output of the tracking model to create a sequence of bounding boxes across video frames. |
| 123 | + - `update_tracker_params()`: Customizes the tracking parameters based on settings in Label Studio’s configuration. |
| 124 | + |
| 125 | +### **Module Interaction** |
| 126 | + |
| 127 | +- **Workflow**: The main workflow begins with `model.py`, which reads tasks and the Label Studio configuration to detect and instantiate the appropriate control models. These control models are responsible for making predictions using the YOLO model and converting the results into a format that Label Studio can use for annotations. |
| 128 | + |
| 129 | +- **Inter-Module Communication**: Each control model inherits from `ControlModel` in `base.py`, ensuring that they all share common methods for loading the YOLO model, handling predictions, and caching. The specific control models (e.g., RectangleLabelsModel, PolygonLabelsModel) implement the abstract methods defined in `ControlModel` to provide the specialized behavior needed for different types of annotations. |
| 130 | + |
| 131 | +This modular structure allows for easy extension and modification, where new control models can be added to handle additional annotation types or new model architectures. |
0 commit comments