Spaces:
Sleeping
Sleeping
File size: 2,534 Bytes
5c11384 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 |
from utils.all_utils import read_yaml, create_dir
import os
import pickle
from mtcnn import MTCNN
from tensorflow.keras.preprocessing import image
from tqdm import tqdm
import cv2
def extract_face_from_image(image_path, required_size=(224, 224)):
img = cv2.imread(image_path)
detector = MTCNN()
faces = detector.detect_faces(img)
if len(faces)>0:
x, y, width, height = faces[0]['box']
face_boundary = img[y:y+ height, x:x+width]
image = cv2.resize(face_boundary, required_size)
return image
def generate_image_pickle_file(config_path, params_path):
config = read_yaml(config_path)
params = read_yaml(params_path)
artifacts = config['artifacts']
artifacts_dir = artifacts['artifacts_dir']
pickle_format_dir = artifacts['pickle_format_data_dir']
img_pickle_filename = artifacts['img_pickle_file_name']
pickle_actors_name = artifacts['pickle_actor_names']
cropped_dir = artifacts['cropped_dir']
raw_local_dir_path = os.path.join(artifacts_dir,pickle_format_dir)
create_dir([raw_local_dir_path])
pickle_file = os.path.join(raw_local_dir_path, img_pickle_filename)
pickle_actor = os.path.join(raw_local_dir_path,pickle_actors_name)
data = params['base']['data_path']
create_dir([os.path.join(data,cropped_dir)])
actors = os.listdir(data)
filenames = []
for actor in tqdm(actors):
count = 0
actor_crop_dir = os.path.join(data, cropped_dir, actor)
create_dir([actor_crop_dir])
for file in os.listdir(os.path.join(data, actor)):
file_dir = os.path.join(data,actor, file)
try:
detected_face = extract_face_from_image(file_dir)
cropped_file_name = actor+ "_" + str(count) + ".jpg"
cv2.imwrite(os.path.join(actor_crop_dir,cropped_file_name), detected_face)
count+=1
except:
pass
for file in os.listdir(actor_crop_dir):
filenames.append(os.path.join(data,cropped_dir,actor,file))
print(f'Total celeb are: {len(actors)}')
print(f'Total celeb images: {len(filenames)}')
pickle.dump(filenames, open(pickle_file, 'wb'))
pickle.dump(actors, open(pickle_actor, 'wb'))
if __name__ == '__main__':
generate_image_pickle_file('config/config.yaml', 'params.yaml')
|