AI/rembg
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/danielgatis/rembg synced 2025-08-05 20:26:10 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix lint and refactored normalizing

Browse Source
This commit is contained in:
Flippchen 2023-04-20 11:36:20 +02:00
parent bb3c58f411
commit ff38b9a377
2 changed files with 11 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
rembg/session_factory.py
View File

@ -88,12 +88,12 @@ def new_session(model_name: str = "u2net") -> BaseSession:
ort.InferenceSession(
str(path / fname_encoder),
providers=ort.get_available_providers(),
sess_options=sess_opts
sess_options=sess_opts,
),
ort.InferenceSession(
str(path / fname_decoder),
providers=ort.get_available_providers(),
sess_options=sess_opts
sess_options=sess_opts,
),
)

18
rembg/session_sam.py
View File

@ -53,7 +53,7 @@ class SamSession(BaseSession):
self,
model_name: str,
encoder: ort.InferenceSession,
decoder: ort.InferenceSession
decoder: ort.InferenceSession,
):
super().__init__(model_name, encoder)
self.decoder = decoder
@ -61,12 +61,12 @@ class SamSession(BaseSession):
def normalize(
self,
img: numpy.ndarray,
mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225),
size=(1024, 1024)
mean=(123.675, 116.28, 103.53),
std=(58.395, 57.12, 57.375),
size=(1024, 1024),
):
pixel_mean = np.array([123.675, 116.28, 103.53]).reshape(1, 1, -1)
pixel_std = np.array([58.395, 57.12, 57.375]).reshape(1, 1, -1)
pixel_mean = np.array([*mean]).reshape(1, 1, -1)
pixel_std = np.array([*std]).reshape(1, 1, -1)
x = (img - pixel_mean) / pixel_std
return x
@ -74,7 +74,7 @@ class SamSession(BaseSession):
self,
img: PILImage,
input_point=np.array([[500, 375]]),
input_label=np.array([1])
input_label=np.array([1]),
) -> List[PILImage]:
# Preprocess image
image = resize_longes_side(img)
@ -90,10 +90,10 @@ class SamSession(BaseSession):
# Add a batch index, concatenate a padding point, and transform.
onnx_coord = np.concatenate([input_point, np.array([[0.0, 0.0]])], axis=0)[
None, :, :
None, :, :
]
onnx_label = np.concatenate([input_label, np.array([-1])], axis=0)[
None, :
None, :
].astype(np.float32)
onnx_coord = apply_coords(onnx_coord, img.size[::1], 1024).astype(np.float32)