app.py

from io import BytesIO
#import base64
import argparse
import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import transforms, models
from PIL import Image
import gradio as gr
import cv2
import time
from collections import Counter
#from functools import reduce

# Combined Code for Beard and Hairstyle Detection and Styling
frames=[]
male_background_image_paths = [
    "Data/AdobeColorFunko/Outfits/MenOutfits/DummyDress1.png",
    "Data/AdobeColorFunko/Outfits/MenOutfits/GlassesDummy.png",
    "Data/AdobeColorFunko/Outfits/MenOutfits/DummyDress3.png"
]

female_background_image_paths = [
    "Data/AdobeColorFunko/Outfits/WomenOutfits/WomenOne.png",
    "Data/AdobeColorFunko/Outfits/WomenOutfits/WomenTwo.png",
    "Data/AdobeColorFunko/Outfits/WomenOutfits/WomenThree.png"
]

def parse_args():
    parser = argparse.ArgumentParser(description='Funko Demo')
    parser.add_argument(
        '--device', type=str, default='cuda:0', help='CPU/CUDA device option.')
    parser.add_argument(
        '--camera-id', type=int, default=0, help='Camera device id.')
    args = parser.parse_args()
    return args
    
def capture_frame_from_webcam(duration=5):
    data_transforms = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
    args = parse_args()
    device = torch.device(args.device)
    
    cap = cv2.VideoCapture(args.camera_id)  # Open the default webcam (usually ID 0)

    frames = []
    start_time = time.time()

    while (time.time() - start_time) < duration:
        ret, frame = cap.read()
        if not ret:
            break

        # Preprocess the frame and store it in the list
        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)  # Convert to RGB format
        frame_pil = Image.fromarray(frame)  # Convert to PIL Image
        frame_tensor = data_transforms(frame_pil)  # Preprocess
        frames.append(frame_tensor)

        # Display the video stream to the user
        cv2.imshow("Video Capture", cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
        cv2.waitKey(1)  # Adjust the delay (milliseconds) as needed for display

    cap.release()
    cv2.destroyAllWindows()  # Close the video stream window

    return frames

class GenderClassifier:
    def __init__(self, model_path, class_names):
        self.model = models.resnet18(pretrained=False)
        num_ftrs = self.model.fc.in_features
        self.model.fc = nn.Linear(num_ftrs, len(class_names))
        self.load_model(model_path)
        self.model.eval()
        self.data_transforms = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
        self.class_names = class_names

    def preprocess_image(self, image_path):
        image = Image.open(image_path).convert("RGB")
        image = self.data_transforms(image)
        image = image.unsqueeze(0)
        return image

    def load_model(self, model_path):
        if torch.cuda.is_available():
            self.model.load_state_dict(torch.load(model_path))
        else:
            self.model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))

    def classify_gender(self, image_path):
        input_image = self.preprocess_image(image_path)

        with torch.no_grad():
            predictions = self.model(input_image)

        probabilities = torch.nn.functional.softmax(predictions[0], dim=0)
        predicted_class = torch.argmax(probabilities).item()
        predicted_label = self.class_names[predicted_class]

        return predicted_label
    
    def classify_from_frames(self, image, image_type):
        input_image = None
        if image_type == True:
            input_image = self.preprocess_image(image)
        else:
            input_image = image.unsqueeze(0)
        with torch.no_grad():
            predictions = self.model(input_image)
        probabilities = torch.nn.functional.softmax(predictions[0], dim=0)
        predicted_class = torch.argmax(probabilities).item()
        predicted_label = self.class_names[predicted_class]
        return predicted_label

    
