kamp/IOPaint
1
0
Fork 0
Code Issues Pull Requests Releases 1 Activity

auto save result image when --output-dir exists

Browse Source
This commit is contained in:
Qing 2023-02-19 14:31:00 +08:00
parent 88a37ea904
commit 774f470e58
8 changed files with 103 additions and 38 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
lama_cleaner/app/src/App.tsx
View File

@ -8,12 +8,14 @@ import {
enableFileManagerState,
fileState,
isDisableModelSwitchState,
isEnableAutoSavingState,
toastState,
} from './store/Atoms'
import { keepGUIAlive } from './utils'
import Header from './components/Header/Header'
import useHotKey from './hooks/useHotkey'
import {
getEnableAutoSaving,
getEnableFileManager,
getIsDisableModelSwitch,
isDesktop,
@ -34,6 +36,7 @@ function App() {
const userInputImage = useInputImage()
const setIsDisableModelSwitch = useSetRecoilState(isDisableModelSwitchState)
const setEnableFileManager = useSetRecoilState(enableFileManagerState)
const setIsEnableAutoSavingState = useSetRecoilState(isEnableAutoSavingState)
// Set Input Image
useEffect(() => {
@ -66,7 +69,17 @@ function App() {
setEnableFileManager(isEnabled === 'true')
}
fetchData2()
}, [setEnableFileManager, setIsDisableModelSwitch])
const fetchData3 = async () => {
const isEnabled = await getEnableAutoSaving().then(res => res.text())
setIsEnableAutoSavingState(isEnabled === 'true')
}
fetchData3()
}, [
setEnableFileManager,
setIsDisableModelSwitch,
setIsEnableAutoSavingState,
])
// Dark Mode Hotkey
useHotKey(

6
lama_cleaner/app/src/adapters/inpainting.ts
View File

@ -122,6 +122,12 @@ export function getEnableFileManager() {
})
}
export function getEnableAutoSaving() {
return fetch(`${API_ENDPOINT}/is_enable_auto_saving`, {
method: 'GET',
})
}
export function switchModel(name: string) {
const fd = new FormData()
fd.append('name', name)

7
lama_cleaner/app/src/components/Editor/Editor.tsx
View File

@ -20,7 +20,6 @@ import { useRecoilState, useRecoilValue, useSetRecoilState } from 'recoil'
import { useWindowSize, useKey, useKeyPressEvent } from 'react-use'
import inpaint, {
downloadToOutput,
makeGif,
postInteractiveSeg,
} from '../../adapters/inpainting'
import Button from '../shared/Button'
@ -45,12 +44,11 @@ import {
imageWidthState,
interactiveSegClicksState,
isDiffusionModelsState,
isEnableAutoSavingState,
isInpaintingState,
isInteractiveSegRunningState,
isInteractiveSegState,
isPaintByExampleState,
isPix2PixState,
isSDState,
negativePropmtState,
propmtState,
runManuallyState,
@ -184,6 +182,7 @@ export default function Editor() {
const [redoCurLines, setRedoCurLines] = useState<Line[]>([])
const [redoLineGroups, setRedoLineGroups] = useState<LineGroup[]>([])
const enableFileManager = useRecoilValue(enableFileManagerState)
const isEnableAutoSaving = useRecoilValue(isEnableAutoSavingState)
const setImageWidth = useSetRecoilState(imageWidthState)
const setImageHeight = useSetRecoilState(imageHeightState)
@ -1101,7 +1100,7 @@ export default function Editor() {
if (file === undefined) {
return
}
if (enableFileManager && renders.length > 0) {
if ((enableFileManager || isEnableAutoSaving) && renders.length > 0) {
try {
downloadToOutput(renders[renders.length - 1], file.name, file.type)
setToastState({

14
lama_cleaner/app/src/store/Atoms.tsx
View File

@ -41,6 +41,7 @@ interface AppState {
disableShortCuts: boolean
isInpainting: boolean
isDisableModelSwitch: boolean
isEnableAutoSaving: boolean
isInteractiveSeg: boolean
isInteractiveSegRunning: boolean
interactiveSegClicks: number[][]
@ -58,6 +59,7 @@ export const appState = atom<AppState>({
disableShortCuts: false,
isInpainting: false,
isDisableModelSwitch: false,
isEnableAutoSaving: false,
isInteractiveSeg: false,
isInteractiveSegRunning: false,
interactiveSegClicks: [],
@ -221,6 +223,18 @@ export const isDisableModelSwitchState = selector({
},
})
export const isEnableAutoSavingState = selector({
key: 'isEnableAutoSavingState',
get: ({ get }) => {
const app = get(appState)
return app.isEnableAutoSaving
},
set: ({ get, set }, newValue: any) => {
const app = get(appState)
set(appState, { ...app, isEnableAutoSaving: newValue })
},
})
export const croperState = atom<Rect>({
key: 'croperState',
default: {

2
lama_cleaner/const.py
View File

@ -59,7 +59,7 @@ Model download directory (by setting XDG_CACHE_HOME environment variable), by de
"""
OUTPUT_DIR_HELP = """
Only required when --input is directory. Result images will be saved to output directory automatically.
Result images will be saved to output directory automatically without confirmation.
"""
INPUT_HELP = """

51
lama_cleaner/file_manager/file_manager.py
View File

@ -14,7 +14,7 @@ from PIL import Image, ImageOps, PngImagePlugin
from loguru import logger
LARGE_ENOUGH_NUMBER = 100
PngImagePlugin.MAX_TEXT_CHUNK = LARGE_ENOUGH_NUMBER * (1024 ** 2)
PngImagePlugin.MAX_TEXT_CHUNK = LARGE_ENOUGH_NUMBER * (1024**2)
from .storage_backends import FilesystemStorageBackend
from .utils import aspect_to_string, generate_filename, glob_img
@ -63,11 +63,11 @@ class FileManager(FileSystemEventHandler):
if event.src_path == str(self.root_directory):
logger.info(f"Image directory {event.src_path} modified")
self.image_dir_filenames = self._media_names(self.root_directory)
self.modified_time['image'] = datetime.utcnow()
self.modified_time["image"] = datetime.utcnow()
elif event.src_path == str(self.output_dir):
logger.info(f"Output directory {event.src_path} modified")
self.output_dir_filenames = self._media_names(self.output_dir)
self.modified_time['output'] = datetime.utcnow()
self.modified_time["output"] = datetime.utcnow()
def init_app(self, app):
if self.app is None:
@ -83,21 +83,15 @@ class FileManager(FileSystemEventHandler):
app.extensions["thumbnail"] = self
app.config.setdefault("THUMBNAIL_MEDIA_ROOT", self._default_root_directory)
app.config.setdefault("THUMBNAIL_MEDIA_THUMBNAIL_ROOT", self._default_thumbnail_directory)
app.config.setdefault(
"THUMBNAIL_MEDIA_THUMBNAIL_ROOT", self._default_thumbnail_directory
)
app.config.setdefault("THUMBNAIL_MEDIA_URL", self._default_root_url)
app.config.setdefault("THUMBNAIL_MEDIA_THUMBNAIL_URL", self._default_thumbnail_root_url)
app.config.setdefault(
"THUMBNAIL_MEDIA_THUMBNAIL_URL", self._default_thumbnail_root_url
)
app.config.setdefault("THUMBNAIL_DEFAULT_FORMAT", self._default_format)
def save_to_output_directory(self, image: np.ndarray, filename: str):
fp = Path(filename)
new_name = fp.stem + f"_{int(time.time())}" + fp.suffix
if image.shape[2] == 3:
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
elif image.shape[2] == 4:
image = cv2.cvtColor(image, cv2.COLOR_RGBA2BGRA)
cv2.imwrite(str(self.output_dir / new_name), image)
@property
def root_directory(self):
path = self.app.config["THUMBNAIL_MEDIA_ROOT"]
@ -137,14 +131,23 @@ class FileManager(FileSystemEventHandler):
for name in names:
path = os.path.join(directory, name)
img = Image.open(path)
res.append({"name": name, "height": img.height, "width": img.width, "ctime": os.path.getctime(path)})
res.append(
{
"name": name,
"height": img.height,
"width": img.width,
"ctime": os.path.getctime(path),
}
)
return res
@property
def thumbnail_url(self):
return self.app.config["THUMBNAIL_MEDIA_THUMBNAIL_URL"]
def get_thumbnail(self, directory: Path, original_filename: str, width, height, **options):
def get_thumbnail(
self, directory: Path, original_filename: str, width, height, **options
):
storage = FilesystemStorageBackend(self.app)
crop = options.get("crop", "fit")
background = options.get("background")
@ -163,13 +166,19 @@ class FileManager(FileSystemEventHandler):
thumbnail_size = (width, height)
thumbnail_filename = generate_filename(
original_filename, aspect_to_string(thumbnail_size), crop, background, quality
original_filename,
aspect_to_string(thumbnail_size),
crop,
background,
quality,
)
thumbnail_filepath = os.path.join(
self.thumbnail_directory, original_path, thumbnail_filename
)
thumbnail_url = os.path.join(self.thumbnail_url, original_path, thumbnail_filename)
thumbnail_url = os.path.join(
self.thumbnail_url, original_path, thumbnail_filename
)
if storage.exists(thumbnail_filepath):
return thumbnail_url, (width, height)
@ -183,7 +192,9 @@ class FileManager(FileSystemEventHandler):
# get original image format
options["format"] = options.get("format", image.format)
image = self._create_thumbnail(image, thumbnail_size, crop, background=background)
image = self._create_thumbnail(
image, thumbnail_size, crop, background=background
)
raw_data = self.get_raw_data(image, **options)
storage.save(thumbnail_filepath, raw_data)

19
lama_cleaner/parse_args.py
View File

@ -170,15 +170,14 @@ def parse_args():
parser.error(
f"invalid --input: {args.input} is a directory, --output-dir is required"
)
else:
output_dir = Path(args.output_dir)
if not output_dir.exists():
logger.info(f"Creating output directory: {output_dir}")
output_dir.mkdir(parents=True)
else:
if not output_dir.is_dir():
parser.error(
f"invalid --output-dir: {output_dir} is not a directory"
)
if args.output_dir is not None:
output_dir = Path(args.output_dir)
if not output_dir.exists():
logger.info(f"Creating output directory: {output_dir}")
output_dir.mkdir(parents=True)
else:
if not output_dir.is_dir():
parser.error(f"invalid --output-dir: {output_dir} is not a directory")
return args

27
lama_cleaner/server.py
View File

@ -83,11 +83,13 @@ CORS(app, expose_headers=["Content-Disposition"])
model: ModelManager = None
thumb: FileManager = None
output_dir: str = None
interactive_seg_model: InteractiveSeg = None
device = None
input_image_path: str = None
is_disable_model_switch: bool = False
is_enable_file_manager: bool = False
is_enable_auto_saving: bool = False
is_desktop: bool = False
@ -124,11 +126,20 @@ def make_gif():
@app.route("/save_image", methods=["POST"])
def save_image():
# all image in output directory
if output_dir is None:
return "--output-dir is None", 500
input = request.files
filename = request.form["filename"]
origin_image_bytes = input["image"].read() # RGB
image, _ = load_img(origin_image_bytes)
thumb.save_to_output_directory(image, request.form["filename"])
if image.shape[2] == 3:
image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
elif image.shape[2] == 4:
image = cv2.cvtColor(image, cv2.COLOR_RGBA2BGRA)
cv2.imwrite(os.path.join(output_dir, filename), image)
return "ok", 200
@ -348,6 +359,12 @@ def get_is_enable_file_manager():
return res, 200
@app.route("/is_enable_auto_saving")
def get_is_enable_auto_saving():
res = "true" if is_enable_auto_saving else "false"
return res, 200
@app.route("/model_downloaded/<name>")
def model_downloaded(name):
return str(model.is_downloaded(name)), 200
@ -408,6 +425,12 @@ def main(args):
global is_enable_file_manager
global is_desktop
global thumb
global output_dir
global is_enable_auto_saving
output_dir = args.output_dir
if output_dir is not None:
is_enable_auto_saving = True
device = torch.device(args.device)
is_disable_model_switch = args.disable_model_switch