Merge pull request #20 from pangxiaobin/dev

release v0.2.4

.gitignore

@@ -49,4 +49,7 @@ backend/dist/
 
 backend/hub_model/briaai/
 
-backend/web/
+backend/web/
+
+# pnpm-lock.yaml
+frontend/pnpm-lock.yaml

README.md

@@ -67,6 +67,12 @@
 
 ### 历史版本记录
 
+#### v0.2.4
+
+- 优化图像边缘处理，减少黑边
+- 增加边缘处理的配置项，可选择是否开启边缘处理，默认开启
+- 增加边缘处理的度
+
 #### v0.2.3
 
 - 修复批量抠图时，webp格式文件被跳过的问题

backend/api/user_api.py

@@ -18,6 +18,8 @@ def put(self, playload: dict):
                 config.save(key, value)
             if key == "tinify.tinify_key":
                 tinify_client.update_key(value)
+            if key.startswith("edge_optimization"):
+                config.save(key, value)
         return res200(dict(config))
 
     def update_window_setting(self, payload):

backend/conf/config.py

@@ -19,6 +19,11 @@
         "on_top": False,
     },
     "tinify": {"tinify_key": "", "preserve": [], "compression_count": 0},
+    # 边缘优化
+    "edge_optimization": {
+        "r": 90,
+        "is_edge_optimization": True,
+    },
 }
 
 

backend/conf/setting.py

@@ -29,7 +29,7 @@ class Settings(BaseSettings):
 
     BASE_DIR: Path = Field(default=BASE_DIR, description="项目基础路径")
 
-    VERSION: str = Field(default="0.2.3", description="版本号")
+    VERSION: str = Field(default="0.2.4", description="版本号")
     # 日志
     LOG: dict = {"max_log_size": "20MB", "backup_count": 20}
 

backend/config.json

@@ -6,14 +6,18 @@
     "window": {
         "width": 1037,
         "height": 800,
-        "x": 236,
-        "y": 25,
+        "x": 42,
+        "y": 23,
         "fullscreen": false,
         "on_top": false
     },
     "tinify": {
         "tinify_key": "",
         "preserve": [],
         "compression_count": 0
+    },
+    "edge_optimization": {
+        "r": 90,
+        "is_edge_optimization": true
     }
 }

backend/hub_model/rmbg_model.py

@@ -8,16 +8,17 @@
 import onnxruntime as ort
 import requests
 from PIL import Image
+from conf.config import config
 
 BASE_DIR = Path(__file__).resolve().parent.parent
 
 try:
     from utilities.log import logger
-    from utilities.utils import img_to_base64
+    from utilities.utils import refine_foreground
 except ImportError:
     sys.path.append(str(BASE_DIR))
     from utilities.log import logger
-    from utilities.utils import img_to_base64
+    from utilities.utils import refine_foreground
 
 
 # 下载模型
@@ -117,17 +118,19 @@ def segment_image(self, img_obj: str) -> Image.Image:
         result = ort_outs[0]
         # 后处理
         result_image = self.postprocess_image(result[0][0], image_size)
-
-        try:
-            pil_im = Image.fromarray(result_image)
-            # pil_im = pil_im.resize(image_size)
-        except Exception as e:
-            raise RuntimeError(f"Error processing images: {e}")
-
-        no_bg_image = Image.new("RGBA", pil_im.size, (0, 0, 0, 0))
-        no_bg_image.paste(orig_image, mask=pil_im)
-
-        return no_bg_image
+        is_edge_optimization = config.get("edge_optimization.is_edge_optimization")
+        r = config.get("edge_optimization.r")
+        pil_im = Image.fromarray(result_image)
+        if is_edge_optimization:
+            try:
+                no_bg_image = refine_foreground(orig_image, pil_im, r=r)
+                return no_bg_image
+            except Exception as e:
+                logger.error(f"Error processing images: {e}")
+        else:
+            no_bg_image = Image.new("RGBA", pil_im.size, (0, 0, 0, 0))
+            no_bg_image.paste(orig_image, mask=pil_im)
+            return no_bg_image
 
 
 if __name__ == "__main__":

backend/utilities/utils.py

@@ -8,6 +8,7 @@
 import requests
 from pathlib import Path
 import numpy as np
+from cv2 import blur
 
 
 def is_image(filename):
@@ -271,3 +272,39 @@ def read_image(img: str) -> Image.Image:
         return Image.open(fp=img, mode="r")
     else:
         raise FileNotFoundError(f"File {img} not found.")
+
+
+def refine_foreground(image: Image.Image, mask: Image.Image, r=90):
+    if mask.size != image.size:
+        mask = mask.resize(image.size)
+    image = image.convert("RGB")
+    image_array = np.array(image) / 255.0
+    alpha_array = np.array(mask) / 255.0
+    estimated_foreground = FB_blur_fusion_foreground_estimator_2(
+        image_array, alpha_array, r=r
+    )
+    image_masked = Image.fromarray((estimated_foreground * 255.0).astype(np.uint8))
+    image_masked.putalpha(mask.resize(image.size))
+    return image_masked
+
+
+def FB_blur_fusion_foreground_estimator_2(image: np.array, alpha: np.array, r=90):
+    # Thanks to the source: https://github.com/Photoroom/fast-foreground-estimation
+    alpha = alpha[:, :, None]
+    F, blur_B = FB_blur_fusion_foreground_estimator(image, image, image, alpha, r)
+    return FB_blur_fusion_foreground_estimator(image, F, blur_B, alpha, r=6)[0]
+
+
+def FB_blur_fusion_foreground_estimator(image, F, B, alpha, r=90):
+    if isinstance(image, Image.Image):
+        image = np.array(image) / 255.0
+    blurred_alpha = blur(alpha, (r, r))[:, :, None]
+
+    blurred_FA = blur(F * alpha, (r, r))
+    blurred_F = blurred_FA / (blurred_alpha + 1e-5)
+
+    blurred_B1A = blur(B * (1 - alpha), (r, r))
+    blurred_B = blurred_B1A / ((1 - blurred_alpha) + 1e-5)
+    F = blurred_F + alpha * (image - alpha * blurred_F - (1 - alpha) * blurred_B)
+    F = np.clip(F, 0, 1)
+    return F, blurred_B

frontend/package.json

@@ -9,19 +9,19 @@
     "preview": "vite preview"
   },
   "dependencies": {
-    "@fortawesome/fontawesome-free": "^6.6.0",
+    "@fortawesome/fontawesome-free": "^6.7.1",
     "cropperjs": "^1.6.2",
-    "daisyui": "^4.12.10",
-    "vue": "^3.5.3",
-    "vue-i18n": "^9.14.0",
-    "vue-router": "^4.4.3",
+    "daisyui": "^4.12.20",
+    "vue": "^3.5.13",
+    "vue-i18n": "^9.14.2",
+    "vue-router": "^4.5.0",
     "vue3-compare-image": "^1.2.5"
   },
   "devDependencies": {
-    "@vitejs/plugin-vue": "^5.1.3",
+    "@vitejs/plugin-vue": "^5.2.1",
     "autoprefixer": "^10.4.20",
-    "postcss": "^8.4.45",
-    "tailwindcss": "^3.4.10",
-    "vite": "^5.4.3"
+    "postcss": "^8.4.49",
+    "tailwindcss": "^3.4.16",
+    "vite": "^5.4.11"
   }
 }