Last active
July 15, 2019 10:54
-
-
Save lsimmons2/6138d3b18ce990fcf172701bde78c762 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import os | |
import cv2 | |
import dlib | |
from matplotlib import pyplot as plt | |
import numpy as np | |
import config | |
detector = dlib.get_frontal_face_detector() | |
def crop_faces(): | |
bad_crop_count = 0 | |
if not os.path.exists(config.CROPPED_IMGS_DIR): | |
os.makedirs(config.CROPPED_IMGS_DIR) | |
print 'Cropping faces and saving to %s' % config.CROPPED_IMGS_DIR | |
good_cropped_images = [] | |
good_cropped_img_file_names = [] | |
detected_cropped_images = [] | |
original_images_detected = [] | |
for file_name in sorted(os.listdir(config.ORIGINAL_IMGS_DIR)): | |
np_img = cv2.imread(os.path.join(config.ORIGINAL_IMGS_DIR,file_name)) | |
detected = detector(np_img, 1) | |
img_h, img_w, _ = np.shape(np_img) | |
original_images_detected.append(np_img) | |
if len(detected) != 1: | |
bad_crop_count += 1 | |
continue | |
d = detected[0] | |
x1, y1, x2, y2, w, h = d.left(), d.top(), d.right() + 1, d.bottom() + 1, d.width(), d.height() | |
xw1 = int(x1 - config.MARGIN * w) | |
yw1 = int(y1 - config.MARGIN * h) | |
xw2 = int(x2 + config.MARGIN * w) | |
yw2 = int(y2 + config.MARGIN * h) | |
cropped_img = crop_image_to_dimensions(np_img, xw1, yw1, xw2, yw2) | |
norm_file_path = '%s/%s' % (config.CROPPED_IMGS_DIR, file_name) | |
cv2.imwrite(norm_file_path, cropped_img) | |
good_cropped_img_file_names.append(file_name) | |
# save info of good cropped images | |
with open(config.ORIGINAL_IMGS_INFO_FILE, 'r') as f: | |
column_headers = f.read().splitlines()[0] | |
all_imgs_info = f.read().splitlines()[1:] | |
cropped_imgs_info = [l for l in all_imgs_info if l.split(',')[-1] in good_cropped_img_file_names] | |
with open(config.CROPPED_IMGS_INFO_FILE, 'w') as f: | |
f.write('%s\n' % column_headers) | |
for l in cropped_imgs_info: | |
f.write('%s\n' % l) | |
print 'Cropped %d images and saved in %s - info in %s' % (len(original_images_detected), config.CROPPED_IMGS_DIR, config.CROPPED_IMGS_INFO_FILE) | |
print 'Error detecting face in %d images - info in Data/unnormalized.txt' % bad_crop_count | |
return good_cropped_images | |
# image cropping function taken from: | |
# https://stackoverflow.com/questions/15589517/how-to-crop-an-image-in-opencv-using-python | |
def crop_image_to_dimensions(img, x1, y1, x2, y2): | |
if x1 < 0 or y1 < 0 or x2 > img.shape[1] or y2 > img.shape[0]: | |
img, x1, x2, y1, y2 = pad_img_to_fit_bbox(img, x1, x2, y1, y2) | |
return img[y1:y2, x1:x2, :] | |
def pad_img_to_fit_bbox(img, x1, x2, y1, y2): | |
img = cv2.copyMakeBorder(img, - min(0, y1), max(y2 - img.shape[0], 0), | |
-min(0, x1), max(x2 - img.shape[1], 0), cv2.BORDER_REPLICATE) | |
y2 += -min(0, y1) | |
y1 += -min(0, y1) | |
x2 += -min(0, x1) | |
x1 += -min(0, x1) | |
return img, x1, x2, y1, y2 | |
if __name__ == '__main__': | |
crop_faces() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment