CODE/Features Extract with Augmentation.py

import os
import numpy as np
import cv2 as cv
import tensorflow as tf
from keras_facenet import FaceNet
import random

# Disable GPU usage to avoid CuDNN errors
tf.config.set_visible_devices([], 'GPU')  # Force TensorFlow to use only CPU
print("Running on CPU to avoid CuDNN errors.")

class FaceEmbeddingGenerator:
    def __init__(self, directory, output_path):
        self.directory = directory
        self.output_path = output_path
        self.target_size = (160, 160)
        self.embedder = FaceNet()
        self.embeddings = []
        self.labels = []

    def load_cropped_face(self, filename):
        """Load an already cropped face image and apply augmentation."""
        img = cv.imread(filename)
        if img is None:
            raise ValueError(f"Image {filename} not found or invalid format.")
        img = cv.cvtColor(img, cv.COLOR_BGR2RGB)

        # Ensure the image is resized to the target size
        img = cv.resize(img, self.target_size)

        # Apply augmentation
        img = self.augment_image(img)
        return img

    def augment_image(self, img):
        """Apply random augmentations to the image."""
        # Random horizontal flip
        if random.random() < 0.5:
            img = cv.flip(img, 1)

        # Random rotation
        angle = random.uniform(-10, 10)
        h, w = img.shape[:2]
        M = cv.getRotationMatrix2D((w // 2, h // 2), angle, 1)
        img = cv.warpAffine(img, M, (w, h))

        # Random brightness adjustment
        brightness_factor = random.uniform(0.8, 1.2)
        img = np.clip(img * brightness_factor, 0, 255).astype(np.uint8)

        # Random noise
        noise = np.random.normal(0, 5, img.shape).astype(np.uint8)
        img = np.clip(img + noise, 0, 255).astype(np.uint8)
        
        return img

    def get_embedding(self, face_img):
        """Generate an embedding from a face image."""
        face_img = face_img.astype('float32')
        face_img = np.expand_dims(face_img, axis=0)
        return self.embedder.embeddings(face_img)[0]  # 512D vector

    def save_batch(self, folder_output_path, embeddings, labels):
        """Save a batch of embeddings to a temporary file."""
        temp_path = folder_output_path.replace(".npz", "_temp.npz")
        np.savez_compressed(temp_path, embeddings=np.asarray(embeddings), labels=np.asarray(labels))
        print(f"Saved intermediate embeddings to {temp_path}")

    def process_folder(self, folder_path, label, batch_size=50):
        """Process a single folder of cropped images in batches."""
        batch_embeddings = []
        batch_labels = []
        count = 0

        for filename in os.listdir(folder_path):
            try:
                filepath = os.path.join(folder_path, filename)
                face = self.load_cropped_face(filepath)
                embedding = self.get_embedding(face)
                batch_embeddings.append(embedding)
                batch_labels.append(label)
                count += 1

                if count % batch_size == 0:
                    # Save batch and clear memory
                    self.save_batch(self.output_path, batch_embeddings, batch_labels)
                    self.embeddings.extend(batch_embeddings)
                    self.labels.extend(batch_labels)
                    batch_embeddings = []
                    batch_labels = []

            except Exception as e:
                print(f"Error processing {filename}: {e}")

        # Save remaining data
        if batch_embeddings:
            self.save_batch(self.output_path, batch_embeddings, batch_labels)
            self.embeddings.extend(batch_embeddings)
            self.labels.extend(batch_labels)

    def process_all_classes(self):
        """Process all folders and save embeddings for each folder separately."""
        for sub_dir in os.listdir(self.directory):
            sub_dir_path = os.path.join(self.directory, sub_dir)
            if not os.path.isdir(sub_dir_path):
                continue

            # Define output file for this folder
            folder_output_path = os.path.join(self.output_path, f"{sub_dir}_embeddings.npz")

            # Skip folder if its embeddings already exist
            if os.path.exists(folder_output_path):
                print(f"Skipping folder {sub_dir} as embeddings already exist.")
                continue

            print(f"Processing folder: {sub_dir}")

            # Clear previous embeddings and labels
            self.embeddings = []
            self.labels = []

            # Process the current folder
            self.process_folder(sub_dir_path, sub_dir)

            # Save the embeddings and labels for the current folder
            np.savez_compressed(folder_output_path, embeddings=np.asarray(self.embeddings), labels=np.asarray(self.labels))
            print(f"Saved embeddings for folder {sub_dir} to {folder_output_path}")

def merge_all_embeddings(output_dir, final_output_file):
    """Merge all folder embeddings into a single NPZ file."""
    all_embeddings = []
    all_labels = []

    # Iterate over all files in the output directory
    for filename in os.listdir(output_dir):
        filepath = os.path.join(output_dir, filename)

        # Skip non-NPZ files
        if not filename.endswith(".npz"):
            continue

        # Load the embeddings and labels from the file
        data = np.load(filepath)
        embeddings = data['embeddings']
        labels = data['labels']

        # Append to the overall list
        all_embeddings.append(embeddings)
        all_labels.append(labels)

    # Combine all embeddings and labels
    all_embeddings = np.vstack(all_embeddings)
    all_labels = np.hstack(all_labels)

    # Save the merged embeddings and labels into a final NPZ file
    np.savez_compressed(final_output_file, embeddings=all_embeddings, labels=all_labels)
    print(f"Final merged embeddings saved to {final_output_file}")

# Usage
data_dir = "/home/shanin/Desktop/SHANIN/MAIN/ALL_CODE/Face_Recognition/shanin"  # Replace with the path to your dataset
output_dir = "/home/shanin/Desktop/SHANIN/MAIN/ALL_CODE/Face_Recognition/tmp"  # Replace with output directory

# Ensure output directory exists
os.makedirs(output_dir, exist_ok=True)

# Initialize and process all classes
face_generator = FaceEmbeddingGenerator(data_dir, output_dir)
face_generator.process_all_classes()

# Path for the final merged embeddings file
final_output_file = "/home/shanin/Desktop/SHANIN/MAIN/ALL_CODE/Face_Recognition/Face_Embedding_Augment.npz"

# Merge all embeddings
merge_all_embeddings(output_dir, final_output_file)