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