Spaces:
Running
Running
File size: 4,815 Bytes
829f2ca | 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 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | 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() |