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

fix zits

Browse Source
This commit is contained in:
Qing 2022-07-15 17:51:27 +08:00
parent eee91e7b82
commit 01c7f3b77d
2 changed files with 4 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
lama_cleaner/app/src/components/Settings/ModelSettingBlock.tsx
View File

@ -10,7 +10,7 @@ import SettingBlock from './SettingBlock'
export enum AIModel { export enum AIModel {
LAMA = 'lama', LAMA = 'lama',
LDM = 'ldm', LDM = 'ldm',
ZITS = 'ZITS', ZITS = 'zits',
} }
function ModelSettingBlock() { function ModelSettingBlock() {

5
lama_cleaner/model/zits.py
View File

@ -229,8 +229,8 @@ class ZITS(InpaintModel):
def wireframe_edge_and_line(self, items, enable: bool): def wireframe_edge_and_line(self, items, enable: bool):
# 最终向 items 中添加 edge 和 line key # 最终向 items 中添加 edge 和 line key
if not enable: if not enable:
items["edge"] = torch.zeros_like(items["mask_256"]) items["edge"] = torch.zeros_like(items["masks"])
items["line"] = torch.zeros_like(items["mask_256"]) items["line"] = torch.zeros_like(items["masks"])
return return
start = time.time() start = time.time()
@ -297,6 +297,7 @@ class ZITS(InpaintModel):
masks: [H, W] masks: [H, W]
return: BGR IMAGE return: BGR IMAGE
""" """
mask = mask[:, :, 0]
items = load_image(image, mask, device=self.device) items = load_image(image, mask, device=self.device)
self.wireframe_edge_and_line(items, config.zits_wireframe) self.wireframe_edge_and_line(items, config.zits_wireframe)