Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fixing RGB2BGR problem with opencv and adding live feed through webcam option #96

Open
wants to merge 3 commits into
base: master
Choose a base branch
from
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
225 changes: 142 additions & 83 deletions demo.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@
from utils.write import write_obj_with_colors, write_obj_with_texture

def main(args):
if args.isShow or args.isTexture:
if args.isShow or args.isTexture or args.isCamera:
import cv2
from utils.cv_plot import plot_kpt, plot_vertices, plot_pose_box

Expand All @@ -36,113 +36,169 @@ def main(args):
image_path_list.extend(glob(os.path.join(image_folder, files)))
total_num = len(image_path_list)

for i, image_path in enumerate(image_path_list):

name = image_path.strip().split('/')[-1][:-4]

# read image
image = imread(image_path)
[h, w, c] = image.shape
if c>3:
image = image[:,:,:3]

# the core: regress position map
if args.isDlib:
max_size = max(image.shape[0], image.shape[1])
if max_size> 1000:
image = rescale(image, 1000./max_size)
image = (image*255).astype(np.uint8)
pos = prn.process(image) # use dlib to detect face
else:
if image.shape[0] == image.shape[1]:
image = resize(image, (256,256))
pos = prn.net_forward(image/255.) # input image has been cropped to 256x256
else:
box = np.array([0, image.shape[1]-1, 0, image.shape[0]-1]) # cropped with bounding box
pos = prn.process(image, box)
if args.isCamera:

# Create a VideoCapture object and read from input file
# If the input is the camera, pass 0 instead of the video file name
cap = cv2.VideoCapture(0)

# Check if camera opened successfully
if (cap.isOpened()== False):
print("Error opening video stream or file")

image = image/255.
if pos is None:
continue

if args.is3d or args.isMat or args.isPose or args.isShow:
# 3D vertices
vertices = prn.get_vertices(pos)
if args.isFront:
save_vertices = frontalize(vertices)
# Read until video is completed
while(cap.isOpened()):
# Capture frame-by-frame
ret, frame = cap.read()
if ret == True:

if args.isDlib:
max_size = max(frame.shape[0], frame.shape[1])
if max_size> 1000:
frame = rescale(frame, 1000./max_size)
frame = (frame*255).astype(np.uint8)
pos = prn.process(frame) # use dlib to detect face
else:
if frame.shape[0] == frame.shape[1]:
frame = resize(frame, (256,256))
pos = prn.net_forward(frame/255.) # input frame has been cropped to 256x256
else:
box = np.array([0, frame.shape[1]-1, 0, frame.shape[0]-1]) # cropped with bounding box
pos = prn.process(frame, box)
# Normalizing the frame and skiping if there was no one in the frame
frame = frame/255.
if pos is None:
continue
# Get landmarks in frame
kpt = prn.get_landmarks(pos)

# Display the resulting frame
cv2.imshow('sparse alignment', plot_kpt(frame, kpt))

# Press Q on keyboard to exit
if cv2.waitKey(25) & 0xFF == ord('q'):
break

# Break the loop
else:
break

# When everything done, release the video capture object
cap.release()

# Closes all the frames
cv2.destroyAllWindows()

else:
for i, image_path in enumerate(image_path_list):

name = image_path.strip().split('/')[-1][:-4]

# read image
image = imread(image_path)
[h, w, c] = image.shape
if c>3:
image = image[:,:,:3]

# the core: regress position map
if args.isDlib:
max_size = max(image.shape[0], image.shape[1])
if max_size> 1000:
image = rescale(image, 1000./max_size)
image = (image*255).astype(np.uint8)
pos = prn.process(image) # use dlib to detect face
else:
save_vertices = vertices.copy()
save_vertices[:,1] = h - 1 - save_vertices[:,1]