class WomenHairStyleClassifier:
    def __init__(self, model_path, class_names):
        self.model = models.resnet18(pretrained=False)
        num_ftrs = self.model.fc.in_features
        self.model.fc = nn.Linear(num_ftrs, len(class_names))
        self.load_model(model_path)
        self.model.eval()
        self.data_transforms = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
        self.class_names = class_names

    def preprocess_image(self, image_path):
        image = Image.open(image_path).convert("RGB")
        image = self.data_transforms(image)
        image = image.unsqueeze(0)
        return image
    
    def load_model(self, model_path):
        if torch.cuda.is_available():
            self.model.load_state_dict(torch.load(model_path))
        else:
            self.model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))

    def classify_hairStyle(self, image_path):
        input_image = self.preprocess_image(image_path)

        with torch.no_grad():
            predictions = self.model(input_image)

        probabilities = torch.nn.functional.softmax(predictions[0], dim=0)
        predicted_class = torch.argmax(probabilities).item()
        predicted_label = self.class_names[predicted_class]

        return predicted_label
    def classify_from_frames(self, image, image_type):
        input_image = None
        if image_type == True:
            input_image = self.preprocess_image(image)
        else:
            input_image = image.unsqueeze(0)
        with torch.no_grad():
            predictions = self.model(input_image)
        probabilities = torch.nn.functional.softmax(predictions[0], dim=0)
        predicted_class = torch.argmax(probabilities).item()
        predicted_label = self.class_names[predicted_class]
        return predicted_label

        # Return a single prediction for the entire video
        # You can choose to use the majority vote or any other method to determine the final prediction
        final_prediction = max(set(predictions), key=predictions.count)
        return final_prediction

    
class WomenHairColorClassifier:
    def __init__(self, model_path, class_names):
        self.model = models.resnet18(pretrained=False)
        num_ftrs = self.model.fc.in_features
        self.model.fc = nn.Linear(num_ftrs, len(class_names))
        self.load_model(model_path)
        self.model.eval()
        self.data_transforms = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
        self.class_names = class_names

    def preprocess_image(self, image_path):
        image = Image.open(image_path).convert("RGB")
        image = self.data_transforms(image)
        image = image.unsqueeze(0)
        return image

    def load_model(self, model_path):
        if torch.cuda.is_available():
            self.model.load_state_dict(torch.load(model_path))
        else:
            self.model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))

    def classify_hairColor(self, image_path):
        input_image = self.preprocess_image(image_path)

        with torch.no_grad():
            predictions = self.model(input_image)

        probabilities = torch.nn.functional.softmax(predictions[0], dim=0)
        predicted_class = torch.argmax(probabilities).item()
        predicted_label = self.class_names[predicted_class]

        return predicted_label
    
    def classify_from_frames(self, image, image_type):
        input_image = None
        if image_type == True:
            input_image = self.preprocess_image(image)
        else:
            input_image = image.unsqueeze(0)
        with torch.no_grad():
            predictions = self.model(input_image)
        probabilities = torch.nn.functional.softmax(predictions[0], dim=0)
        predicted_class = torch.argmax(probabilities).item()
        predicted_label = self.class_names[predicted_class]
        return predicted_label

# Function to classify beard style
class BeardClassifier:
    def __init__(self, model_path, class_names):
        self.model = torch.hub.load('pytorch/vision', 'resnet50', pretrained=False)
        num_ftrs = self.model.fc.in_features
        self.model.fc = torch.nn.Linear(num_ftrs, len(class_names))
        self.load_model(model_path)
        self.model.eval()
        self.data_transforms = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
        self.class_names = class_names

    def preprocess_image(self, image):
        image = Image.open(image).convert("RGB")
        image = self.data_transforms(image)
        image = image.unsqueeze(0)
        return image

    def load_model(self, model_path):
        if torch.cuda.is_available():
            self.model.load_state_dict(torch.load(model_path))
        else:
            self.model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))

    def classify_beard(self, image):
        input_image = self.preprocess_image(image)
        with torch.no_grad():
            predictions = self.model(input_image)
        probabilities = torch.nn.functional.softmax(predictions[0], dim=0)
        predicted_class = torch.argmax(probabilities).item()
        predicted_label = self.class_names[predicted_class]
        return predicted_label

    def classify_from_frames(self, image, image_type):
        input_image = None
        if image_type == True:
            input_image = self.preprocess_image(image)
        else:
            input_image = image.unsqueeze(0)
        with torch.no_grad():
            predictions = self.model(input_image)
        probabilities = torch.nn.functional.softmax(predictions[0], dim=0)
        predicted_class = torch.argmax(probabilities).item()
        predicted_label = self.class_names[predicted_class]
        return predicted_label


