app.py

from io import BytesIO
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
#from functools import reduce

# Combined Code for Beard and Hairstyle Detection and Styling
male_background_image_paths = [
    "Data/AdobeColorFunko/Outfits/MenOutfits/MenOne.png",
    "Data/AdobeColorFunko/Outfits/MenOutfits/MenTwo.png",
    "Data/AdobeColorFunko/Outfits/MenOutfits/MenThree.png"
]

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

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 = models.resnet50(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_beard(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




# Function to classify beard color
class BeardColorClassifier:
    def __init__(self, model_path, class_names):
        self.model = models.resnet50(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_beard_color(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



# Function to classify hairstyle
class HairStyleClassifier:
    def __init__(self, model_path, class_names):
        self.model = models.resnet50(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_hair(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



class MenHairColorClassifier:
    def __init__(self, model_path, class_names):
        self.model = models.resnet50(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_menHair_color(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


# 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




    
    
# 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)
        if predicted_gender == 'Male':
            # 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, 460,
                                     f"Data/AdobeColorFunko/Beard/Bandholz/{predicted_color_label}.png",
                                     -20, 55)

            if predicted_style_label == 'ShortBeard':
                process_image_Beard(background_image, 405,
                                     f"Data/AdobeColorFunko/Beard/ShortBeard/{predicted_color_label}.png",
                                     10, 56)

            if predicted_style_label == 'FullGoatee':
                process_image_Beard(background_image, 180,
                                     f"Data/AdobeColorFunko/Beard/Goatee/{predicted_color_label}.png",
                                     121, 176)

            if predicted_style_label == 'RapIndustryStandards':
                process_image_Beard(background_image, 400,
                                     f"Data/AdobeColorFunko/Beard/RapIndustry/{predicted_color_label}.png",
                                     14, 62)

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

            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, 434, 530,
                                       f"Data/AdobeColorFunko/MenHairstyle/Afro/{predicted_menhairColor_label}.png",
                                       -7, -23)

            if predicted_hairStyle_label == 'Puff':
                process_image_menHair(background_image, 410, 520,
                                       f"Data/AdobeColorFunko/MenHairstyle/Puff/{predicted_menhairColor_label}.png",
                                       2, -23)

            if predicted_hairStyle_label == 'Spike':
                process_image_menHair(background_image, 419, 530,
                                       f"Data/AdobeColorFunko/MenHairstyle/Spike/{predicted_menhairColor_label}.png",
                                       -2,-22)

            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':
            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, 400,660,
                                     f"Data/AdobeColorFunko/WomenHairstyle/MediumLength/{predicted_WomenHairColor}.png",
                                     5, -45)

            if predicted_WomenHairStyle == 'ShortHair':
                process_image_WomanHair(background_image, 370,660,
                                     f"Data/AdobeColorFunko/WomenHairstyle/ShortHair/{predicted_WomenHairColor}.png",
                                     5, -45)

            if predicted_WomenHairStyle == 'SidePlait':
                process_image_WomanHair(background_image, 405,660,
                                     f"Data/AdobeColorFunko/WomenHairstyle/SidePlait/{predicted_WomenHairColor}.png",
                                     7, -45)


        # 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




with gr.Blocks() as demo:
    gr.Markdown(
    """
    # Funko POP! Figure Creation
    ### Enabling Streamlined Automation with Artificial Intelligence 
    """)
    with gr.Row():
        imageComponent = gr.Image(type="filepath", height=300, width=300)
        gr.Markdown(
        """
        # Please Consider these points when uploading your picture.
        ### a) The image should be a selfie, ideally resembling a passport-size picture.
        ### b) The background in the image should be clear, devoid of people or any visual clutter.
        ### c) Ensure the selfie has proper exposure or is in a well-lit room. 
        """)

    with gr.Row():
        MyOutputs = [gr.Image(type="pil", label="Generated Image " + str(i + 1), height=450, width=300) for i in range(3)]
    submitButton = gr.Button(value="Submit")
    submitButton.click(generate_funko_figurines, inputs=imageComponent, outputs=MyOutputs)

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