Spaces:
Running
Running
File size: 4,696 Bytes
5f57808 c8f8b0e 5f57808 c8f8b0e 5f57808 c8f8b0e 5f57808 c8f8b0e 5f57808 c8f8b0e 5f57808 c8f8b0e 5f57808 c8f8b0e 5f57808 c8f8b0e 5f57808 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 |
import asyncio
import io
import json
import os
import sys
from typing import IO
import click
from PIL.Image import Image as PILImage
from ..bg import remove
from ..session_factory import new_session
from ..sessions import sessions_names
@click.command( # type: ignore
name="b",
help="for a byte stream as input",
)
@click.option(
"-m",
"--model",
default="u2net",
type=click.Choice(sessions_names),
show_default=True,
show_choices=True,
help="model name",
)
@click.option(
"-a",
"--alpha-matting",
is_flag=True,
show_default=True,
help="use alpha matting",
)
@click.option(
"-af",
"--alpha-matting-foreground-threshold",
default=240,
type=int,
show_default=True,
help="trimap fg threshold",
)
@click.option(
"-ab",
"--alpha-matting-background-threshold",
default=10,
type=int,
show_default=True,
help="trimap bg threshold",
)
@click.option(
"-ae",
"--alpha-matting-erode-size",
default=10,
type=int,
show_default=True,
help="erode size",
)
@click.option(
"-om",
"--only-mask",
is_flag=True,
show_default=True,
help="output only the mask",
)
@click.option(
"-ppm",
"--post-process-mask",
is_flag=True,
show_default=True,
help="post process the mask",
)
@click.option(
"-bgc",
"--bgcolor",
default=(0, 0, 0, 0),
type=(int, int, int, int),
nargs=4,
help="Background color (R G B A) to replace the removed background with",
)
@click.option("-x", "--extras", type=str)
@click.option(
"-o",
"--output_specifier",
type=str,
help="printf-style specifier for output filenames (e.g. 'output-%d.png'))",
)
@click.argument(
"image_width",
type=int,
)
@click.argument(
"image_height",
type=int,
)
def b_command(
model: str,
extras: str,
image_width: int,
image_height: int,
output_specifier: str,
**kwargs
) -> None:
"""
Command-line interface for processing images by removing the background using a specified model and generating a mask.
This CLI command takes several options and arguments to configure the background removal process and save the processed images.
Parameters:
model (str): The name of the model to use for background removal.
extras (str): Additional options in JSON format that can be passed to customize the background removal process.
image_width (int): The width of the input images in pixels.
image_height (int): The height of the input images in pixels.
output_specifier (str): A printf-style specifier for the output filenames. If specified, the processed images will be saved to the specified output directory with filenames generated using the specifier.
**kwargs: Additional keyword arguments that can be used to customize the background removal process.
Returns:
None
"""
try:
kwargs.update(json.loads(extras))
except Exception:
pass
session = new_session(model, **kwargs)
bytes_per_img = image_width * image_height * 3
if output_specifier:
output_dir = os.path.dirname(
os.path.abspath(os.path.expanduser(output_specifier))
)
if not os.path.isdir(output_dir):
os.makedirs(output_dir, exist_ok=True)
def img_to_byte_array(img: PILImage) -> bytes:
buff = io.BytesIO()
img.save(buff, format="PNG")
return buff.getvalue()
async def connect_stdin_stdout():
loop = asyncio.get_event_loop()
reader = asyncio.StreamReader()
protocol = asyncio.StreamReaderProtocol(reader)
await loop.connect_read_pipe(lambda: protocol, sys.stdin)
w_transport, w_protocol = await loop.connect_write_pipe(
asyncio.streams.FlowControlMixin, sys.stdout
)
writer = asyncio.StreamWriter(w_transport, w_protocol, reader, loop)
return reader, writer
async def main():
reader, writer = await connect_stdin_stdout()
idx = 0
while True:
try:
img_bytes = await reader.readexactly(bytes_per_img)
if not img_bytes:
break
img = PILImage.frombytes("RGB", (image_width, image_height), img_bytes)
output = remove(img, session=session, **kwargs)
if output_specifier:
output.save((output_specifier % idx), format="PNG")
else:
writer.write(img_to_byte_array(output))
idx += 1
except asyncio.IncompleteReadError:
break
asyncio.run(main())
|