add batch_processing

This commit is contained in:
Qing 2024-01-04 21:39:34 +08:00
parent cd82b21cd9
commit 4a9f2ab03c
3 changed files with 150 additions and 8 deletions

View File

@ -0,0 +1,120 @@
import json
import cv2
from pathlib import Path
from typing import Dict, Optional
from PIL import Image
from loguru import logger
from rich.console import Console
from rich.progress import (
Progress,
SpinnerColumn,
TimeElapsedColumn,
MofNCompleteColumn,
TextColumn,
BarColumn,
TaskProgressColumn,
TimeRemainingColumn,
)
from lama_cleaner.helper import pil_to_bytes
from lama_cleaner.model_manager import ModelManager
from lama_cleaner.schema import InpaintRequest
def glob_images(path: Path) -> Dict[str, Path]:
# png/jpg/jpeg
if path.is_file():
return {path.stem: path}
elif path.is_dir():
res = {}
for it in path.glob("*.*"):
if it.suffix.lower() in [".png", ".jpg", ".jpeg"]:
res[it.stem] = it
return res
def batch_inpaint(
model: str,
device,
image: Path,
mask: Path,
output: Path,
config: Optional[Path] = None,
concat: bool = False,
):
if image.is_dir() and output.is_file():
logger.error(
f"invalid --output: when image is a directory, output should be a directory"
)
exit(-1)
image_paths = glob_images(image)
mask_paths = glob_images(mask)
if len(image_paths) == 0:
logger.error(f"invalid --image: empty image folder")
exit(-1)
if len(mask_paths) == 0:
logger.error(f"invalid --mask: empty mask folder")
exit(-1)
if config is None:
inpaint_request = InpaintRequest()
logger.info(f"Using default config: {inpaint_request}")
else:
with open(config, "r", encoding="utf-8") as f:
inpaint_request = InpaintRequest(**json.load(f))
model_manager = ModelManager(name=model, device=device)
first_mask = list(mask_paths.values())[0]
console = Console()
with Progress(
SpinnerColumn(),
TextColumn("[progress.description]{task.description}"),
BarColumn(),
TaskProgressColumn(),
MofNCompleteColumn(),
TimeElapsedColumn(),
console=console,
transient=False,
) as progress:
task = progress.add_task("Batch processing...", total=len(image_paths))
for stem, image_p in image_paths.items():
if stem not in mask_paths and mask.is_dir():
progress.log(f"mask for {image_p} not found")
progress.update(task, advance=1)
continue
mask_p = mask_paths.get(stem, first_mask)
infos = Image.open(image_p).info
img = cv2.imread(str(image_p))
img = cv2.cvtColor(img, cv2.COLOR_BGRA2RGB)
mask_img = cv2.imread(str(mask_p), cv2.IMREAD_GRAYSCALE)
if mask_img.shape[:2] != img.shape[:2]:
progress.log(
f"resize mask {mask_p.name} to image {image_p.name} size: {img.shape[:2]}"
)
mask_img = cv2.resize(
mask_img,
(img.shape[1], img.shape[0]),
interpolation=cv2.INTER_NEAREST,
)
mask_img[mask_img >= 127] = 255
mask_img[mask_img < 127] = 0
# bgr
inpaint_result = model_manager(img, mask_img, inpaint_request)
inpaint_result = cv2.cvtColor(inpaint_result, cv2.COLOR_BGR2RGB)
if concat:
mask_img = cv2.cvtColor(mask_img, cv2.COLOR_GRAY2RGB)
inpaint_result = cv2.hconcat([img, mask_img, inpaint_result])
img_bytes = pil_to_bytes(Image.fromarray(inpaint_result), "png", 100, infos)
save_p = output / f"{stem}.png"
with open(save_p, "wb") as fw:
fw.write(img_bytes)
progress.update(task, advance=1)

View File

@ -1,4 +1,5 @@
from pathlib import Path
from typing import Dict
import typer
from fastapi import FastAPI
@ -39,15 +40,36 @@ def list_model(
print(it.name)
@typer_app.command(help="Processing image with lama cleaner")
@typer_app.command(help="Batch processing images")
def run(
input: Path = Option(..., help="Image file or folder containing images"),
output_dir: Path = Option(..., help="Output directory"),
config_path: Path = Option(..., help="Config file path"),
model: str = Option("lama"),
device: Device = Option(Device.cpu),
image: Path = Option(..., help="Image folders or file path"),
mask: Path = Option(
...,
help="Mask folders or file path. "
"If it is a directory, the mask images in the directory should have the same name as the original image."
"If it is a file, all images will use this mask."
"Mask will automatically resize to the same size as the original image.",
),
output: Path = Option(..., help="Output directory or file path"),
config: Path = Option(
None, help="Config file path. You can use dump command to create a base config."
),
concat: bool = Option(
False, help="Concat original image, mask and output images into one image"
),
model_dir: Path = Option(DEFAULT_MODEL_DIR, help=MODEL_DIR_HELP, file_okay=False),
):
setup_model_dir(model_dir)
pass
scanned_models = scan_models()
if model not in [it.name for it in scanned_models]:
logger.info(f"{model} not found in {model_dir}, try to downloading")
cli_download_model(model, model_dir)
from lama_cleaner.batch_processing import batch_inpaint
batch_inpaint(model, device, image, mask, output, config, concat)
@typer_app.command(help="Start lama cleaner server")

View File

@ -65,7 +65,7 @@ class InpaintModel:
mask, mod=self.pad_mod, square=self.pad_to_square, min_size=self.min_size
)
logger.info(f"final forward pad size: {pad_image.shape}")
# logger.info(f"final forward pad size: {pad_image.shape}")
image, mask = self.forward_pre_process(image, mask, config)
@ -93,7 +93,7 @@ class InpaintModel:
return: BGR IMAGE
"""
inpaint_result = None
logger.info(f"hd_strategy: {config.hd_strategy}")
# logger.info(f"hd_strategy: {config.hd_strategy}")
if config.hd_strategy == HDStrategy.CROP:
if max(image.shape) > config.hd_strategy_crop_trigger_size:
logger.info(f"Run crop strategy")
@ -189,7 +189,7 @@ class InpaintModel:
crop_img = image[t:b, l:r, :]
crop_mask = mask[t:b, l:r]
logger.info(f"box size: ({box_h},{box_w}) crop size: {crop_img.shape}")
# logger.info(f"box size: ({box_h},{box_w}) crop size: {crop_img.shape}")
return crop_img, crop_mask, [l, t, r, b]