AI/rembg
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/danielgatis/rembg synced 2025-08-16 18:45:54 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #715 from lukas-buergi/fix_b

Browse Source 
fixed b_command
This commit is contained in:
Daniel Gatis 2025-01-30 12:57:40 -03:00 committed by GitHub
commit e1537eeea9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
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: