Skip to content

Commit 94f9441

Browse files
Taehun KimTaehun Kim
authored andcommitted
persisting background image fixed
1 parent 24fc382 commit 94f9441

File tree

2 files changed

+15
-6
lines changed

2 files changed

+15
-6
lines changed

setup.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77

88
setuptools.setup(
99
name="transparent-background",
10-
version="1.2.8",
10+
version="1.2.9",
1111
author="Taehun Kim",
1212
author_email="taehoon1018@postech.ac.kr",
1313
description="Make images with transparent background",

transparent_background/Remover.py

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -121,7 +121,7 @@ def __init__(self, mode="base", jit=False, device=None, ckpt=None, fast=None):
121121
]
122122
)
123123

124-
self.background = None
124+
self.background = {'img': None, 'name': None, 'shape': None}
125125
desc = "Mode={}, Device={}, Torchscript={}".format(
126126
mode, self.device, "enabled" if jit else "disabled"
127127
)
@@ -198,10 +198,19 @@ def process(self, img, type="rgba", threshold=None):
198198
img[border != 0] = [120, 255, 155]
199199

200200
elif type.lower().endswith((".jpg", ".jpeg", ".png")):
201-
if self.background is None:
202-
self.background = cv2.cvtColor(cv2.imread(type), cv2.COLOR_BGR2RGB)
203-
self.background = cv2.resize(self.background, img.shape[:2][::-1])
204-
img = img * pred[..., np.newaxis] + self.background * (
201+
if self.background['name'] != type:
202+
background_img = cv2.cvtColor(cv2.imread(type), cv2.COLOR_BGR2RGB)
203+
background_img = cv2.resize(background_img, img.shape[:2][::-1])
204+
205+
self.background['img'] = background_img
206+
self.background['shape'] = img.shape[:2][::-1]
207+
self.background['name'] = type
208+
209+
elif self.background['shape'] != img.shape[:2][::-1]:
210+
self.background['img'] = cv2.resize(self.background['img'], img.shape[:2][::-1])
211+
self.background['shape'] = img.shape[:2][::-1]
212+
213+
img = img * pred[..., np.newaxis] + self.background['img'] * (
205214
1 - pred[..., np.newaxis]
206215
)
207216

0 commit comments

Comments
 (0)