app.py

import cv2
import cvzone
import numpy as np
import os
import gradio as gr
import mediapipe as mp
from datetime import datetime

# Load the YuNet model
model_path = 'face_detection_yunet_2023mar.onnx'
face_detector = cv2.FaceDetectorYN.create(model_path, "", (320, 320))

# Initialize MediaPipe Face Mesh
mp_face_mesh = mp.solutions.face_mesh
face_mesh = mp_face_mesh.FaceMesh(static_image_mode=False, max_num_faces=1, min_detection_confidence=0.5)

# Initialize the glass number
num = 1
overlay = cv2.imread(f'glasses/glass{num}.png', cv2.IMREAD_UNCHANGED)

# Count glasses files
def count_files_in_directory(directory):
    file_count = 0
    for root, dirs, files in os.walk(directory):
        file_count += len(files)
    return file_count

# Determine face shape
def determine_face_shape(landmarks):
    # Example logic to determine face shape based on landmarks
    # This is a simplified version and may need adjustments
    jaw_width = np.linalg.norm(landmarks[0] - landmarks[16])
    face_height = np.linalg.norm(landmarks[8] - landmarks[27])
    if jaw_width / face_height > 1.5:
        return "Round"
    elif jaw_width / face_height < 1.2:
        return "Oval"
    else:
        return "Square"

# Recommend glass shape based on face shape
def recommend_glass_shape(face_shape):
    if face_shape == "Round":
        return "Square"
    elif face_shape == "Oval":
        return "Round"
    else:
        return "Square"

directory_path = 'glasses'
total_glass_num = count_files_in_directory(directory_path)

# Change glasses
def change_glasses():
    global num, overlay
    num += 1
    if num > total_glass_num:
        num = 1
    overlay = cv2.imread(f'glasses/glass{num}.png', cv2.IMREAD_UNCHANGED)
    return overlay

# Process frame for overlay and face shape detection
def process_frame(frame):
    global overlay
    frame = np.array(frame, copy=True)
    height, width = frame.shape[:2]
    face_detector.setInputSize((width, height))
    _, faces = face_detector.detect(frame)

    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    results = face_mesh.process(frame_rgb)

    face_shape = "Unknown"
    glass_shape = "Unknown"
    if faces is not None and results.multi_face_landmarks:
        for face in faces:
            x, y, w, h = face[:4].astype(int)
            face_landmarks = face[4:14].reshape(5, 2).astype(int)

            left_eye_x, left_eye_y = face_landmarks[0].astype(int)
            right_eye_x, right_eye_y = face_landmarks[1].astype(int)

            eye_center_x = (left_eye_x + right_eye_x) // 2
            eye_center_y = (left_eye_y + right_eye_y) // 2

            delta_x = right_eye_x - left_eye_x
            delta_y = right_eye_y - left_eye_y
            angle = np.degrees(np.arctan2(delta_y, delta_x))
            angle = -angle

            overlay_resize = cv2.resize(overlay, (int(w * 1.15), int(h * 0.8)))
            overlay_center = (overlay_resize.shape[1] // 2, overlay_resize.shape[0] // 2)
            rotation_matrix = cv2.getRotationMatrix2D(overlay_center, angle, 1.0)
            overlay_rotated = cv2.warpAffine(
                overlay_resize, rotation_matrix,
                (overlay_resize.shape[1], overlay_resize.shape[0]),
                flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_CONSTANT, borderValue=(0, 0, 0, 0)
            )

            overlay_x = eye_center_x - overlay_rotated.shape[1] // 2
            overlay_y = eye_center_y - overlay_rotated.shape[0] // 2

            try:
                frame = cvzone.overlayPNG(frame, overlay_rotated, [overlay_x, overlay_y])
            except Exception as e:
                print(f"Error overlaying glasses: {e}")

            for face_landmarks_mp in results.multi_face_landmarks:
                landmarks = np.array([(lm.x * frame.shape[1], lm.y * frame.shape[0]) for lm in face_landmarks_mp.landmark])
                face_shape = determine_face_shape(landmarks)
                glass_shape = recommend_glass_shape(face_shape)

    return frame, face_shape, glass_shape

# Transform function
def transform_cv2(frame, transform):
    if transform == "cartoon":
        # prepare color
        img_color = cv2.pyrDown(cv2.pyrDown(frame))
        for _ in range(6):
            img_color = cv2.bilateralFilter(img_color, 9, 9, 7)
        img_color = cv2.pyrUp(cv2.pyrUp(img_color))

        # prepare edges
        img_edges = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY)
        img_edges = cv2.adaptiveThreshold(
            cv2.medianBlur(img_edges, 7),
            255,
            cv2.ADAPTIVE_THRESH_MEAN_C,
            cv2.THRESH_BINARY,
            9,
            2,
        )
        img_edges = cv2.cvtColor(img_edges, cv2.COLOR_GRAY2RGB)
        # combine color and edges
        img = cv2.bitwise_and(img_color, img_edges)
        return img
    elif transform == "edges":
        # perform edge detection
        img = cv2.cvtColor(cv2.Canny(frame, 100, 200), cv2.COLOR_GRAY2BGR)
        return img
    
    elif transform == "sepia":
        # apply sepia effect
        kernel = np.array([[0.272, 0.534, 0.131],
                           [0.349, 0.686, 0.168],
                           [0.393, 0.769, 0.189]])
        img = cv2.transform(frame, kernel)
        img = np.clip(img, 0, 255)  # ensure values are within byte range
        # Convert BGR to RGB if necessary (for display purposes)
        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        return img_rgb

    elif transform == "negative":
        # apply negative effect
        img = cv2.bitwise_not(frame)
        return img

    elif transform == "sketch":
        # apply sketch effect
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        inv_gray = cv2.bitwise_not(gray)
        blur = cv2.GaussianBlur(inv_gray, (21, 21), 0)
        inv_blur = cv2.bitwise_not(blur)
        img = cv2.divide(gray, inv_blur, scale=256.0)
        img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
        return img

    elif transform == "blur":
        # apply blur effect
        img = cv2.GaussianBlur(frame, (15, 15), 0)
        return img

    else:
        return frame
    