# Function to classify beard color
class BeardColorClassifier:
    def __init__(self, model_path, class_names):
        self.model = torch.hub.load('pytorch/vision', 'resnet50', pretrained=False)
        num_ftrs = self.model.fc.in_features
        self.model.fc = torch.nn.Linear(num_ftrs, len(class_names))
        self.load_model(model_path)
        self.model.eval()
        self.data_transforms = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
        self.class_names = class_names

    def preprocess_image(self, image):
        image = Image.open(image).convert("RGB")
        image = self.data_transforms(image)
        image = image.unsqueeze(0)
        return image

    def load_model(self, model_path):
        if torch.cuda.is_available():
            self.model.load_state_dict(torch.load(model_path))
        else:
            self.model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))

    def classify_beard_color(self, image):
        input_image = self.preprocess_image(image)
        with torch.no_grad():
            predictions = self.model(input_image)
        probabilities = torch.nn.functional.softmax(predictions[0], dim=0)
        predicted_class = torch.argmax(probabilities).item()
        predicted_label = self.class_names[predicted_class]
        return predicted_label

    def classify_from_frames(self, image, image_type):
        input_image = None
        if image_type == True:
            input_image = self.preprocess_image(image)
        else:
            input_image = image.unsqueeze(0)
        with torch.no_grad():
            predictions = self.model(input_image)
        probabilities = torch.nn.functional.softmax(predictions[0], dim=0)
        predicted_class = torch.argmax(probabilities).item()
        predicted_label = self.class_names[predicted_class]
        return predicted_label



# Function to classify hairstyle
class HairStyleClassifier:
    def __init__(self, model_path, class_names):
        self.model = torch.hub.load('pytorch/vision', 'resnet50', pretrained=False)
        num_ftrs = self.model.fc.in_features
        self.model.fc = torch.nn.Linear(num_ftrs, len(class_names))
        self.load_model(model_path)
        self.model.eval()
        self.data_transforms = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
        self.class_names = class_names

    def preprocess_image(self, image):
        image = Image.open(image).convert("RGB")
        image = self.data_transforms(image)
        image = image.unsqueeze(0)
        return image

    def load_model(self, model_path):
        if torch.cuda.is_available():
            self.model.load_state_dict(torch.load(model_path))
        else:
            self.model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))

    def classify_hair(self, image):
        input_image = self.preprocess_image(image)
        with torch.no_grad():
            predictions = self.model(input_image)
        probabilities = torch.nn.functional.softmax(predictions[0], dim=0)
        predicted_class = torch.argmax(probabilities).item()
        predicted_label = self.class_names[predicted_class]
        return predicted_label

    def classify_from_frames(self, image, image_type):
        input_image = None
        if image_type == True:
            input_image = self.preprocess_image(image)
        else:
            input_image = image.unsqueeze(0)
        with torch.no_grad():
            predictions = self.model(input_image)
        probabilities = torch.nn.functional.softmax(predictions[0], dim=0)
        predicted_class = torch.argmax(probabilities).item()
        predicted_label = self.class_names[predicted_class]
        return predicted_label

class MenHairColorClassifier:
    def __init__(self, model_path, class_names):
        self.model = torch.hub.load('pytorch/vision', 'resnet50', pretrained=False)
        num_ftrs = self.model.fc.in_features
        self.model.fc = torch.nn.Linear(num_ftrs, len(class_names))
        self.load_model(model_path)
        self.model.eval()
        self.data_transforms = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
        ])
        self.class_names = class_names

    def preprocess_image(self, image):
        image = Image.open(image).convert("RGB")
        image = self.data_transforms(image)
        image = image.unsqueeze(0)
        return image

    def load_model(self, model_path):
        if torch.cuda.is_available():
            self.model.load_state_dict(torch.load(model_path))
        else:
            self.model.load_state_dict(torch.load(model_path, map_location=torch.device('cpu')))

    def classify_menHair_color(self, image):
        input_image = self.preprocess_image(image)
        with torch.no_grad():
            predictions = self.model(input_image)
        probabilities = torch.nn.functional.softmax(predictions[0], dim=0)
        predicted_class = torch.argmax(probabilities).item()
        predicted_label = self.class_names[predicted_class]
        return predicted_label
    
    def classify_from_frames(self, image, image_type):
        input_image = None
        if image_type == True:
            input_image = self.preprocess_image(image)
        else:
            input_image = image.unsqueeze(0)
        with torch.no_grad():
            predictions = self.model(input_image)
        probabilities = torch.nn.functional.softmax(predictions[0], dim=0)
        predicted_class = torch.argmax(probabilities).item()
        predicted_label = self.class_names[predicted_class]
        return predicted_label




