Spaces:
Sleeping
Sleeping
| 8import os | |
| import face_recognition | |
| import pickle | |
| from tensorflow.keras.preprocessing.image import ImageDataGenerator | |
| import cv2 | |
| def generateDatasetImages(): | |
| # Define ImageDataGenerator with desired augmentation parameters | |
| datagen = ImageDataGenerator( | |
| rotation_range=20, # Rotate images randomly within the range of -20 to +20 degrees | |
| width_shift_range=0.1, # Shift images horizontally (as a fraction of total width) | |
| height_shift_range=0.1, # Shift images vertically (as a fraction of total height) | |
| shear_range=0.1, # Shear intensity (shear angle in counter-clockwise direction in radians) | |
| zoom_range=0.1, # Zoom images randomly by a factor of 0.1 | |
| horizontal_flip=True, # Flip images horizontally | |
| vertical_flip=True, # Flip images vertically | |
| ) | |
| # Directory containing example images | |
| input_dir = 'static/data/dataset/' | |
| # Iterate over each image file in the directory | |
| for filename in os.listdir(input_dir): | |
| if filename.endswith(".jpg") or filename.endswith(".jpeg") or filename.endswith(".png"): | |
| output_dir = os.path.join(input_dir, "augmented_images", os.path.splitext(filename)[0]) | |
| if not os.path.exists(output_dir): | |
| # Load the image using OpenCV (required for face_recognition) | |
| img_path = os.path.join(input_dir, filename) | |
| img = cv2.imread(img_path) | |
| rgb_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) # Convert BGR to RGB | |
| # Detect faces in the image using face_recognition | |
| face_locations = face_recognition.face_locations(rgb_img) | |
| if len(face_locations) == 0: | |
| print(f"No faces detected in '{filename}'. Skipping augmentation.") | |
| continue | |
| # Create a directory for augmented images based on the image filename | |
| os.makedirs(output_dir, exist_ok=True) | |
| # Convert pixel values to [0, 1] | |
| rgb_img = rgb_img / 255.0 | |
| # Generate augmented images | |
| img = rgb_img.reshape((1,) + rgb_img.shape) | |
| # Generate augmented images | |
| num_samples = 100 # Number of augmented samples to generate | |
| augmented_images = [] | |
| # Generate augmented images | |
| for batch in datagen.flow(img, batch_size=1): | |
| augmented_image = (batch[0] * 255).astype('uint8') # Convert back to uint8 | |
| augmented_images.append(augmented_image) | |
| num_samples -= 1 | |
| if num_samples == 0: | |
| break | |
| # Save augmented images | |
| for i, augmented_image in enumerate(augmented_images): | |
| aug_filename = f'{os.path.splitext(filename)[0]}_{i+1}.jpg' | |
| aug_filepath = os.path.join(output_dir, aug_filename) | |
| cv2.imwrite(aug_filepath, cv2.cvtColor(augmented_image, cv2.COLOR_RGB2BGR)) | |
| print(f"{len(augmented_images)} augmented images saved in '{output_dir}' directory.") | |
| known_face_encodings = [] | |
| known_face_names = [] | |
| dataset_path = "static/data/dataset" | |
| # Load images and augment dataset | |
| image_files = [file for file in os.listdir(dataset_path) if file.endswith(('.jpg', '.jpeg', '.png'))] | |
| for file in image_files: | |
| name = os.path.splitext(os.path.basename(file))[0] | |
| images = os.listdir(os.path.join(dataset_path, "augmented_images", name)) | |
| file_path = face_recognition.load_image_file(os.path.join(dataset_path, file)) | |
| face_locations = face_recognition.face_locations(file_path, model="hog") | |
| encoding = face_recognition.face_encodings(file_path, face_locations) | |
| known_face_encodings.append(encoding) | |
| known_face_names.append(name) | |
| # for image in images: | |
| # subfile = face_recognition.load_image_file(os.path.join(dataset_path, "augmented_images", name, image)) | |
| # face_locations = face_recognition.face_locations(subfile, model="hog") | |
| # face_encodings = face_recognition.face_encodings(subfile, face_locations) | |
| # for encoding in face_encodings: | |
| # known_face_names.append(name) | |
| # known_face_encodings.append(encoding) | |
| # Save the trained model for future use | |
| with open('static/data/dataset/trained_faces.pkl', 'wb') as f: | |
| pickle.dump((known_face_encodings, known_face_names), f) | |
| generateDatasetImages() |