You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
66 lines
1.8 KiB
Python
66 lines
1.8 KiB
Python
"""Command for colorizing images with AI"""
|
|
|
|
import logging
|
|
|
|
import click
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@click.argument("image_filepaths", nargs=-1)
|
|
@click.option(
|
|
"--outdir",
|
|
default="./outputs/colorized",
|
|
show_default=True,
|
|
type=click.Path(),
|
|
help="Where to write results to.",
|
|
)
|
|
@click.option(
|
|
"-r",
|
|
"--repeats",
|
|
default=1,
|
|
show_default=True,
|
|
type=int,
|
|
help="How many times to repeat the renders. If you provide two prompts and --repeat=3 then six images will be generated.",
|
|
)
|
|
@click.option(
|
|
"--caption",
|
|
default="",
|
|
show_default=False,
|
|
help="Description of the photo. If not provided, it will be generated automatically.",
|
|
)
|
|
@click.command("colorize")
|
|
def colorize_cmd(image_filepaths, outdir, repeats, caption):
|
|
"""
|
|
Colorize images using AI. Doesn't work very well yet.
|
|
"""
|
|
import os.path
|
|
|
|
from tqdm import tqdm
|
|
|
|
from imaginairy.api.colorize import colorize_img
|
|
from imaginairy.schema import LazyLoadingImage
|
|
from imaginairy.utils.log_utils import configure_logging
|
|
|
|
configure_logging()
|
|
|
|
os.makedirs(outdir, exist_ok=True)
|
|
base_count = len(os.listdir(outdir))
|
|
for _ in range(repeats):
|
|
for p in tqdm(image_filepaths):
|
|
base_count += 1
|
|
filename = f"{base_count:06d}_{os.path.basename(p)}".lower()
|
|
savepath = os.path.join(outdir, filename)
|
|
if p.startswith("http"):
|
|
img = LazyLoadingImage(url=p)
|
|
elif os.path.isdir(p):
|
|
print(f"Skipping directory: {p}")
|
|
continue
|
|
else:
|
|
img = LazyLoadingImage(filepath=p)
|
|
logger.info(f"Colorizing {p} and saving it to {savepath}")
|
|
|
|
img = colorize_img(img, caption=caption)
|
|
|
|
img.save(savepath)
|