AI/rembg
1
0
Fork 0
mirror of https://git.mirrors.martin98.com/https://github.com/danielgatis/rembg synced 2025-08-15 23:36:03 +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

15
rembg/commands/b_command.py
View File

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