if args.isImage:
imsave(os.path.join(save_folder, name + '.jpg'), image)
if image.shape[0] == image.shape[1]:
image = resize(image, (256,256))
pos = prn.net_forward(image/255.) # input image has been cropped to 256x256
else:
box = np.array([0, image.shape[1]-1, 0, image.shape[0]-1]) # cropped with bounding box
pos = prn.process(image, box)

image = image/255.
if pos is None:
continue

if args.is3d or args.isMat or args.isPose or args.isShow:
# 3D vertices
vertices = prn.get_vertices(pos)
if args.isFront:
save_vertices = frontalize(vertices)
else:
save_vertices = vertices.copy()
save_vertices[:,1] = h - 1 - save_vertices[:,1]

if args.isImage:
imsave(os.path.join(save_folder, name + '.jpg'), image)

if args.is3d:
# corresponding colors
colors = prn.get_colors(image, vertices)

if args.isTexture:
if args.texture_size != 256:
pos_interpolated = resize(pos, (args.texture_size, args.texture_size), preserve_range = True)
else:
pos_interpolated = pos.copy()
texture = cv2.remap(image, pos_interpolated[:,:,:2].astype(np.float32), None, interpolation=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT,borderValue=(0))
if args.isMask:
vertices_vis = get_visibility(vertices, prn.triangles, h, w)
uv_mask = get_uv_mask(vertices_vis, prn.triangles, prn.uv_coords, h, w, prn.resolution_op)
uv_mask = resize(uv_mask, (args.texture_size, args.texture_size), preserve_range = True)
texture = texture*uv_mask[:,:,np.newaxis]
write_obj_with_texture(os.path.join(save_folder, name + '.obj'), save_vertices, prn.triangles, texture, prn.uv_coords/prn.resolution_op)#save 3d face with texture(can open with meshlab)
else:
write_obj_with_colors(os.path.join(save_folder, name + '.obj'), save_vertices, prn.triangles, colors) #save 3d face(can open with meshlab)

if args.is3d:
# corresponding colors
colors = prn.get_colors(image, vertices)
if args.isDepth:
depth_image = get_depth_image(vertices, prn.triangles, h, w, True)
depth = get_depth_image(vertices, prn.triangles, h, w)
imsave(os.path.join(save_folder, name + '_depth.jpg'), depth_image)
sio.savemat(os.path.join(save_folder, name + '_depth.mat'), {'depth':depth})

if args.isTexture:
if args.texture_size != 256:
pos_interpolated = resize(pos, (args.texture_size, args.texture_size), preserve_range = True)
else:
pos_interpolated = pos.copy()
texture = cv2.remap(image, pos_interpolated[:,:,:2].astype(np.float32), None, interpolation=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT,borderValue=(0))
if args.isMask:
vertices_vis = get_visibility(vertices, prn.triangles, h, w)
uv_mask = get_uv_mask(vertices_vis, prn.triangles, prn.uv_coords, h, w, prn.resolution_op)
uv_mask = resize(uv_mask, (args.texture_size, args.texture_size), preserve_range = True)
texture = texture*uv_mask[:,:,np.newaxis]
write_obj_with_texture(os.path.join(save_folder, name + '.obj'), save_vertices, prn.triangles, texture, prn.uv_coords/prn.resolution_op)#save 3d face with texture(can open with meshlab)
else:
write_obj_with_colors(os.path.join(save_folder, name + '.obj'), save_vertices, prn.triangles, colors) #save 3d face(can open with meshlab)
if args.isMat:
sio.savemat(os.path.join(save_folder, name + '_mesh.mat'), {'vertices': vertices, 'colors': colors, 'triangles': prn.triangles})

if args.isDepth:
depth_image = get_depth_image(vertices, prn.triangles, h, w, True)
depth = get_depth_image(vertices, prn.triangles, h, w)
imsave(os.path.join(save_folder, name + '_depth.jpg'), depth_image)
sio.savemat(os.path.join(save_folder, name + '_depth.mat'), {'depth':depth})
if args.isKpt or args.isShow:
# get landmarks
kpt = prn.get_landmarks(pos)
np.savetxt(os.path.join(save_folder, name + '_kpt.txt'), kpt)