def dummy_eye(background_image, x, y, placeholder_image_path, x_coordinate, y_coordinate):
    placeholder_image = Image.open(placeholder_image_path)
    target_size = (x, y)
    placeholder_image = placeholder_image.resize(target_size, Image.LANCZOS)
    placeholder_array = np.array(placeholder_image)
    placeholder_width, placeholder_height = placeholder_image.size
    region_box = (x_coordinate, y_coordinate, x_coordinate + placeholder_width, y_coordinate + placeholder_height)
    placeholder_mask = placeholder_image.split()[3] if placeholder_image.mode == 'RGBA' else None
    background_image.paste(placeholder_image, region_box, mask=placeholder_mask)
    background_array = np.array(background_image)

# Function to overlay a beard on a background image
def process_image_Beard(background_image, x, placeholder_image_path, x_coordinate, y_coordinate):
    placeholder_image = Image.open(placeholder_image_path)
    target_size = (x, x)
    placeholder_image = placeholder_image.resize(target_size, Image.LANCZOS)
    placeholder_array = np.array(placeholder_image)
    placeholder_width, placeholder_height = placeholder_image.size
    region_box = (x_coordinate, y_coordinate, x_coordinate + placeholder_width, y_coordinate + placeholder_height)
    placeholder_mask = placeholder_image.split()[3] if placeholder_image.mode == 'RGBA' else None
    background_image.paste(placeholder_image, region_box, mask=placeholder_mask)
    background_array = np.array(background_image)
    placeholder_alpha = placeholder_image.split()[3] if placeholder_image.mode == 'RGBA' else None

def process_image_WomanHair(background_image, x, y, placeholder_image_path, x_coordinate, y_coordinate):
    placeholder_image = Image.open(placeholder_image_path)
    target_size = (x, y)
    placeholder_image = placeholder_image.resize(target_size, Image.LANCZOS)
    placeholder_array = np.array(placeholder_image)
    placeholder_width, placeholder_height = placeholder_image.size
    region_box = (x_coordinate, y_coordinate, x_coordinate + placeholder_width, y_coordinate + placeholder_height)
    placeholder_mask = placeholder_image.split()[3] if placeholder_image.mode == 'RGBA' else None
    background_image.paste(placeholder_image, region_box, mask=placeholder_mask)
    background_array = np.array(background_image)
    placeholder_alpha = placeholder_image.split()[3] if placeholder_image.mode == 'RGBA' else None


def add_eyebrow(background_image, x_coordinate, y_coordinate, eyebrow_image_path):
    eyebrow_image = Image.open(eyebrow_image_path)
    target_size = (200, 200)  # Adjust the size as needed
    eyebrow_image = eyebrow_image.resize(target_size, Image.LANCZOS)
    region_box = (x_coordinate, y_coordinate, x_coordinate + eyebrow_image.width, y_coordinate + eyebrow_image.height)
    eyebrow_mask = eyebrow_image.split()[3] if eyebrow_image.mode == 'RGBA' else None
    background_image.paste(eyebrow_image, region_box, mask=eyebrow_mask)
    background_array = np.array(background_image)


    
    
# Function to overlay a hairstyle on a background image
def process_image_menHair(background_image, x, y, placeholder_image_path, x_coordinate, y_coordinate):
    placeholder_image = Image.open(placeholder_image_path)
    target_size = (x, y)
    placeholder_image = placeholder_image.resize(target_size, Image.LANCZOS)
    placeholder_array = np.array(placeholder_image)
    placeholder_width, placeholder_height = placeholder_image.size
    region_box = (x_coordinate, y_coordinate, x_coordinate + placeholder_width, y_coordinate + placeholder_height)
    placeholder_mask = placeholder_image.split()[3] if placeholder_image.mode == 'RGBA' else None
    background_image.paste(placeholder_image, region_box, mask=placeholder_mask)
    background_array = np.array(background_image)
    placeholder_alpha = placeholder_image.split()[3] if placeholder_image.mode == 'RGBA' else None



