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)