From 2fc1bf042a46c04b72f8663ea2d099f9c189d9a6 Mon Sep 17 00:00:00 2001
From: Qing <cwq1913@gmail.com>
Date: Thu, 4 Jan 2024 21:39:59 +0800
Subject: [PATCH] fix PIL exif rotate

---
 lama_cleaner/helper.py        | 26 +++++++++++++++++++++++++-
 lama_cleaner/model_manager.py |  5 +++--
 2 files changed, 28 insertions(+), 3 deletions(-)

diff --git a/lama_cleaner/helper.py b/lama_cleaner/helper.py
index 00acc19..551c457 100644
--- a/lama_cleaner/helper.py
+++ b/lama_cleaner/helper.py
@@ -309,11 +309,12 @@ def decode_base64_to_image(
     image = Image.open(io.BytesIO(base64.b64decode(encoding)))
 
     alpha_channel = None
-    infos = image.info
     try:
         image = ImageOps.exif_transpose(image)
     except:
         pass
+    # exif_transpose will remove exif rotate info，we must call image.info after exif_transpose
+    infos = image.info
 
     if gray:
         image = image.convert("L")
@@ -352,6 +353,29 @@ def concat_alpha_channel(rgb_np_img, alpha_channel) -> np.ndarray:
     return rgb_np_img
 
 
+def adjust_mask(mask: np.ndarray, kernel_size: int, operate):
+    # kernel_size = kernel_size*2+1
+    mask[mask >= 127] = 255
+    mask[mask < 127] = 0
+    # fronted brush color "ffcc00bb"
+    if operate == "expand":
+        mask = cv2.dilate(
+            mask,
+            np.ones((kernel_size, kernel_size), np.uint8),
+            iterations=1,
+        )
+    else:
+        mask = cv2.erode(
+            mask,
+            np.ones((kernel_size, kernel_size), np.uint8),
+            iterations=1,
+        )
+    res_mask = np.zeros((mask.shape[0], mask.shape[1], 4), dtype=np.uint8)
+    res_mask[mask > 128] = [255, 203, 0, int(255 * 0.73)]
+    res_mask = cv2.cvtColor(res_mask, cv2.COLOR_BGRA2RGBA)
+    return res_mask
+
+
 def gen_frontend_mask(bgr_or_gray_mask):
     if len(bgr_or_gray_mask.shape) == 3 and bgr_or_gray_mask.shape[2] != 1:
         bgr_or_gray_mask = cv2.cvtColor(bgr_or_gray_mask, cv2.COLOR_BGR2GRAY)
diff --git a/lama_cleaner/model_manager.py b/lama_cleaner/model_manager.py
index f31bdaa..e73b8b7 100644
--- a/lama_cleaner/model_manager.py
+++ b/lama_cleaner/model_manager.py
@@ -2,6 +2,7 @@ from typing import List, Dict
 
 import torch
 from loguru import logger
+import numpy as np
 
 from lama_cleaner.download import scan_models
 from lama_cleaner.helper import switch_mps_device
@@ -77,12 +78,12 @@ class ModelManager:
             config:
 
         Returns:
-
+            BGR image
         """
         self.switch_controlnet_method(config)
         self.enable_disable_freeu(config)
         self.enable_disable_lcm_lora(config)
-        return self.model(image, mask, config)
+        return self.model(image, mask, config).astype(np.uint8)
 
     def scan_models(self) -> List[ModelInfo]:
         available_models = scan_models()