[Model] Added support RMBG v1.4 model for image matting

CVHub520 · Sep 8, 2024 · 5a599e8 · 5a599e8
1 parent 3393134
commit 5a599e8
Show file tree

Hide file tree

Showing 29 changed files with 15,998 additions and 15,730 deletions.
diff --git a/README.md b/README.md
@@ -35,6 +35,7 @@
 
 - Sep. 2024:
   - Release version [2.4.2](https://github.com/CVHub520/X-AnyLabeling/releases/tag/v2.4.2)
+  - 🧸🧸🧸 Added support for image matting based on [RMBG v1.4 model](https://huggingface.co/briaai/RMBG-1.4).
   - 🔥🔥🔥 Added support for interactive video object tracking based on [Segment-Anything-2](https://github.com/CVHub520/segment-anything-2). [[Tutorial](examples/interactive_video_object_segmentation/README.md)]
 
 <br>
@@ -201,6 +202,8 @@ For more details, please refer to 👉 [model_zoo](./docs/en/model_zoo.md) 👈
   - [Tracking by Instance Segmentation](./examples/multiple_object_tracking/README.md)
   - [Tracking by Pose Estimation](./examples/multiple_object_tracking/README.md)
 - [iVOS](./examples/interactive_video_object_segmentation/README.md)
+- [Matting](./examples/matting/)
+  - [Image Matting](./examples/matting/image_matting/README.md)
 
 ## Contact
 

diff --git a/README_zh-CN.md b/README_zh-CN.md
@@ -34,6 +34,7 @@
 
 - 2024年9月:
   - 发布[X-AnyLabeling v2.4.2](https://github.com/CVHub520/X-AnyLabeling/releases/tag/v2.4.2)版本。
+  - 🧸🧸🧸 支持基于[RMBG v1.4](https://github.com/facebookresearch/segment-anything)的图像抠图功能。
   - 🔥🔥🔥 支持基于[Segment-Anything-2](https://github.com/CVHub520/segment-anything-2)交互式视频目标追踪功能。【[教程](examples/interactive_video_object_segmentation/README.md) | [B站](https://www.bilibili.com/video/BV1kMpwedE8M/)】
 
 <br>
@@ -204,6 +205,8 @@
   - [Tracking by Instance Segmentation](./examples/multiple_object_tracking/README.md)
   - [Tracking by Pose Estimation](./examples/multiple_object_tracking/README.md)
 - [iVOS](./examples/interactive_video_object_segmentation/README.md)
+- [Matting](./examples/matting/)
+  - [Image Matting](./examples/matting/image_matting/README.md)
 
 ## 联系
 

diff --git a/anylabeling/configs/auto_labeling/models.yaml b/anylabeling/configs/auto_labeling/models.yaml
@@ -2,6 +2,8 @@
   config_file: ":/sam2_hiera_base.yaml"
 - model_name: "sam2_hiera_large_video-r20240901"
   config_file: ":/sam2_hiera_large_video.yaml"
+- model_name: "rmbg_v14-r20240908"
+  config_file: ":/rmbg_v14.yaml"
 - model_name: "yolov5s-r20230520"
   config_file: ":/yolov5s.yaml"
 - model_name: "yolov5_car_plate-r20230112"

diff --git a/anylabeling/configs/auto_labeling/rmbg_v14.yaml b/anylabeling/configs/auto_labeling/rmbg_v14.yaml
@@ -0,0 +1,4 @@
+type: rmbg
+name: rmbg_v14-r20240908
+display_name: RMBG v1.4 (BRIA AI)
+model_path: https://github.com/CVHub520/X-AnyLabeling/releases/download/v2.4.3/bria-rmbg-1.4.onnx
diff --git a/anylabeling/resources/resources.py b/anylabeling/resources/resources.py
diff --git a/anylabeling/resources/translations/en_US.ts b/anylabeling/resources/translations/en_US.ts
@@ -1975,67 +1975,67 @@ Check the results in: {save_path}.</source>
 <context>
     <name>ModelManager</name>
     <message>
-        <location filename="../../services/auto_labeling/model_manager.py" line="176"/>
+        <location filename="../../services/auto_labeling/model_manager.py" line="177"/>
         <source>Model loaded. Ready for labeling.</source>
         <translation type="unfinished"></translation>
     </message>
     <message>
-        <location filename="../../services/auto_labeling/model_manager.py" line="279"/>
+        <location filename="../../services/auto_labeling/model_manager.py" line="280"/>
         <source>No model selected.</source>
         <translation type="unfinished"></translation>
     </message>
     <message>
-        <location filename="../../services/auto_labeling/model_manager.py" line="295"/>
+        <location filename="../../services/auto_labeling/model_manager.py" line="296"/>
         <source>Loading model: {model_name}. Please wait...</source>
         <translation type="unfinished"></translation>
     </message>
     <message>
-        <location filename="../../services/auto_labeling/model_manager.py" line="1694"/>
+        <location filename="../../services/auto_labeling/model_manager.py" line="1720"/>
         <source>Model is not loaded. Choose a mode to continue.</source>
         <translation type="unfinished"></translation>
     </message>
     <message>
-        <location filename="../../services/auto_labeling/model_manager.py" line="1676"/>
+        <location filename="../../services/auto_labeling/model_manager.py" line="1700"/>
         <source>Finished inferencing AI model. Check the result.</source>
         <translation type="unfinished"></translation>
     </message>
     <message>
-        <location filename="../../services/auto_labeling/model_manager.py" line="1698"/>
+        <location filename="../../services/auto_labeling/model_manager.py" line="1724"/>
         <source>Inferencing AI model. Please wait...</source>
         <translation type="unfinished"></translation>
     </message>
     <message>
-        <location filename="../../services/auto_labeling/model_manager.py" line="1708"/>
+        <location filename="../../services/auto_labeling/model_manager.py" line="1734"/>
         <source>Another model is being executed. Please wait for it to finish.</source>
         <translation type="unfinished"></translation>
     </message>
     <message>
-        <location filename="../../services/auto_labeling/model_manager.py" line="201"/>
+        <location filename="../../services/auto_labeling/model_manager.py" line="202"/>
         <source>Error in loading custom model: Invalid path.</source>
         <translation type="unfinished"></translation>
     </message>
     <message>
-        <location filename="../../services/auto_labeling/model_manager.py" line="212"/>
+        <location filename="../../services/auto_labeling/model_manager.py" line="213"/>
         <source>Error in loading custom model: Invalid config file.</source>
         <translation type="unfinished"></translation>
     </message>
     <message>
-        <location filename="../../services/auto_labeling/model_manager.py" line="222"/>
+        <location filename="../../services/auto_labeling/model_manager.py" line="223"/>
         <source>Error in loading custom model: Invalid config file format.</source>
         <translation type="unfinished"></translation>
     </message>
     <message>
-        <location filename="../../services/auto_labeling/model_manager.py" line="289"/>
+        <location filename="../../services/auto_labeling/model_manager.py" line="290"/>
         <source>Error in loading model: Invalid model name.</source>
         <translation type="unfinished"></translation>
     </message>
     <message>
-        <location filename="../../services/auto_labeling/model_manager.py" line="1473"/>
+        <location filename="../../services/auto_labeling/model_manager.py" line="1497"/>
         <source>Error in loading model: {error_message}</source>
         <translation type="unfinished"></translation>
     </message>
     <message>
-        <location filename="../../services/auto_labeling/model_manager.py" line="1681"/>
+        <location filename="../../services/auto_labeling/model_manager.py" line="1705"/>
         <source>Error in model prediction: {e}. Please check the model.</source>
         <translation type="unfinished"></translation>
     </message>
@@ -2253,22 +2253,22 @@ and
 <context>
     <name>parent</name>
     <message>
-        <location filename="../../views/labeling/utils/video.py" line="35"/>
+        <location filename="../../views/labeling/utils/video.py" line="37"/>
         <source>Frame Interval</source>
         <translation type="unfinished"></translation>
     </message>
     <message>
-        <location filename="../../views/labeling/utils/video.py" line="35"/>
+        <location filename="../../views/labeling/utils/video.py" line="37"/>
         <source>Enter the frame interval (FPS: {fps}):</source>
         <translation type="unfinished"></translation>
     </message>
     <message>
-        <location filename="../../views/labeling/utils/video.py" line="68"/>
+        <location filename="../../views/labeling/utils/video.py" line="73"/>
         <source>Extracting frames. Please wait...</source>
         <translation type="unfinished"></translation>
     </message>
     <message>
-        <location filename="../../views/labeling/utils/video.py" line="68"/>
+        <location filename="../../views/labeling/utils/video.py" line="73"/>
         <source>Cancel</source>
         <translation type="unfinished"></translation>
     </message>

diff --git a/anylabeling/resources/translations/zh_CN.ts b/anylabeling/resources/translations/zh_CN.ts
@@ -2393,47 +2393,47 @@ Check the results in: {save_path}.</source>
 <context>
     <name>ModelManager</name>
     <message>
-        <location filename="../../services/auto_labeling/model_manager.py" line="176"/>
+        <location filename="../../services/auto_labeling/model_manager.py" line="177"/>
         <source>Model loaded. Ready for labeling.</source>
         <translation>模型已加载。准备好进行数据标注。</translation>
     </message>
     <message>
-        <location filename="../../services/auto_labeling/model_manager.py" line="201"/>
+        <location filename="../../services/auto_labeling/model_manager.py" line="202"/>
         <source>Error in loading custom model: Invalid path.</source>
         <translation>加载自定义模型时出错：无效路径。</translation>
     </message>
     <message>
-        <location filename="../../services/auto_labeling/model_manager.py" line="212"/>
+        <location filename="../../services/auto_labeling/model_manager.py" line="213"/>
         <source>Error in loading custom model: Invalid config file.</source>
         <translation>加载自定义模型时出错：无效的配置文件。</translation>
     </message>
     <message>
-        <location filename="../../services/auto_labeling/model_manager.py" line="222"/>
+        <location filename="../../services/auto_labeling/model_manager.py" line="223"/>
         <source>Error in loading custom model: Invalid config file format.</source>
         <translation>加载自定义模型时出错：无效的配置文件格式。</translation>
     </message>
     <message>
-        <location filename="../../services/auto_labeling/model_manager.py" line="279"/>
+        <location filename="../../services/auto_labeling/model_manager.py" line="280"/>
         <source>No model selected.</source>
         <translation>未选择模型。</translation>
     </message>
     <message>
-        <location filename="../../services/auto_labeling/model_manager.py" line="289"/>
+        <location filename="../../services/auto_labeling/model_manager.py" line="290"/>
         <source>Error in loading model: Invalid model name.</source>
         <translation>加载模型时出错：无效的模型名称。</translation>
     </message>
     <message>
-        <location filename="../../services/auto_labeling/model_manager.py" line="295"/>
+        <location filename="../../services/auto_labeling/model_manager.py" line="296"/>
         <source>Loading model: {model_name}. Please wait...</source>
         <translation>加载模型 {model_name} 中，请稍等。。。</translation>
     </message>
     <message>
-        <location filename="../../services/auto_labeling/model_manager.py" line="1473"/>
+        <location filename="../../services/auto_labeling/model_manager.py" line="1497"/>
         <source>Error in loading model: {error_message}</source>
         <translation>加载模型时出错：{error_message}。</translation>
     </message>
     <message>
-        <location filename="../../services/auto_labeling/model_manager.py" line="1694"/>
+        <location filename="../../services/auto_labeling/model_manager.py" line="1720"/>
         <source>Model is not loaded. Choose a mode to continue.</source>
         <translation>模型尚未加载。选择一个模式以继续。</translation>
     </message>
@@ -2443,22 +2443,22 @@ Check the results in: {save_path}.</source>
         <translation type="obsolete">模型预测出错。请检查模型。</translation>
     </message>
     <message>
-        <location filename="../../services/auto_labeling/model_manager.py" line="1676"/>
+        <location filename="../../services/auto_labeling/model_manager.py" line="1700"/>
         <source>Finished inferencing AI model. Check the result.</source>
         <translation>完成 AI 模型推理，请检查结果。</translation>
     </message>
     <message>
-        <location filename="../../services/auto_labeling/model_manager.py" line="1698"/>
+        <location filename="../../services/auto_labeling/model_manager.py" line="1724"/>
         <source>Inferencing AI model. Please wait...</source>
         <translation>正在进行 AI 模型推理，请稍等...</translation>
     </message>
     <message>
-        <location filename="../../services/auto_labeling/model_manager.py" line="1708"/>
+        <location filename="../../services/auto_labeling/model_manager.py" line="1734"/>
         <source>Another model is being executed. Please wait for it to finish.</source>
         <translation>另一个模型正在执行，请耐心等待。</translation>
     </message>
     <message>
-        <location filename="../../services/auto_labeling/model_manager.py" line="1681"/>
+        <location filename="../../services/auto_labeling/model_manager.py" line="1705"/>
         <source>Error in model prediction: {e}. Please check the model.</source>
         <translation>模型推理失败：{e}. 请检查模型。</translation>
     </message>
@@ -2684,22 +2684,22 @@ and
 <context>
     <name>parent</name>
     <message>
-        <location filename="../../views/labeling/utils/video.py" line="35"/>
+        <location filename="../../views/labeling/utils/video.py" line="37"/>
         <source>Frame Interval</source>
         <translation>帧率</translation>
     </message>
     <message>
-        <location filename="../../views/labeling/utils/video.py" line="35"/>
+        <location filename="../../views/labeling/utils/video.py" line="37"/>
         <source>Enter the frame interval (FPS: {fps}):</source>
         <translation>请输入抽帧间隔（FPS: {fps}）：</translation>
     </message>
     <message>
-        <location filename="../../views/labeling/utils/video.py" line="68"/>
+        <location filename="../../views/labeling/utils/video.py" line="73"/>
         <source>Extracting frames. Please wait...</source>
         <translation>正在解析中，请稍等。。。</translation>
     </message>
     <message>
-        <location filename="../../views/labeling/utils/video.py" line="68"/>
+        <location filename="../../views/labeling/utils/video.py" line="73"/>
         <source>Cancel</source>
         <translation>取消</translation>
     </message>

diff --git a/anylabeling/services/auto_labeling/depth_anything.py b/anylabeling/services/auto_labeling/depth_anything.py
@@ -322,7 +322,7 @@ def predict_shapes(self, image, image_path=None):
         depth = self.postprocess(output, orig_shape)
 
         image_dir_path = os.path.dirname(image_path)
-        save_path = os.path.join(image_dir_path, "..", "depth")
+        save_path = os.path.join(image_dir_path, "..", "x-anylabeling-depth")
         save_path = os.path.realpath(save_path)
         os.makedirs(save_path, exist_ok=True)
         image_file_name = os.path.basename(image_path)

diff --git a/anylabeling/services/auto_labeling/depth_anything_v2.py b/anylabeling/services/auto_labeling/depth_anything_v2.py
@@ -107,7 +107,7 @@ def predict_shapes(self, image, image_path=None):
         depth = self.postprocess(outputs, orig_shape)
 
         image_dir_path = os.path.dirname(image_path)
-        save_path = os.path.join(image_dir_path, "..", "depth")
+        save_path = os.path.join(image_dir_path, "..", "x-anylabeling-depth")
         save_path = os.path.realpath(save_path)
         os.makedirs(save_path, exist_ok=True)
         image_file_name = os.path.basename(image_path)

diff --git a/anylabeling/services/auto_labeling/model_manager.py b/anylabeling/services/auto_labeling/model_manager.py
@@ -60,10 +60,11 @@ class ModelManager(QObject):
         "yolov8_obb",
         "yolov5_car_plate",
         "rtmdet_pose",
-        "depth_anything",
         "yolov9",
         "yolow",
         "yolov10",
+        "rmbg",
+        "depth_anything",
         "depth_anything_v2",
         "yolow_ram",
         "rtdetrv2",
@@ -1440,6 +1441,28 @@ def _load_model(self, model_id):
                     )
                 )
                 return
+        elif model_config["type"] == "rmbg":
+            from .rmbg import RMBG
+
+            try:
+                model_config["model"] = RMBG(
+                    model_config, on_message=self.new_model_status.emit
+                )
+                self.auto_segmentation_model_unselected.emit()
+            except Exception as e:  # noqa
+                self.new_model_status.emit(
+                    self.tr(
+                        "Error in loading model: {error_message}".format(
+                            error_message=str(e)
+                        )
+                    )
+                )
+                print(
+                    "Error in loading model: {error_message}".format(
+                        error_message=str(e)
+                    )
+                )
+                return
         elif model_config["type"] == "depth_anything":
             from .depth_anything import DepthAnything