# Function to generate Funko figurines
def generate_funko_figurines(input_image):
    # Detect and classify gender
    gender_classifier = GenderClassifier('Data/FunkoSavedModels/Gender.pt', ['Female', 'Male'])
    predicted_gender = gender_classifier.classify_gender(input_image)
    # Process background images and apply beard style and color along with hair style and color
    final_images = []

    if predicted_gender == 'Male':
        background_image_paths = male_background_image_paths
    if predicted_gender == 'Female':
        background_image_paths = female_background_image_paths
        
    for background_image_paths in background_image_paths:
        background_image = Image.open(background_image_paths)
        x_coordinate = 90
        y_coordinate = 50
        add_eyebrow(background_image, 115, 80, "Data/AdobeColorFunko/EyezBrowz/Eyebrow.png")
        #dummy_eye(background_image, 245, 345, 'Data/AdobeColorFunko/EyezBrowz/MaleEye.png', x_coordinate, y_coordinate)
        if predicted_gender == 'Male':
            x = 245
            y = 345
            placeholder_image_path = f"Data/AdobeColorFunko/EyezBrowz/{predicted_gender}Eye.png"
            x_coordinate = 90
            y_coordinate = 50
            dummy_eye(background_image, x, y, placeholder_image_path, x_coordinate, y_coordinate)
            # Detect and classify beard style
            beard_classifier = BeardClassifier('Data/FunkoSavedModels/FunkoResnet50BeardStyle.pt', ['Bandholz', 'CleanShave', 'FullGoatee', 'Moustache', 'RapIndustryStandards', 'ShortBeard'])
            predicted_style_label = beard_classifier.classify_beard(input_image)

            # Detect and classify beard color
            beard_color_classifier = BeardColorClassifier('Data/FunkoSavedModels/FunkoResnet50BeardColor.pt', ['Black', 'DarkBrown', 'Ginger', 'LightBrown', 'SaltAndPepper', 'White'])
            predicted_color_label = beard_color_classifier.classify_beard_color(input_image)

            # Classify hairstyle
            hair_style_classifier = HairStyleClassifier('Data/FunkoSavedModels/FunkoResnet50MenHairStyle.pt', ['Afro', 'Bald', 'Puff', 'Spike'])
            predicted_hairStyle_label = hair_style_classifier.classify_hair(input_image)

            #classify menHairColor
            menhair_color_classifier = MenHairColorClassifier('Data/FunkoSavedModels/FunkoResnet50MenHairColor.pt', ['Black', 'DarkBrown', 'Ginger', 'LightBrown', 'SaltAndPepper', 'White'])
            predicted_menhairColor_label = menhair_color_classifier.classify_menHair_color(input_image)

            if predicted_style_label == 'Bandholz':
                process_image_Beard(background_image, 320,
                                     f"Data/AdobeColorFunko/Beard/Bandholz/{predicted_color_label}.png",
                                     50, 142)

            if predicted_style_label == 'ShortBeard':
                process_image_Beard(background_image, 300,
                                     f"Data/AdobeColorFunko/Beard/ShortBeard/{predicted_color_label}.png",
                                     62, 118)

            if predicted_style_label == 'FullGoatee':
                process_image_Beard(background_image, 230,
                                     f"Data/AdobeColorFunko/Beard/Goatee/{predicted_color_label}.png",
                                     96, 168)

            if predicted_style_label == 'RapIndustryStandards':
                process_image_Beard(background_image, 290,
                                     f"Data/AdobeColorFunko/Beard/RapIndustry/{predicted_color_label}.png",
                                     67, 120)

            if predicted_style_label == 'Moustache':
                process_image_Beard(background_image, 220,
                                     f"Data/AdobeColorFunko/Beard/Moustache/{predicted_color_label}.png",
                                     100, 160)

            if predicted_style_label == 'CleanShave':
                process_image_Beard(background_image, 220,
                                     f"Data/AdobeColorFunko/Beard/CleanShave/{predicted_color_label}.png",
                                     100, 160)

            # Add other conditions for different beard styles

            # Overlay hairstyle
            if predicted_hairStyle_label == 'Afro':
                process_image_menHair(background_image, 336, 420,
                                       f"Data/AdobeColorFunko/MenHairstyle/Afro/{predicted_menhairColor_label}.png",
                                       41, 76)

            if predicted_hairStyle_label == 'Puff':
                process_image_menHair(background_image, 305, 420,
                                       f"Data/AdobeColorFunko/MenHairstyle/Puff/{predicted_menhairColor_label}.png",
                                       56, 68)

            if predicted_hairStyle_label == 'Spike':
                process_image_menHair(background_image, 310, 420,
                                       f"Data/AdobeColorFunko/MenHairstyle/Spike/{predicted_menhairColor_label}.png",
                                       52, 70)

            if predicted_hairStyle_label == 'Bald':
                process_image_menHair(background_image, 310, 420,
                                       f"Data/AdobeColorFunko/MenHairstyle/Bald/{predicted_menhairColor_label}.png",
                                       67, 120)


        if predicted_gender == 'Female':
            x = 245
            y = 345
            placeholder_image_path = f"Data/AdobeColorFunko/EyezBrowz/{predicted_gender}Eye.png"
            x_coordinate = 90
            y_coordinate = 50
            dummy_eye(background_image, x, y, placeholder_image_path, x_coordinate, y_coordinate)
            WomenHairStyle_classifier = WomenHairStyleClassifier('Data/FunkoSavedModels/WomenHairStyle.pt', ['MediumLength', 'ShortHair', 'SidePlait'])
            predicted_WomenHairStyle = WomenHairStyle_classifier.classify_hairStyle(input_image)

            WomenHairColor_classifier = WomenHairColorClassifier('Data/FunkoSavedModels/WomenHairColor.pt', ['Black', 'Brown', 'Ginger', 'White'])
            predicted_WomenHairColor = WomenHairColor_classifier.classify_hairColor(input_image)
            if predicted_WomenHairStyle == 'MediumLength':
                process_image_WomanHair(background_image, 300,460,
                                     f"Data/AdobeColorFunko/WomenHairstyle/MediumLength/{predicted_WomenHairColor}.png",
                                     56, 50)

            if predicted_WomenHairStyle == 'ShortHair':
                process_image_WomanHair(background_image, 270,460,
                                     f"Data/AdobeColorFunko/WomenHairstyle/ShortHair/{predicted_WomenHairColor}.png",
                                     61, 49)

            if predicted_WomenHairStyle == 'SidePlait':
                process_image_WomanHair(background_image, 300,450,
                                     f"Data/AdobeColorFunko/WomenHairstyle/SidePlait/{predicted_WomenHairColor}.png",
                                     54, 56)


        # Convert the resulting image to base64
        buffered = BytesIO()
        background_image.save(buffered, format="PNG")
        #base64_image = base64.b64encode(buffered.getvalue()).decode("utf-8")
        final_images.append(background_image)

    return final_images

