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

fixed b_command (and raise error if -x is given but invalid)

Browse Source
This commit is contained in:
Lukas Bürgi 2025-01-27 12:13:07 +01:00
parent 37bbf63500
commit c4c4e60acf
1 changed files with 8 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
rembg/commands/b_command.py
View File

@ -6,7 +6,7 @@ import sys
from typing import IO from typing import IO
import click import click
from PIL.Image import Image as PILImage import PIL
from ..bg import remove from ..bg import remove
from ..session_factory import new_session from ..session_factory import new_session
@ -118,10 +118,11 @@ def b_command(
Returns: Returns:
None None
""" """
if extras:
try: try:
kwargs.update(json.loads(extras)) kwargs.update(json.loads(extras))
except Exception: except Exception:
pass raise click.BadParameter("extras must be a valid JSON string")
session = new_session(model, **kwargs) session = new_session(model, **kwargs)
bytes_per_img = image_width * image_height * 3 bytes_per_img = image_width * image_height * 3
@ -134,7 +135,7 @@ def b_command(
if not os.path.isdir(output_dir): if not os.path.isdir(output_dir):
os.makedirs(output_dir, exist_ok=True) os.makedirs(output_dir, exist_ok=True)
def img_to_byte_array(img: PILImage) -> bytes: def img_to_byte_array(img: PIL.Image) -> bytes:
buff = io.BytesIO() buff = io.BytesIO()
img.save(buff, format="PNG") img.save(buff, format="PNG")
return buff.getvalue() return buff.getvalue()
@ -162,7 +163,7 @@ def b_command(
if not img_bytes: if not img_bytes:
break break
img = PILImage.frombytes("RGB", (image_width, image_height), img_bytes) img = PIL.Image.frombytes("RGB", (image_width, image_height), img_bytes)
output = remove(img, session=session, **kwargs) output = remove(img, session=session, **kwargs)
if output_specifier: if output_specifier: