Add semitransparency model into removebg (kubeflow#106)

* 1.initial submit of adding semitransparency into removebg pipeline * merged new semitransparency model into removebg, local tested thro rabbitmq, all test passed * 1.remove testing the new model in test_carinterior_semitransparency 2.modified test_car_semitransparency for api * 1.modified some changes that i made when testing locally but forgot to change back. mainly in removebg.demo and remove.py * changed removebg resize related code, before when testing locally, it has issue with DA lib, so need to squeeze it before resizing. now testing tro building the whole pipeline pass * 1.modified server code and related test files on adding a new flag for the new semitransparency model. 2. updated semitransparency-core lib version * remove the squeeze and unsqueeze code which is needed to test locally. change the version requirement for semitransparency-core * putting model into eval and half when initializing instead of in semitransparency func * upgraded semitransparency-core lib to remove google cloud related and updated the requirement * pillow requirement removed from requirementdeploy * update semitransparencycore lib version to 0.9.1 * update semitransparency-core-lib to the newest version 0.9.3 Co-authored-by: Alexander Graf <[email protected]> Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> Co-authored-by: jerome <[email protected]>
georgkaleido · Mar 17, 2022 · f373b1c · f373b1c
1 parent 51406cb
commit f373b1c
Show file tree

Hide file tree

Showing 15 changed files with 162 additions and 68 deletions.
diff --git a/.gitignore b/.gitignore
@@ -326,3 +326,7 @@ core/sweep_local/sample_imgs/
 # Python Dist Tools
 core/ChangeLog
 core/AUTHORS
+
+# local pycharm setting files
+.idea/*
+core/.idea/*
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,23 +1,5 @@
 # Changelog
 
-## 1.21.0 (14/03/2022)
-
-#### 🚀 Enhancements:
-
-- Update kaleido-models version. #trivial [#117](https://github.com/remove-bg/kaleido-removebg/pull/117)
-
-#### 🔀 Dependencies
-
-- Update dependencies [#114](https://github.com/remove-bg/kaleido-removebg/pull/114) 
-
-### Docker Images
-
-* `eu.gcr.io/removebg-226919/removebg-api:1.21.0`
-* `eu.gcr.io/removebg-226919/removebg-core:1.21.0-cc75`
-* `eu.gcr.io/removebg-226919/removebg-core:1.21.0-cc61`
-* `eu.gcr.io/removebg-226919/removebg-core:1.21.0-cc86`
----
-
 ## 1.20.0 (08/03/2022)
 
 #### 🚀 Enhancements:

diff --git a/README.md b/README.md
@@ -174,15 +174,15 @@ COMPUTE_CAPABILITY=[your GPU compute capability] DOCKER_BUILDKIT=1 FURY_AUTH_TOK
 ```
 3. Set the following environment variables in your shell:
 ```bash
-REQUEST_QUEUE=remove_bg
-RABBITMQ_HOST=127.0.0.1
-RABBITMQ_PORT=5672
-RABBITMQ_USER=rabbitmq
-RABBITMQ_PASSWORD=rabbitmq
+export REQUEST_QUEUE=remove_bg
+export RABBITMQ_HOST=127.0.0.1
+export RABBITMQ_PORT=5672
+export RABBITMQ_USER=rabbitmq
+export RABBITMQ_PASSWORD=rabbitmq
 ```
 4. Execute tests as the following:
 ```bash
 # change to test directory
-cd core/test/
-pytest core/
+cd core/test/core/
+pytest .
 ```
diff --git a/core/bin/removebg-demo.py b/core/bin/removebg-demo.py
@@ -53,6 +53,7 @@ def handle_args():
     parser.add_argument("--subject_crop_margin", default=0, type=int, help="relative crop margin")
 
     parser.add_argument("--compute_device", default="cuda", type=str, choices=["cuda", "cpu"], help="Choose device between CUDA and CPU. Default: CUDA")
+    parser.add_argument("--semitransparency_experimental", action="store_true", help="use semitransparent matting model to the car window")
 
     args = parser.parse_args()
     return args
@@ -105,9 +106,7 @@ def demo(args):
         t_start = time.time()
 
         for file in files:
-
             is_folder = not os.path.isfile(file)
-
             if is_folder:
                 path_save = os.path.join(file, (args.caption + "_" if args.caption else "") + "result.png")
             else:
@@ -184,6 +183,7 @@ def demo(args):
                     color_enabled=(not args.no_color),
                     shadow_enabled=args.shadow,
                     trimap_confidence_thresh=trimap_confidence_thresh,
+                    semitranspareny_new_enabled=(cls == 'car' and args.semitransparency_experimental and not args.no_semitransparency)
                 )
             except UnknownForegroundException as e:
                 print("could not detect foreground: {}".format(e))
@@ -204,12 +204,13 @@ def demo(args):
 
             if not args.evaluate:
                 if cls == "car":
-                    image.fill_holes(
-                        255 if args.no_semitransparency else 200,
-                        mode="car",
-                        average=(not args.no_semitransparency),
-                        im_rgb_precolorcorr=im_rgb_precolorcorr,
-                    )
+                    if not (not args.no_semitransparency and args.semitransparency_experimental):
+                        image.fill_holes(
+                            255 if args.no_semitransparency else 200,
+                            mode="car",
+                            average=(not args.no_semitransparency),
+                            im_rgb_precolorcorr=im_rgb_precolorcorr,
+                        )
 
                 elif cls == "car_interior":
                     image.fill_holes(

diff --git a/core/bin/removebg-server.py b/core/bin/removebg-server.py
@@ -56,6 +56,8 @@ def _preprocess(self, processing_data: ImageProcessingData, message: Dict[bytes,
         roi = message.get(b"roi", None)
         shadow = message.get(b"shadow", False)
         semitransparency = message.get(b"semitransparency", True)
+        semitransparency_experimental = message.get(b"semitransparency_experimental", False)
+
         image_bg = None
 
         times = []
@@ -151,6 +153,7 @@ def compute_crop_roi(roi, scale):
             "crop_margin": crop_margin,
             "crop_roi": crop_roi_image,
             "semitransparency": semitransparency,
+            "semitransparency_experimental": semitransparency_experimental,
         }
 
     def _post_process(self, processing_data: ImageProcessingData) -> None:
@@ -163,6 +166,7 @@ def _post_process(self, processing_data: ImageProcessingData) -> None:
         times = data["times"]
         image = data["image"]
         semitransparency = data["semitransparency"]
+        semitransparency_experimental = data["semitransparency_experimental"]
         crop_margin = data["crop_margin"]
         scale_param = data["scale_param"]
         position_param = data["position_param"]
@@ -182,14 +186,15 @@ def _post_process(self, processing_data: ImageProcessingData) -> None:
         im_rgb_precolorcorr = image.get("rgb")
         image.set(im_res, "bgra", limit_alpha=True, crop=data["crop_roi"])
 
-        # car windows
+        # car windows, new model only works for car but not car interior
         if data["api"] == "car":
-            image.fill_holes(
-                200 if semitransparency else 255,
-                mode="car",
-                average=semitransparency,
-                im_rgb_precolorcorr=im_rgb_precolorcorr,
-            )
+            if not(semitransparency_experimental and semitransparency): #only if both are true, use new model
+                image.fill_holes(
+                    200 if semitransparency else 255,
+                    mode="car",
+                    average=semitransparency,
+                    im_rgb_precolorcorr=im_rgb_precolorcorr,
+                )
         elif data["api"] == "car_interior":
             image.fill_holes(
                 200 if semitransparency else 0,
@@ -367,6 +372,7 @@ def __init__(
     def _process(self, data: ImageProcessingData) -> None:
         processing_data = data.data
         im, shadow = processing_data["im_cv"], processing_data["shadow"]
+        semitransparency, semitransparency_experimental = processing_data["semitransparency"], processing_data["semitransparency_experimental"]
         im_for_trimap = processing_data["im_for_trimap"]
         result = data.result
 
@@ -396,6 +402,7 @@ def _process(self, data: ImageProcessingData) -> None:
                     im_for_trimap_tr,
                     color_enabled=(processing_data["api"] == "person" or processing_data["api"] == "animal"),
                     shadow_enabled=(processing_data["api"] == "car" and shadow),
+                    semitranspareny_new_enabled=(processing_data["api"] == "car" and semitransparency and semitransparency_experimental),
                     trimap_confidence_thresh=trimap_confidence_thresh,
                 )
             except UnknownForegroundException:

diff --git a/core/kaleido-models.yml b/core/kaleido-models.yml
@@ -1,4 +1,4 @@
-version: "1.33.0"
+version: "1.30.0"
 files:
   - name: identifier-mobilenetv2-c9.pth.tar
     hash: 25bf45cb8cd7f3165395ac1194a4fac0aec4db0e0d8b818c3216b899794d3720
@@ -8,3 +8,7 @@ files:
     hash: fc3ca55b394ca30c85588c3f93fb18eb9ad6bc069cbab16a91c91b06dcc3b5d9
   - name: trimap513-deeplab-res2net.pth.tar
     hash: c9ea9fa9deda7eb7b84c2a27dbc3813580555a773f8ea292363aa22e1191c1f1
+  - name: window_segmentation_deeplab_best_e40.pth.tar
+    hash: 04e6ceacc8376b01738950876bd0a5b301ac1820652761c01e992ac88a63d419
+  - name: semimatting_windinput50_aftergamma_windloss_fix2_5_change2_last_e27.pth.tar
+    hash: 0db520d33f8f55c733b6642114c18390e81d916cdb7a90660c76ab3917109b3e
diff --git a/core/removebg/removebg.py b/core/removebg/removebg.py
@@ -13,6 +13,7 @@
 from kaleido.image.resize import resize_tensor
 from removebg.models import Classifier, Matting, Trimap
 from shadowgen.models.oriented_shadow import OrientedShadow
+from semitransparency.semitransparency import Semitransparency
 
 
 class UnknownForegroundException(Exception):
@@ -90,12 +91,20 @@ def load_model(model, path, k="state_dict", strict=True):
             self.model_shadow, os.path.join(model_paths, "shadowgen256_car.pth.tar"), k=None, strict=False
         )
 
-    def __call__(self, im, im_for_trimap_tr=None, color_enabled=False, shadow_enabled=False, trimap_confidence_thresh=0.5, extra_trimap_output=False):
+        # load semi_transparency models (segmentation+matting) to cpu first
+        self.semitransparency_model = Semitransparency(matting_weight_path=os.path.join(model_paths,
+                                                                                        'semimatting_windinput50_aftergamma_windloss_fix2_5_change2_last_e27.pth.tar'),
+                                                       segmentation_weight_path=os.path.join(model_paths,
+                                                                                             'window_segmentation_deeplab_best_e40.pth.tar'),
+                                                       device='cpu')
+        # semitransparency models are already in eval mode, float. here make segmentation model half
+        self.semitransparency_model.model_segmentation.half()
+
+    def __call__(self, im, im_for_trimap_tr=None, color_enabled=False, shadow_enabled=False, semitranspareny_new_enabled=False, trimap_confidence_thresh=0.5, extra_trimap_output=False):
 
         has_trimap_optimized_image = im_for_trimap_tr is not None
 
         # all operations are done with half precision
-
         if "cuda" == self.compute_device:
             im = im.half()
             if has_trimap_optimized_image:
@@ -121,12 +130,12 @@ def __call__(self, im, im_for_trimap_tr=None, color_enabled=False, shadow_enable
             if color_enabled:
                 im = im_color
 
-            # shadow
-
             if shadow_enabled:
                 im, im_alpha = self.shadow(im, im_alpha)
 
-            # return
+            # semi-transparency matting for window
+            if semitranspareny_new_enabled:
+                im, im_alpha = self.semitransparency(im, im_alpha)
 
             if extra_trimap_output:
                 return im[0], im_alpha[0], trimap[0]
@@ -366,6 +375,34 @@ def shadow(self, im, alpha, darkness=0.9):
 
         return im, alpha_new
 
+    def semitransparency(self, im, alpha):
+
+        # unload removebg and load semitransparency model as cant fit into 1 GPU
+        # removebg GPU to CPU
+        if "cuda" == self.compute_device:
+            self.model_trimap.cpu()
+            self.model_matting.cpu()
+
+            #Load semitransparency model to GPU
+            self.semitransparency_model.model_segmentation.cuda()
+            self.semitransparency_model.model_semitransparency.cuda()
+
+        # semitransparency core
+        removebg_result_4ch = torch.cat((im, alpha), dim=1).float()
+        alpha_new, im_color = self.semitransparency_model(im[0], removebg_result_4ch[0])
+
+        # removebg model CPU to GPU
+        if "cuda" == self.compute_device:
+            # semitransparency model GPU to CPU
+            self.semitransparency_model.model_segmentation.cpu()
+            self.semitransparency_model.model_semitransparency.cpu()
+
+            # removebg model CPU TO GPU
+            self.model_trimap.cuda()
+            self.model_matting.cuda()
+
+        return im_color, alpha_new
+
 
 class Identifier:
     def __init__(self, model_paths, require_models=True, compute_device="cuda"):

diff --git a/core/requirements-deploy.in b/core/requirements-deploy.in
@@ -12,3 +12,4 @@ pika>=1.0.0,<2.0.0
 msgpack<1.0.0
 pandas<1.4.0
 nibabel<2.5.0,>=2.4.0
+semitransparency-core>=0.9.3,<1.0.0
diff --git a/core/requirements-deploy.txt b/core/requirements-deploy.txt
@@ -27,6 +27,7 @@ idna==3.3
 kaleido-core==2.17.0
     # via
     #   -r requirements-deploy.in
+    #   semitransparency-core
     #   shadowgen-core
 msgpack==0.6.2
     # via
@@ -45,12 +46,14 @@ numpy==1.21.5
     #   opencv-python
     #   pandas
     #   scipy
+    #   semitransparency-core
     #   shadowgen-core
     #   torchvision
-opencv-python==4.5.5.62
+opencv-python==4.5.5.64
     # via
     #   -r requirements-deploy.in
     #   kaleido-core
+    #   semitransparency-core
     #   shadowgen-core
 pandas==1.3.5
     # via
@@ -87,7 +90,11 @@ requests==2.27.1
     #   kaleido-core
     #   wandb
 scipy==1.8.0
-    # via kaleido-core
+    # via
+    #   kaleido-core
+    #   semitransparency-core
+semitransparency-core==0.9.3
+    # via -r requirements-deploy.in
 sentry-sdk==1.5.7
     # via wandb
 shadowgen-core==1.10.3
@@ -111,12 +118,14 @@ torch==1.10.0+cu113
     # via
     #   -r requirements-deploy.in
     #   kaleido-core
+    #   semitransparency-core
     #   shadowgen-core
     #   torchvision
 torchvision==0.11.1+cu113
     # via
     #   -r requirements-deploy.in
     #   kaleido-core
+    #   semitransparency-core
     #   shadowgen-core
 tqdm==4.63.0
     # via kaleido-core