class Photo2Cartoon:
def __init__(self):
self.pre = Preprocess()
self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
self.net = ResnetGenerator(ngf=32, img_size=256, light=True).to(self.device)
params = torch.load('./models/photo2cartoon_10000.pt', map_location=self.device)
self.net.load_state_dict(params['genA2B'])
def inference(self, img):
# face alignment and segmentation
face_rgba = self.pre.process(img)
if face_rgba is None:
print('can not detect face!!!')
return None
face_rgba = cv2.resize(face_rgba, (256, 256), interpolation=cv2.INTER_AREA)
face = face_rgba[:, :, :3].copy()
mask = face_rgba[:, :, 3][:, :, np.newaxis].copy() / 255.
face = (face * mask + (1 - mask) * 255) / 127.5 - 1
face = np.transpose(face[np.newaxis, :, :, :], (0, 3, 1, 2)).astype(np.float32)
face = torch.from_numpy(face).to(self.device)
# inference
with torch.no_grad():
cartoon = self.net(face)[0][0]
# post-process
cartoon = np.transpose(cartoon.cpu().numpy(), (1, 2, 0))
cartoon = (cartoon + 1) * 127.5
cartoon = (cartoon * mask + 255 * (1 - mask)).astype(np.uint8)
cartoon = cv2.cvtColor(cartoon, cv2.COLOR_RGB2BGR)
return cartoon
class Photo2Cartoon:
def __init__(self):
self.pre = Preprocess()
self.device = torch.device(
"cuda:0" if torch.cuda.is_available() else "cpu")
self.net = ResnetGenerator(ngf=32, img_size=256,
light=True).to(self.device)
assert os.path.exists(
'./models/photo2blackjack_weights.pt'
), "[Step1: load weights] Can not find 'photo2blackjack_weights.pt' in folder 'models!!!'"
params = torch.load('./models/photo2blackjack_weights.pt',
map_location=self.device)
self.net.load_state_dict(params['genA2B'])
print('[Step1: load weights] success!')
def inference(self, img):
# face alignment and segmentation
face_rgba = self.pre.process(img)
if face_rgba is None:
print('[Step2: face detect] can not detect face!!!')
return None
print('[Step2: face detect] success!')
face_rgba = cv2.resize(face_rgba, (256, 256),
interpolation=cv2.INTER_AREA)
face = face_rgba[:, :, :3].copy()
mask = face_rgba[:, :, 3][:, :, np.newaxis].copy() / 255.
face = (face * mask + (1 - mask) * 255)
img = Image.fromarray(np.uint8(face))
img = img.convert("L")
img = img.convert("RGB")
face = np.asarray(img)
face = np.transpose(face[np.newaxis, :, :, :],
(0, 3, 1, 2)).astype(np.float32)
face = torch.from_numpy(face).to(self.device)
# inference
with torch.no_grad():
cartoon = self.net(face)[0][0]
# post-process
cartoon = np.transpose(cartoon.cpu().numpy(), (1, 2, 0))
cartoon = (cartoon + 1) * 127.5
cartoon = (cartoon * mask + 255 * (1 - mask)).astype(np.uint8)
cartoon = cv2.cvtColor(cartoon, cv2.COLOR_RGB2BGR)
print('[Step3: face detect] success!')
return cartoon
def gen_cartoon(img):
pre = Preprocess()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
net = ResnetGenerator(ngf=32, img_size=256, light=True).to(device)
assert os.path.exists(
'./models/photo2cartoon_weights.pt'
), "[Step1: load weights] Can not find 'photo2cartoon_weights.pt' in folder 'models!!!'"
params = torch.load('./models/photo2cartoon_weights.pt',
map_location=device)
net.load_state_dict(params['genA2B'])
# face alignment and segmentation
face_rgba = pre.process(img)
if face_rgba is None:
return None
face_rgba = cv2.resize(face_rgba, (256, 256), interpolation=cv2.INTER_AREA)
face = face_rgba[:, :, :3].copy()
mask = face_rgba[:, :, 3][:, :, np.newaxis].copy() / 255.
face = (face * mask + (1 - mask) * 255) / 127.5 - 1
face = np.transpose(face[np.newaxis, :, :, :],
(0, 3, 1, 2)).astype(np.float32)
face = torch.from_numpy(face).to(device)
# inference
with torch.no_grad():
cartoon = net(face)[0][0]
# post-process
cartoon = np.transpose(cartoon.cpu().numpy(), (1, 2, 0))
cartoon = (cartoon + 1) * 127.5
cartoon = (cartoon * mask + 255 * (1 - mask)).astype(np.uint8)
cartoon = cv2.cvtColor(cartoon, cv2.COLOR_RGB2BGR)
out_path = Path(tempfile.mkdtemp()) / "out.png"
cv2.imwrite(str(out_path), cartoon)
return out_path
class Photo2Cartoon:
def __init__(self):
self.pre = Preprocess()
self.device = torch.device(
"cuda:0" if torch.cuda.is_available() else "cpu")
self.net = ResnetGenerator(ngf=32, img_size=256,
light=True).to(self.device)
assert os.path.exists(
'./models/photo2cartoon_weights.pt'
), "[Step1: load weights] Can not find 'photo2cartoon_weights.pt' in folder 'models!!!'"
params = torch.load('./models/photo2cartoon_weights.pt',
map_location=self.device)
self.net.load_state_dict(params['genA2B'])
print('[Step1: load weights] success!')
def inference(self, img):
# face alignment and segmentation
start = time.time()
face_rgba = self.pre.process_optimize(img)
# cv2.imwrite("/home/onepiece/GZX/photo2cartoon/photo2cartoon-master/images/face_rgba.png", face_rgba)
end = time.time()
print("preprocess time spent: ", end - start)
if face_rgba is None:
print('[Step2: face detect] can not detect face!!!')
return None, None, None
print('[Step2: face detect] success!')
size = [face_rgba.shape[0], face_rgba.shape[1]]
face_rgba = cv2.resize(face_rgba, (256, 256),
interpolation=cv2.INTER_AREA)
face = face_rgba[:, :, :3].copy()
mask = face_rgba[:, :, 3][:, :, np.newaxis].copy() / 255.
# mask = np.ones((256, 256, 1)) # cancel front background segmentation module
face = (face * mask + (1 - mask) * 255) / 127.5 - 1
face = np.transpose(face[np.newaxis, :, :, :],
(0, 3, 1, 2)).astype(np.float32)
face = torch.from_numpy(face).to(self.device)
# inference
with torch.no_grad():
cartoon = self.net(face)[0][0]
# post-process
cartoon = np.transpose(cartoon.cpu().numpy(), (1, 2, 0))
cartoon = (cartoon + 1) * 127.5
cartoon = (cartoon * mask + 255 * (1 - mask)).astype(np.uint8)
cartoon = cv2.cvtColor(cartoon, cv2.COLOR_RGB2BGR)
print('[Step3: photo to cartoon] success!')
# cv2.imwrite("/home/onepiece/GZX/photo2cartoon/photo2cartoon-master/images/cartoon_before.png", cartoon)
# cv2.imwrite("/home/onepiece/GZX/photo2cartoon/photo2cartoon-master/images/mask_before.png", mask * 255)
# resize back to origin size
cartoon = cv2.resize(cartoon, (size[0], size[1]),
interpolation=cv2.INTER_AREA)
mask = cv2.resize(mask * 1.1, (size[0], size[1]),
interpolation=cv2.INTER_AREA)[:, :,
np.newaxis].astype(
np.uint8)
return cartoon, mask, self.pre.rect
parser.add_argument('--path',
default="inference",
type=str,
metavar='DIR',
help='path to get images')
device = torch.device(
'cuda') if torch.cuda.is_available() else torch.device('cpu')
args = parser.parse_args()
file_names = sorted(os.listdir(args.path))
mymodel = ResnetGenerator()
mymodel.to(device)
os.makedirs(os.path.join("result"), exist_ok=True)
mymodel.load_state_dict(
torch.load(os.path.join("model_weight", 'best_weight.pt'),
map_location=device)['G_state_dict'])
mymodel.eval()
for i in range(len(file_names)):
video = read_video(os.path.join(args.path, file_names[i]),
inference=True).to(device)
with torch.no_grad():
reconstructed = []
for j in range(video.size(0)):
reconstructed.append(mymodel(video[j][None]).cpu().numpy())
reconstructed = np.concatenate(reconstructed)
save_video(reconstructed,
os.path.join("result", file_names[i], '_filled.avi'))
# read source image
img_source = Image.open("assets/3_color_net_neg23570_ori.png")
# read target image
img_target = Image.open("assets/3_color_net_neg23570_target.png")
# load pretrained StainNet
model_Net = StainNet().cuda()
model_Net.load_state_dict(
torch.load("checkpoints/StainNet/StainNet-3x0_best_psnr_layer3_ch32.pth"))
# load pretrained StainGAN
model_GAN = ResnetGenerator(3,
3,
ngf=64,
norm_layer=nn.InstanceNorm2d,
n_blocks=9).cuda()
model_GAN.load_state_dict(
torch.load("checkpoints/StainGAN/latest_net_G_A.pth"))
def test_deeplearning_fps(model, n_iters, batchsize):
data = torch.rand(batchsize, 3, 512, 512).cuda()
start_time = time.time()
for i in tqdm(range(n_iters)):
with torch.no_grad():
outputs = model(data)
process_time = time.time() - start_time
print("FPS is ", n_iters * batchsize / process_time)
def test_traditional_fps(source_img, ref_img, method, n_iters):
ref_img = np.array(ref_img)
source_img = np.array(ref_img)