Spaces:
Running
Running
File size: 2,624 Bytes
3faa99b c8f8b0e 3faa99b c8f8b0e 3faa99b c8f8b0e 3faa99b c8f8b0e |
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 |
import json
import sys
from typing import IO
import click
from ..bg import remove
from ..session_factory import new_session
from ..sessions import sessions_names
@click.command( # type: ignore
name="i",
help="for a file 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.argument(
"input", default=(None if sys.stdin.isatty() else "-"), type=click.File("rb")
)
@click.argument(
"output",
default=(None if sys.stdin.isatty() else "-"),
type=click.File("wb", lazy=True),
)
def i_command(model: str, extras: str, input: IO, output: IO, **kwargs) -> None:
"""
Click command line interface function to process an input file based on the provided options.
This function is the entry point for the CLI program. It reads an input file, applies image processing operations based on the provided options, and writes the output to a file.
Parameters:
model (str): The name of the model to use for image processing.
extras (str): Additional options in JSON format.
input: The input file to process.
output: The output file to write the processed image to.
**kwargs: Additional keyword arguments corresponding to the command line options.
Returns:
None
"""
try:
kwargs.update(json.loads(extras))
except Exception:
pass
output.write(remove(input.read(), session=new_session(model, **kwargs), **kwargs))
|