rembg/commands/b_command.py

import asyncio
import io
import json
import os
import sys
from typing import IO

import click
from PIL import Image

from ..bg import remove
from ..session_factory import new_session
from ..sessions import sessions_names


@click.command(
    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=None,
    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 rs_command(
    model: str,
    extras: str,
    image_width: int,
    image_height: int,
    output_specifier: str,
    **kwargs
) -> None:
    try:
        kwargs.update(json.loads(extras))
    except Exception:
        pass

    session = new_session(model)
    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: Image) -> 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 = Image.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())