Closed Ravindu-Yasas-Nagasinghe closed 2 years ago
Thank you for your quick response. Here I get the output '.mat' file as a 4-D Uint8 output. Is there a way to get a 2-D int32 output ".mat" file. given below is the code for my output.
` ''' This script enhance all faces in one image with PSFR-GAN and paste it back to the original place. ''' import dlib import os import cv2 import numpy as np from tqdm import tqdm from skimage import transform as trans from skimage import io
import torch from utils import utils from options.test_options import TestOptions from models import create_model from scipy.io import savemat
def detect_and_align_faces(img, face_detector, lmk_predictor, template_path, template_scale=2, size_threshold=999): align_out_size = (512, 512) ref_points = np.load(template_path) / template_scale
# Detect landmark points
face_dets = face_detector(img, 1)
assert len(face_dets) > 0, 'No faces detected'
aligned_faces = []
tform_params = []
for det in face_dets:
if isinstance(face_detector, dlib.cnn_face_detection_model_v1):
rec = det.rect # for cnn detector
else:
rec = det
if rec.width() > size_threshold or rec.height() > size_threshold:
print('Face is too large')
break
landmark_points = lmk_predictor(img, rec)
single_points = []
for i in range(5):
single_points.append([landmark_points.part(i).x, landmark_points.part(i).y])
single_points = np.array(single_points)
tform = trans.SimilarityTransform()
tform.estimate(single_points, ref_points)
tmp_face = trans.warp(img, tform.inverse, output_shape=align_out_size, order=3)
aligned_faces.append(tmp_face*255)
tform_params.append(tform)
return [aligned_faces, tform_params]
def def_models(opt): model = create_model(opt) model.load_pretrain_models() model.netP.to(opt.device) model.netG.to(opt.device) return model
def enhance_faces(LQ_faces, model): hq_faces = [] lq_parse_maps = [] for lq_face in tqdm(LQ_faces): with torch.no_grad(): lq_tensor = torch.tensor(lq_face.transpose(2, 0, 1)) / 255. * 2 - 1 lq_tensor = lq_tensor.unsqueeze(0).float().to(model.device) parsemap, = model.netP(lq_tensor) parse_map_onehot = (parse_map == parse_map.max(dim=1, keepdim=True)[0]).float() output_SR = model.netG(lq_tensor, parse_map_onehot) hq_faces.append(utils.tensor_to_img(output_SR)) lq_parse_maps.append(utils.color_parse_map(parse_map_onehot)[0]) return hq_faces, lq_parse_maps
def past_faces_back(img, hq_faces, tform_params, upscale=1): h, w = img.shape[:2] img = cv2.resize(img, (int(wupscale), int(hupscale)), interpolation=cv2.INTER_CUBIC) for hq_img, tform in tqdm(zip(hq_faces, tform_params), total=len(hq_faces)): tform.params[0:2,0:2] /= upscale back_img = trans.warp(hq_img/255., tform, output_shape=[int(hupscale), int(wupscale)], order=3) * 255
# blur mask to avoid border artifacts
mask = (back_img == 0)
mask = cv2.blur(mask.astype(np.float32), (5,5))
mask = (mask > 0)
img = img * mask + (1 - mask) * back_img
return img.astype(np.uint8)
def save_imgs(img_list, save_dir): for idx, img in enumerate(img_list): save_path = os.path.join(save_dir, '{:03d}.jpg'.format(idx)) io.imsave(save_path, img.astype(np.uint8))
if name == 'main': opt = TestOptions().parse()
face_detector = dlib.cnn_face_detection_model_v1('./pretrain_models/mmod_human_face_detector.dat')
lmk_predictor = dlib.shape_predictor('./pretrain_models/shape_predictor_5_face_landmarks.dat')
template_path = './pretrain_models/FFHQ_template.npy'
print('======> Loading images, crop and align faces.')
img_path = opt.test_img_path
img = dlib.load_rgb_image('/home/zebra/Deep3dPortrait/examples/1_1234.jpg')
aligned_faces, tform_params = detect_and_align_faces(img, face_detector, lmk_predictor, template_path)
# Save aligned LQ faces
save_lq_dir = os.path.join(opt.results_dir, 'LQ_faces')
os.makedirs(save_lq_dir, exist_ok=True)
print('======> Saving aligned LQ faces to', save_lq_dir)
save_imgs(aligned_faces, save_lq_dir)
enhance_model = def_models(opt)
hq_faces, lq_parse_maps = enhance_faces(aligned_faces, enhance_model)
# Save LQ parsing maps and enhanced faces
save_parse_dir = os.path.join(opt.results_dir, 'ParseMaps')
save_hq_dir = os.path.join(opt.results_dir, 'HQ')
os.makedirs(save_parse_dir, exist_ok=True)
os.makedirs(save_hq_dir, exist_ok=True)
print('======> Save parsing map and the enhanced faces.')
arr = np.array(lq_parse_maps)
print(arr)
mdic = {"lq_parse_maps": lq_parse_maps}
savemat("custom_input.mat", mdic)
print('######################################')
print(mdic)
save_imgs(lq_parse_maps, save_parse_dir)
save_imgs(hq_faces, save_hq_dir)
print('======> Paste the enhanced faces back to the original image.')
hq_img = past_faces_back(img, hq_faces, tform_params, upscale=opt.test_upscale)
final_save_path = os.path.join(opt.results_dir, 'hq_final.jpg')
print('======> Save final result to', final_save_path)
io.imsave(final_save_path, hq_img)
`
You may try tools like
scipy.io.savemat
to save the numpy array to.mat
files using python script.
new code part added to create .mat file are in the last 8-13 lines
You may try tools like
scipy.io.savemat
to save the numpy array to.mat
files using python script.