if args.isMat:
sio.savemat(os.path.join(save_folder, name + '_mesh.mat'), {'vertices': vertices, 'colors': colors, 'triangles': prn.triangles})
if args.isPose or args.isShow:
# estimate pose
camera_matrix, pose = estimate_pose(vertices)
np.savetxt(os.path.join(save_folder, name + '_pose.txt'), pose)
np.savetxt(os.path.join(save_folder, name + '_camera_matrix.txt'), camera_matrix)

if args.isKpt or args.isShow:
# get landmarks
kpt = prn.get_landmarks(pos)
np.savetxt(os.path.join(save_folder, name + '_kpt.txt'), kpt)
np.savetxt(os.path.join(save_folder, name + '_pose.txt'), pose)

if args.isPose or args.isShow:
# estimate pose
camera_matrix, pose = estimate_pose(vertices)
np.savetxt(os.path.join(save_folder, name + '_pose.txt'), pose)
np.savetxt(os.path.join(save_folder, name + '_camera_matrix.txt'), camera_matrix)
if args.isShow:
# ---------- Plot
image_pose = plot_pose_box(image, camera_matrix, kpt)
cv2.imshow('sparse alignment', cv2.cvtColor(np.float32(plot_kpt(image, kpt)), cv2.COLOR_RGB2BGR))
cv2.imshow('dense alignment', cv2.cvtColor(np.float32(plot_vertices(image, vertices)), cv2.COLOR_RGB2BGR))
cv2.imshow('pose', cv2.cvtColor(np.float32(plot_pose_box(image, camera_matrix, kpt)), cv2.COLOR_RGB2BGR))
cv2.waitKey(0)

np.savetxt(os.path.join(save_folder, name + '_pose.txt'), pose)

if args.isShow:
# ---------- Plot
image_pose = plot_pose_box(image, camera_matrix, kpt)
cv2.imshow('sparse alignment', plot_kpt(image, kpt))
cv2.imshow('dense alignment', plot_vertices(image, vertices))
cv2.imshow('pose', plot_pose_box(image, camera_matrix, kpt))
cv2.waitKey(0)


if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Joint 3D Face Reconstruction and Dense Alignment with Position Map Regression Network')

parser.add_argument('-i', '--inputDir', default='TestImages/', type=str,
help='path to the input directory, where input images are stored.')
help='path to the input directory, where input images/videos are stored.')
parser.add_argument('-o', '--outputDir', default='TestImages/results', type=str,
help='path to the output directory, where results(obj,txt files) will be stored.')
parser.add_argument('--gpu', default='0', type=str,
help='set gpu id, -1 for CPU')
parser.add_argument('--isDlib', default=True, type=ast.literal_eval,
help='whether to use dlib for detecting face, default is True, if False, the input image should be cropped in advance')
parser.add_argument('--is3d', default=True, type=ast.literal_eval,
parser.add_argument('--is3d', default=False, type=ast.literal_eval,
help='whether to output 3D face(.obj). default save colors.')
parser.add_argument('--isMat', default=False, type=ast.literal_eval,
help='whether to save vertices,color,triangles as mat for matlab showing')
parser.add_argument('--isKpt', default=False, type=ast.literal_eval,
parser.add_argument('--isKpt', default=True, type=ast.literal_eval,
help='whether to output key points(.txt)')
parser.add_argument('--isPose', default=False, type=ast.literal_eval,
help='whether to output estimated pose(.txt)')
Expand All @@ -164,4 +220,7 @@ def main(args):
# update in 2017/7/19
parser.add_argument('--texture_size', default=256, type=int,
help='size of texture map, default is 256. need isTexture is True')
# update in 2019/01/03
parser.add_argument('--isCamera', default=False, type=ast.literal_eval,
help='whether to show live 3d face reconstruction (need opencv)')
main(parser.parse_args())