def Igenerate_funko_figurines():
    
    predicted_women_hairstyle = None
    predicted_women_haircolor = None
    predicted_gender = None
    predicted_style_label = None
    predicted_color_label = None
    predicted_hairstyle_label = None
    predicted_menhaircolor_label = None
    background_image_paths = None

    # Capture video from the webcam for 7 seconds
    # Initialize variables to store frames and track time


    # Classify women hairstyle
    women_hairstyle_classifier = WomenHairStyleClassifier('Data/FunkoSavedModels/WomenHairStyle.pt', ['MediumLength', 'ShortHair', 'SidePlait'])
    
    # Classify women hair color
    women_hair_color_classifier = WomenHairColorClassifier('Data/FunkoSavedModels/WomenHairColor.pt', ['Black', 'Brown', 'Ginger', 'White'])
    
    # Detect and classify gender
    gender_classifier = GenderClassifier('Data/FunkoSavedModels/Gender.pt', ['Female', 'Male'])

    # Detect and classify beard style
    beard_classifier = BeardClassifier('Data/FunkoSavedModels/FunkoResnet50BeardStyle.pt', ['Bandholz', 'CleanShave', 'FullGoatee', 'Moustache', 'RapIndustryStandards', 'ShortBeard'])
    
    # Detect and classify beard color
    beard_color_classifier = BeardColorClassifier('Data/FunkoSavedModels/FunkoResnet50BeardColor.pt', ['Black', 'DarkBrown', 'Ginger', 'LightBrown', 'SaltAndPepper', 'White'])

    # Classify hairstyle
    hair_style_classifier = HairStyleClassifier('Data/FunkoSavedModels/FunkoResnet50MenHairStyle.pt', ['Afro', 'Bald', 'Puff', 'Spike'])
    
    #classify menHairColor
    menhair_color_classifier = MenHairColorClassifier('Data/FunkoSavedModels/FunkoResnet50MenHairColor.pt', ['Black', 'DarkBrown', 'Ginger', 'LightBrown', 'SaltAndPepper', 'White'])
    
    def predict_male_features_from_frames(frame):
        return  [
            beard_classifier.classify_from_frames(image=frame,image_type=False),
            beard_color_classifier.classify_from_frames(image=frame,image_type=False),
            hair_style_classifier.classify_from_frames(image=frame,image_type=False),
            menhair_color_classifier.classify_from_frames(image=frame,image_type=False)
        ]
    
    def predict_female_features_from_frames(frame):
        return  [
            women_hairstyle_classifier.classify_from_frames(image=frame,image_type=False),
            women_hair_color_classifier.classify_from_frames(image=frame,image_type=False),
        ]
    
    def predict_gender_from_frames(frame):
        return gender_classifier.classify_from_frames(image=frame,image_type=False)
    

    
    print("Capturing frames")
    frames = capture_frame_from_webcam(duration=7)
    print("Frames captured")
    print("Predictions started")
    # time counting
    gp_start = time.time()
    gender_predictions = map(predict_gender_from_frames, frames)
    gender_counter = Counter(gender_predictions)
    predicted_gender = gender_counter.most_common(1)[0][0]
    # time counting
    gp_end = time.time()
    print(f'Predicted Gender: {predicted_gender} and it took {round(gp_end - gp_start)}s')

    if predicted_gender == 'Male':
        # time counting
        mp_start = time.time()
        background_image_paths = male_background_image_paths
     
        facial_feature_predictions = map(predict_male_features_from_frames, frames)
        beard_style_counter = Counter()
        beard_color_counter = Counter()
        hair_style_label_counter = Counter()
        menhair_color_counter = Counter()


        for pred in facial_feature_predictions:
            beard_style_counter[pred[0]] += 1
            beard_color_counter[pred[1]] += 1
            hair_style_label_counter[pred[2]] += 1
            menhair_color_counter[pred[3]] += 1

        predicted_style_label = beard_style_counter.most_common(1)[0][0]
        predicted_color_label = beard_color_counter.most_common(1)[0][0]
        predicted_hairstyle_label = hair_style_label_counter.most_common(1)[0][0]
        predicted_menhaircolor_label = menhair_color_counter.most_common(1)[0][0]
        # time counting
        mp_end = time.time()


        print("Predictions are:\n")
        print(predicted_style_label,predicted_color_label,predicted_hairstyle_label,predicted_menhaircolor_label)
        print(f'\nand it took {round(mp_end - mp_start)}s')

    if predicted_gender == 'Female':
        background_image_paths = female_background_image_paths
        women_hairstyle_counter = women_haircolor_counter = Counter()
        facial_feature_predictions = map(predict_female_features_from_frames, frames)
        
        for pred in facial_feature_predictions:
            women_hairstyle_counter[pred[0]] += 1 
            women_haircolor_counter[pred[1]] += 1 
            
        predicted_women_hairstyle = women_hairstyle_counter.most_common(1)[0][0]
        predicted_women_haircolor = women_haircolor_counter.most_common(1)[0][0]
        
    # Process background images and apply beard style and color along with hair style and color
    final_images = []


    for path in background_image_paths:
        background_image = Image.open(path)
        x_coordinate = 90
        y_coordinate = 50
        add_eyebrow(background_image, 115, 80, "Data/AdobeColorFunko/EyezBrowz/Eyebrow.png")
        #dummy_eye(background_image, 245, 345, 'Data/AdobeColorFunko/EyezBrowz/MaleEye.png', x_coordinate, y_coordinate)
        if predicted_gender == 'Male':
            x = 245
            y = 345
            placeholder_image_path = f"Data/AdobeColorFunko/EyezBrowz/{predicted_gender}Eye.png"
            x_coordinate = 90
            y_coordinate = 50
            dummy_eye(background_image, x, y, placeholder_image_path, x_coordinate, y_coordinate)

            if predicted_style_label == 'Bandholz':
                process_image_Beard(background_image, 320,
                                     f"Data/AdobeColorFunko/Beard/Bandholz/{predicted_color_label}.png",
                                     50, 142)

            if predicted_style_label == 'ShortBeard':
                process_image_Beard(background_image, 300,
                                     f"Data/AdobeColorFunko/Beard/ShortBeard/{predicted_color_label}.png",
                                     62, 118)

            if predicted_style_label == 'FullGoatee':
                process_image_Beard(background_image, 230,
                                     f"Data/AdobeColorFunko/Beard/Goatee/{predicted_color_label}.png",
                                     96, 168)

            if predicted_style_label == 'RapIndustryStandards':
                process_image_Beard(background_image, 290,
                                     f"Data/AdobeColorFunko/Beard/RapIndustry/{predicted_color_label}.png",
                                     67, 120)

            if predicted_style_label == 'Moustache':
                process_image_Beard(background_image, 220,
                                     f"Data/AdobeColorFunko/Beard/Moustache/{predicted_color_label}.png",
                                     100, 160)

            if predicted_style_label == 'CleanShave':
                process_image_Beard(background_image, 220,
                                     f"Data/AdobeColorFunko/Beard/CleanShave/{predicted_color_label}.png",
                                     100, 160)

            # Add other conditions for different beard styles

            # Overlay hairstyle
            if predicted_hairstyle_label == 'Afro':
                process_image_menHair(background_image, 336, 420,
                                       f"Data/AdobeColorFunko/MenHairstyle/Afro/{predicted_menhaircolor_label}.png",
                                       41, 76)

            if predicted_hairstyle_label == 'Puff':
                process_image_menHair(background_image, 305, 420,
                                       f"Data/AdobeColorFunko/MenHairstyle/Puff/{predicted_menhaircolor_label}.png",
                                       56, 68)

            if predicted_hairstyle_label == 'Spike':
                process_image_menHair(background_image, 310, 420,
                                       f"Data/AdobeColorFunko/MenHairstyle/Spike/{predicted_menhaircolor_label}.png",
                                       52, 70)

            if predicted_hairstyle_label == 'Bald':
                process_image_menHair(background_image, 310, 420,
                                       f"Data/AdobeColorFunko/MenHairstyle/Bald/{predicted_menhaircolor_label}.png",
                                       67, 120)


        if predicted_gender == 'Female':
            x = 245
            y = 345
            placeholder_image_path = f"Data/AdobeColorFunko/EyezBrowz/{predicted_gender}Eye.png"
            x_coordinate = 90
            y_coordinate = 50
            dummy_eye(background_image, x, y, placeholder_image_path, x_coordinate, y_coordinate)
            if predicted_women_hairstyle == 'MediumLength':
                process_image_WomanHair(background_image, 300,460,
                                     f"Data/AdobeColorFunko/WomenHairstyle/MediumLength/{predicted_women_haircolor}.png",
                                     56, 50)

            if predicted_women_hairstyle == 'ShortHair':
                process_image_WomanHair(background_image, 270,460,
                                     f"Data/AdobeColorFunko/WomenHairstyle/ShortHair/{predicted_women_haircolor}.png",
                                     61, 49)

            if predicted_women_hairstyle == 'SidePlait':
                process_image_WomanHair(background_image, 300,450,
                                     f"Data/AdobeColorFunko/WomenHairstyle/SidePlait/{predicted_women_haircolor}.png",
                                     54, 56)


        # Convert the resulting image to base64
        buffered = BytesIO()
        # background_image.resize((200,400))
        background_image.save(buffered, format="PNG")
        #base64_image = base64.b64encode(buffered.getvalue()).decode("utf-8")
        final_images.append(background_image)

    return final_images



with gr.Blocks() as demo:
    gr.Markdown(
    """
    # Funko POP! Figure Creation
    Enabling Streamlined Automation with Artificial Intelligence 
    """)
    imageComponent = gr.Image(type="filepath", height=350, width=350)
    with gr.Row():
        MyOutputs = [gr.Image(type="pil", label="Generated Image " + str(i + 1), height=400, width=400) for i in range(3)]
    submitButton = gr.Button(value="Submit")
    submitButton.click(generate_funko_figurines, inputs=imageComponent, outputs=MyOutputs)
    RecordButton = gr.Button(value="Generate My Custom Funko POP")
    RecordButton.click(Igenerate_funko_figurines, outputs=MyOutputs)

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