dataCollection.py

import cv2
import math
import numpy as np
import os

from config import DATA_PATH, IMAGE_SIZE, OFFSET, HAARCASCADE_PATH

class HaarFaceModule:
    def __init__(self, img_size=IMAGE_SIZE[0], offset=OFFSET):
        self.img_size = img_size 
        self.offset = offset

        # Haar Cascades Face Detection Initialization
        cascade_path = cv2.data.haarcascades + HAARCASCADE_PATH
        self.face_cascade = cv2.CascadeClassifier(cascade_path)

    def find_faces(self, img, draw=True):
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        self.faces = self.face_cascade.detectMultiScale(
            gray,
            scaleFactor=1.3,
            minNeighbors=5,
            minSize=(100, 100)
        )

        img_draw = img.copy()
        for (x, y, w, h) in self.faces:
            if draw:
                cv2.rectangle(
                    img_draw, 
                    (x - self.offset, y - self.offset),
                    (x + w + self.offset, y + h + self.offset),
                    (0, 255, 0), 
                    2
                )
        return img_draw

    def get_bounding_boxes(self):
        faces_data = []
        # Store width, height for bounding box
        for (x, y, w, h) in self.faces:
            faces_data.append({'bbox': (x, y, w, h)})
        return faces_data

    def crop_and_resize(self, img, bbox):
        x, y, w, h = bbox
        img_white = np.ones((self.img_size, self.img_size, 3), np.uint8) * 255

        x_start = max(0, x - self.offset)
        y_start = max(0, y - self.offset)
        x_end = min(img.shape[1], x + w + self.offset)
        y_end = min(img.shape[0], y + h + self.offset)

        img_crop = img[y_start:y_end, x_start:x_end]
        if img_crop.size == 0:
            print("Empty cropped image.")
            return img_white

        # We need the dimensions of the crop to calculate the aspect ratio correctly
        crop_h, crop_w, _ = img_crop.shape

        aspect_ratio = crop_h / max(crop_w, 1)  # Prevent division by zero
        if aspect_ratio > 1:
            k = self.img_size / crop_h
            w_cal = math.floor(k * crop_w)
            img_resize = cv2.resize(img_crop, (w_cal, self.img_size))
            w_gap = math.ceil((self.img_size - w_cal) / 2)
            
            # Make sure we don't go out of bounds due to rounding
            end_col = min(w_gap + w_cal, self.img_size)
            img_white[:, w_gap:end_col] = img_resize[:, :end_col - w_gap]
        else:
            k = self.img_size / crop_w
            h_cal = math.floor(k * crop_h)
            img_resize = cv2.resize(img_crop, (self.img_size, h_cal))
            h_gap = math.ceil((self.img_size - h_cal) / 2)
            
            # Make sure we don't go out of bounds due to rounding
            end_row = min(h_gap + h_cal, self.img_size)
            img_white[h_gap:end_row, :] = img_resize[:end_row - h_gap, :]

        return img_white

    def save_image(self, img_white, folder, counter):
        if not os.path.exists(folder):
            os.makedirs(folder)
            
        filename = os.path.join(folder, f"{counter}.jpg")
        cv2.imwrite(filename, img_white)
        print(f"Image saved: {filename}")

def main():
    person_name = input("Enter the name of the person you are collecting data for: ").strip()
    if not person_name:
        print("Name cannot be empty.")
        return
        
    target_folder = os.path.join(DATA_PATH, person_name)
    print(f"Data will be saved to: {target_folder}")

    cap = cv2.VideoCapture(0)
    if not cap.isOpened():
        print("Error: Could not open webcam.")
        return

    detector = HaarFaceModule()
    
    # Check current counter
    if not os.path.exists(target_folder):
        os.makedirs(target_folder)
        counter = 0
    else:
        # Start from the number of existing files
        existing_files = [f for f in os.listdir(target_folder) if f.casefold().endswith((".jpg", ".png", ".jpeg"))]
        counter = len(existing_files)

    print("Press 's' to save a face. Press 'q' or '0' to quit.")

    while True:
        success, img = cap.read()
        if not success:
            break
        
        # Real-time processing
        img = cv2.flip(img, 1)
        
        # Find faces using Haar cascades
        img_display = detector.find_faces(img)
        faces_data = detector.get_bounding_boxes()

        if faces_data:
            face_bbox = faces_data[0]['bbox']
            # Crop the aligned face region
            img_white = detector.crop_and_resize(img, face_bbox)
            
            # Show preview
            cv2.imshow("Cropped Face Preview", img_white)

            key = cv2.waitKey(1)
            # Press 's' to save the extracted face
            if key == ord('s'):
                counter += 1
                detector.save_image(img_white, target_folder, counter)

        cv2.imshow("Face Detection", img_display)

        key = cv2.waitKey(1)
        if key == ord('0') or key == ord('q') or counter >= 2000:
            break

    cap.release()
    cv2.destroyAllWindows()

if __name__ == "__main__":
    main()