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()