|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
import logging |
|
|
import os.path |
|
|
|
|
|
import PIL.Image |
|
|
import cv2 |
|
|
import numpy as np |
|
|
import torch |
|
|
from torchvision.transforms import transforms |
|
|
|
|
|
from scripts.erasescratches.models import Pix2PixHDModel_Mapping |
|
|
from scripts.erasescratches.options import Options |
|
|
from scripts.maskscratches import ScratchesDetector |
|
|
from scripts.util import irregular_hole_synthesize, tensor_to_ndarray |
|
|
|
|
|
REPO_ID = "leonelhs/zeroscratches" |
|
|
|
|
|
|
|
|
class EraseScratches: |
|
|
|
|
|
def __init__(self): |
|
|
|
|
|
snapshot_folder = './models/zeroscratches' |
|
|
model_path = os.path.join(snapshot_folder, "restoration") |
|
|
self.detector = ScratchesDetector(snapshot_folder) |
|
|
gpu_ids = [] |
|
|
if torch.cuda.is_available(): |
|
|
gpu_ids = [d for d in range(torch.cuda.device_count())] |
|
|
self.options = Options(model_path, gpu_ids) |
|
|
self.model_scratches = Pix2PixHDModel_Mapping() |
|
|
self.model_scratches.initialize(self.options) |
|
|
self.model_scratches.eval() |
|
|
|
|
|
def erase(self, image) -> np.array: |
|
|
transformed, mask = self.detector.process(image) |
|
|
logging.info("Start erasing scratches") |
|
|
|
|
|
img_transform = transforms.Compose( |
|
|
[transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))] |
|
|
) |
|
|
mask_transform = transforms.ToTensor() |
|
|
|
|
|
if self.options.mask_dilation != 0: |
|
|
kernel = np.ones((3, 3), np.uint8) |
|
|
mask = np.array(mask) |
|
|
mask = cv2.dilate(mask, kernel, iterations=self.options.mask_dilation) |
|
|
mask = PIL.Image.fromarray(mask.astype('uint8')) |
|
|
|
|
|
transformed = irregular_hole_synthesize(transformed, mask) |
|
|
mask = mask_transform(mask) |
|
|
mask = mask[:1, :, :] |
|
|
mask = mask.unsqueeze(0) |
|
|
transformed = img_transform(transformed) |
|
|
transformed = transformed.unsqueeze(0) |
|
|
|
|
|
try: |
|
|
with torch.no_grad(): |
|
|
generated = self.model_scratches.inference(transformed, mask) |
|
|
except Exception as ex: |
|
|
raise TypeError("Skip photo due to an error:\n%s" % str(ex)) |
|
|
|
|
|
tensor_restored = (generated.data.cpu() + 1.0) / 2.0 |
|
|
return tensor_to_ndarray(tensor_restored) |
|
|
|