def refresh_interface():
    # Reset the image to an empty state or a default image
    input_img.update(value=None)

    # Return a message indicating the interface has been refreshed
    return "Interface refreshed!"

def save_frame(frame):
       # Convert frame to RGB
       frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

       # Create a unique filename using the current timestamp
       filename = f"saved_frame_{datetime.now().strftime('%Y%m%d_%H%M%S')}.png"

       # Save the frame
       cv2.imwrite(filename, frame)

       # Refresh the interfaceq
       refresh_interface()

       return f"Frame saved as '{filename}'"



# Gradio webcam input
def webcam_input(frame, transform):
    frame, face_shape, glass_shape = process_frame(frame)
    frame = transform_cv2(frame, transform)
    return frame, face_shape, glass_shape

# Gradio Interface
with gr.Blocks(theme=gr.themes.Soft(primary_hue="purple", secondary_hue="blue")) as demo:
    gr.Markdown("<h1 style='text-align: center; font-weight: bold;'>🤓 Glasses Virtual Try-On 🕶️👓</h1>")
    with gr.Column(elem_classes=["my-column"]):
        with gr.Group(elem_classes=["my-group"]):
            transform = gr.Dropdown(choices=["cartoon", "edges", "sepia", "negative", "sketch", "blur", "none"],
                                    value="none", label="Select Filter")
            gr.Markdown("Click the Webcam icon to start the camera, and then press the record button to start the virtual try-on.")
            input_img = gr.Image(sources=["webcam"], type="numpy", streaming=True)
            gr.Markdown("Face Shape and Recommended Glass Shape")
            face_shape_output = gr.Textbox(label="Detected Face Shape")
            glass_shape_output = gr.Textbox(label="Recommended Glass Shape")
            next_button = gr.Button("Next Glasses➡️")
            save_button = gr.Button("Save as a Picture📌")

    input_img.stream(webcam_input, [input_img, transform], [input_img, face_shape_output, glass_shape_output], stream_every=0.1)
    with gr.Row():
        next_button.click(change_glasses, [], [])
    with gr.Row():
        save_button.click(save_frame, [input_img], [])

if __name__ == "__main__":
    demo.launch(share=True)