diff --git a/inference.py b/inference.py index e698164..8b8a009 100644 --- a/inference.py +++ b/inference.py @@ -1,86 +1,86 @@ -import os, sys, argparse, warnings -import torch, torch.nn as nn, torch.nn.functional as F -import numpy as np, cv2 -from PIL import Image -from model import DocScanner -from seg import U2NETP - -warnings.filterwarnings('ignore') - -# Запоминаем CWD пользователя ДО смены директории -USER_CWD = os.getcwd() -SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) -os.chdir(SCRIPT_DIR) - - -class Net(nn.Module): - def __init__(self): - super().__init__() - self.msk = U2NETP(3, 1) - self.bm = DocScanner() - - def forward(self, x): - msk, *_ = self.msk(x) - bm = self.bm((msk > 0.5).float() * x, iters=12, test_mode=True) - return (2 * (bm / 286.8) - 1) * 0.99 - - -def load_model(model, path, strip_prefix=False): - state_dict = model.state_dict() - pretrained = torch.load(path, map_location='cuda:0') - if strip_prefix: - pretrained = {k[6:]: v for k, v in pretrained.items() if k[6:] in state_dict} - else: - pretrained = {k: v for k, v in pretrained.items() if k in state_dict} - state_dict.update(pretrained) - model.load_state_dict(state_dict) - return model - - -def resolve_user_path(path): - """Если путь относительный — считаем его относительно CWD пользователя, не скрипта""" - if os.path.isabs(path): - return path - return os.path.join(USER_CWD, path) - - -def main(): - parser = argparse.ArgumentParser() - parser.add_argument('-i', '--input', required=True) - parser.add_argument('-o', '--output', required=True) - opt = parser.parse_args() - - input_path = resolve_user_path(opt.input) - output_path = resolve_user_path(opt.output) - - net = Net().cuda().eval() - load_model(net.msk, f'{SCRIPT_DIR}/model_pretrained/seg.pth', strip_prefix=True) - load_model(net.bm, f'{SCRIPT_DIR}/model_pretrained/DocScanner-L.pth') - - image = np.array(Image.open(input_path))[:, :, :3] / 255.0 - height, width = image.shape[:2] - tensor = torch.from_numpy(cv2.resize(image, (288, 288)).transpose(2, 0, 1)).float().unsqueeze(0) - - with torch.no_grad(): - bm = net(tensor.cuda()).cpu() - - flow = torch.from_numpy(np.stack([ - cv2.blur(cv2.resize(bm[0, 0].numpy(), (width, height)), (3, 3)), - cv2.blur(cv2.resize(bm[0, 1].numpy(), (width, height)), (3, 3)) - ], axis=2)).unsqueeze(0) - - out = F.grid_sample( - torch.from_numpy(image).permute(2, 0, 1).unsqueeze(0).float(), - flow, align_corners=True - ) - result = ((out[0] * 255).permute(1, 2, 0).numpy()[:, :, ::-1]).astype(np.uint8) - - out_dir = os.path.dirname(output_path) - if out_dir: - os.makedirs(out_dir, exist_ok=True) - cv2.imwrite(output_path, result) - print(f"[OK] {output_path}") - - -if __name__ == "__main__": - main() +import os, sys, argparse, warnings +import torch, torch.nn as nn, torch.nn.functional as F +import numpy as np, cv2 +from PIL import Image +from model import DocScanner +from seg import U2NETP + +warnings.filterwarnings('ignore') + +# Запоминаем CWD пользователя ДО смены директории +USER_CWD = os.getcwd() +SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) +os.chdir(SCRIPT_DIR) + + +class Net(nn.Module): + def __init__(self): + super().__init__() + self.msk = U2NETP(3, 1) + self.bm = DocScanner() + + def forward(self, x): + msk, *_ = self.msk(x) + bm = self.bm((msk > 0.5).float() * x, iters=12, test_mode=True) + return (2 * (bm / 286.8) - 1) * 0.99 + + +def load_model(model, path, strip_prefix=False): + state_dict = model.state_dict() + pretrained = torch.load(path, map_location='cuda:0') + if strip_prefix: + pretrained = {k[6:]: v for k, v in pretrained.items() if k[6:] in state_dict} + else: + pretrained = {k: v for k, v in pretrained.items() if k in state_dict} + state_dict.update(pretrained) + model.load_state_dict(state_dict) + return model + + +def resolve_user_path(path): + """Если путь относительный — считаем его относительно CWD пользователя, не скрипта""" + if os.path.isabs(path): + return path + return os.path.join(USER_CWD, path) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('-i', '--input', required=True) + parser.add_argument('-o', '--output', required=True) + opt = parser.parse_args() + + input_path = resolve_user_path(opt.input) + output_path = resolve_user_path(opt.output) + + net = Net().cuda().eval() + load_model(net.msk, f'{SCRIPT_DIR}/model_pretrained/seg.pth', strip_prefix=True) + load_model(net.bm, f'{SCRIPT_DIR}/model_pretrained/DocScanner-L.pth') + + image = np.array(Image.open(input_path))[:, :, :3] / 255.0 + height, width = image.shape[:2] + tensor = torch.from_numpy(cv2.resize(image, (288, 288)).transpose(2, 0, 1)).float().unsqueeze(0) + + with torch.no_grad(): + bm = net(tensor.cuda()).cpu() + + flow = torch.from_numpy(np.stack([ + cv2.blur(cv2.resize(bm[0, 0].numpy(), (width, height)), (3, 3)), + cv2.blur(cv2.resize(bm[0, 1].numpy(), (width, height)), (3, 3)) + ], axis=2)).unsqueeze(0) + + out = F.grid_sample( + torch.from_numpy(image).permute(2, 0, 1).unsqueeze(0).float(), + flow, align_corners=True + ) + result = ((out[0] * 255).permute(1, 2, 0).numpy()[:, :, ::-1]).astype(np.uint8) + + out_dir = os.path.dirname(output_path) + if out_dir: + os.makedirs(out_dir, exist_ok=True) + cv2.imwrite(output_path, result) + print(f"[OK] {output_path}") + + +if __name__ == "__main__": + main()