From 2119a5f90522157fdcfc1bbd54d8c3814463692c Mon Sep 17 00:00:00 2001 From: Qing Date: Sun, 4 Sep 2022 16:00:42 +0800 Subject: [PATCH] FcF use unique resize strategy --- .../components/Settings/HDSettingBlock.tsx | 3 ++ .../components/Settings/ModelSettingBlock.tsx | 42 +++++++++++++------ lama_cleaner/app/src/store/Atoms.tsx | 13 ++++++ lama_cleaner/model/base.py | 23 +++++++--- lama_cleaner/model/fcf.py | 34 ++++++++++++++- lama_cleaner/tests/test_model.py | 8 ++++ 6 files changed, 105 insertions(+), 18 deletions(-) diff --git a/lama_cleaner/app/src/components/Settings/HDSettingBlock.tsx b/lama_cleaner/app/src/components/Settings/HDSettingBlock.tsx index 5ce001b..d3bb401 100644 --- a/lama_cleaner/app/src/components/Settings/HDSettingBlock.tsx +++ b/lama_cleaner/app/src/components/Settings/HDSettingBlock.tsx @@ -18,6 +18,9 @@ export enum LDMSampler { function HDSettingBlock() { const [hdSettings, setHDSettings] = useRecoilState(hdSettingsState) + if (!hdSettings.enabled) { + return <> + } const onStrategyChange = (value: HDStrategy) => { setHDSettings({ hdStrategy: value }) diff --git a/lama_cleaner/app/src/components/Settings/ModelSettingBlock.tsx b/lama_cleaner/app/src/components/Settings/ModelSettingBlock.tsx index af55fdf..386202c 100644 --- a/lama_cleaner/app/src/components/Settings/ModelSettingBlock.tsx +++ b/lama_cleaner/app/src/components/Settings/ModelSettingBlock.tsx @@ -30,17 +30,6 @@ function ModelSettingBlock() { ) => { return (
- {/* - - Paper - - */} - + + {/* + + Paper + + */}
) } @@ -123,6 +123,16 @@ function ModelSettingBlock() { ) } + const renderFCFModelDesc = () => { + return ( +
+ FcF only support fixed size(512x512) image input. Lama Cleaner will take + care of resize and crop process, it still recommended applies to small + defects. +
+ ) + } + const renderOptionDesc = (): ReactNode => { switch (setting.model) { case AIModel.LAMA: @@ -133,6 +143,8 @@ function ModelSettingBlock() { return renderZITSModelDesc() case AIModel.MAT: return undefined + case AIModel.FCF: + return renderFCFModelDesc() default: return <> } @@ -161,9 +173,15 @@ function ModelSettingBlock() { case AIModel.MAT: return renderModelDesc( 'Mask-Aware Transformer for Large Hole Image Inpainting', - 'https://arxiv.org/pdf/2203.15270.pdf', + 'https://arxiv.org/abs/2203.15270', 'https://github.com/fenglinglwb/MAT' ) + case AIModel.FCF: + return renderModelDesc( + 'Keys to Better Image Inpainting: Structure and Texture Go Hand in Hand', + 'https://arxiv.org/abs/2208.03382', + 'https://github.com/SHI-Labs/FcF-Inpainting' + ) default: return <> } diff --git a/lama_cleaner/app/src/store/Atoms.tsx b/lama_cleaner/app/src/store/Atoms.tsx index 927427c..05c6c08 100644 --- a/lama_cleaner/app/src/store/Atoms.tsx +++ b/lama_cleaner/app/src/store/Atoms.tsx @@ -8,6 +8,7 @@ export enum AIModel { LDM = 'ldm', ZITS = 'zits', MAT = 'mat', + FCF = 'fcf', } export const fileState = atom({ @@ -42,6 +43,7 @@ export interface HDSettings { hdStrategyResizeLimit: number hdStrategyCropTrigerSize: number hdStrategyCropMargin: number + enabled: boolean } type ModelsHDSettings = { [key in AIModel]: HDSettings } @@ -68,24 +70,35 @@ const defaultHDSettings: ModelsHDSettings = { hdStrategyResizeLimit: 2048, hdStrategyCropTrigerSize: 2048, hdStrategyCropMargin: 128, + enabled: true, }, [AIModel.LDM]: { hdStrategy: HDStrategy.CROP, hdStrategyResizeLimit: 1080, hdStrategyCropTrigerSize: 1080, hdStrategyCropMargin: 128, + enabled: true, }, [AIModel.ZITS]: { hdStrategy: HDStrategy.CROP, hdStrategyResizeLimit: 1024, hdStrategyCropTrigerSize: 1024, hdStrategyCropMargin: 128, + enabled: true, }, [AIModel.MAT]: { hdStrategy: HDStrategy.CROP, hdStrategyResizeLimit: 1024, hdStrategyCropTrigerSize: 512, hdStrategyCropMargin: 128, + enabled: true, + }, + [AIModel.FCF]: { + hdStrategy: HDStrategy.CROP, + hdStrategyResizeLimit: 512, + hdStrategyCropTrigerSize: 512, + hdStrategyCropMargin: 128, + enabled: false, }, } diff --git a/lama_cleaner/model/base.py b/lama_cleaner/model/base.py index df769a6..1cb8b57 100644 --- a/lama_cleaner/model/base.py +++ b/lama_cleaner/model/base.py @@ -4,7 +4,6 @@ from typing import Optional import cv2 import torch from loguru import logger -import numpy as np from lama_cleaner.helper import boxes_from_mask, resize_max_size, pad_img_to_modulo from lama_cleaner.schema import Config, HDStrategy @@ -92,7 +91,6 @@ class InpaintModel: inpaint_result = cv2.resize(inpaint_result, (origin_size[1], origin_size[0]), interpolation=cv2.INTER_CUBIC) - original_pixel_indices = mask < 127 inpaint_result[original_pixel_indices] = image[:, :, ::-1][original_pixel_indices] @@ -101,7 +99,7 @@ class InpaintModel: return inpaint_result - def _run_box(self, image, mask, box, config: Config): + def _crop_box(self, image, mask, box, config: Config): """ Args: @@ -110,7 +108,7 @@ class InpaintModel: box: [left,top,right,bottom] Returns: - BGR IMAGE + BGR IMAGE, (l, r, r, b) """ box_h = box[3] - box[1] box_w = box[2] - box[0] @@ -131,7 +129,7 @@ class InpaintModel: t = max(_t, 0) b = min(_b, img_h) - # try to get more context when crop around image edge + # try to get more context when crop around image edge if _l < 0: r += abs(_l) if _r > img_w: @@ -151,4 +149,19 @@ class InpaintModel: logger.info(f"box size: ({box_h},{box_w}) crop size: {crop_img.shape}") + return crop_img, crop_mask, [l, t, r, b] + + def _run_box(self, image, mask, box, config: Config): + """ + + Args: + image: [H, W, C] RGB + mask: [H, W, 1] + box: [left,top,right,bottom] + + Returns: + BGR IMAGE + """ + crop_img, crop_mask, [l, t, r, b] = self._crop_box(image, mask, box, config) + return self._pad_forward(crop_img, crop_mask, config), [l, t, r, b] diff --git a/lama_cleaner/model/fcf.py b/lama_cleaner/model/fcf.py index ce52b21..06f1356 100644 --- a/lama_cleaner/model/fcf.py +++ b/lama_cleaner/model/fcf.py @@ -8,7 +8,7 @@ import torch.fft as fft from lama_cleaner.schema import Config -from lama_cleaner.helper import load_model, get_cache_path_by_url, norm_img +from lama_cleaner.helper import load_model, get_cache_path_by_url, norm_img, boxes_from_mask, resize_max_size from lama_cleaner.model.base import InpaintModel from torch import conv2d, nn import torch.nn.functional as F @@ -1154,6 +1154,38 @@ class FcF(InpaintModel): def is_downloaded() -> bool: return os.path.exists(get_cache_path_by_url(FCF_MODEL_URL)) + @torch.no_grad() + def __call__(self, image, mask, config: Config): + """ + images: [H, W, C] RGB, not normalized + masks: [H, W] + return: BGR IMAGE + """ + boxes = boxes_from_mask(mask) + crop_result = [] + config.hd_strategy_crop_margin = 128 + for box in boxes: + crop_image, crop_mask, crop_box = self._crop_box(image, mask, box, config) + origin_size = crop_image.shape[:2] + resize_image = resize_max_size(crop_image, size_limit=512) + resize_mask = resize_max_size(crop_mask, size_limit=512) + inpaint_result = self._pad_forward(resize_image, resize_mask, config) + + # only paste masked area result + inpaint_result = cv2.resize(inpaint_result, (origin_size[1], origin_size[0]), interpolation=cv2.INTER_CUBIC) + + original_pixel_indices = crop_mask < 127 + inpaint_result[original_pixel_indices] = crop_image[:, :, ::-1][original_pixel_indices] + + crop_result.append((inpaint_result, crop_box)) + + inpaint_result = image[:, :, ::-1] + for crop_image, crop_box in crop_result: + x1, y1, x2, y2 = crop_box + inpaint_result[y1:y2, x1:x2, :] = crop_image + + return inpaint_result + def forward(self, image, mask, config: Config): """Input images and output images have same size images: [H, W, C] RGB diff --git a/lama_cleaner/tests/test_model.py b/lama_cleaner/tests/test_model.py index 3eb98be..3050891 100644 --- a/lama_cleaner/tests/test_model.py +++ b/lama_cleaner/tests/test_model.py @@ -143,3 +143,11 @@ def test_fcf(strategy): fx=2, fy=2 ) + + assert_equal( + model, + cfg, + f"fcf_{strategy.capitalize()}_result.png", + fx=3.8, + fy=2 + )