let5see/IOPaint
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

auto switch mps device to cpu device

Browse Source
This commit is contained in:
Qing
2023-02-11 13:30:09 +08:00
parent f9b5dcbfd7
commit 8f8bcfe0f4
15 changed files with 52 additions and 19 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
lama_cleaner/model/opencv2.py
View File

@@ -2,12 +2,11 @@ import cv2
from lama_cleaner.model.base import InpaintModel
from lama_cleaner.schema import Config
flag_map = {
"INPAINT_NS": cv2.INPAINT_NS,
"INPAINT_TELEA": cv2.INPAINT_TELEA
}
flag_map = {"INPAINT_NS": cv2.INPAINT_NS, "INPAINT_TELEA": cv2.INPAINT_TELEA}
class OpenCV2(InpaintModel):
name = "cv2"
pad_mod = 1
@staticmethod
@@ -20,5 +19,10 @@ class OpenCV2(InpaintModel):
mask: [H, W, 1]
return: BGR IMAGE
"""
cur_res = cv2.inpaint(image[:,:,::-1], mask, inpaintRadius=config.cv2_radius, flags=flag_map[config.cv2_flag])
cur_res = cv2.inpaint(
image[:, :, ::-1],
mask,
inpaintRadius=config.cv2_radius,
flags=flag_map[config.cv2_flag],
)
return cur_res