Update app.py

app.py

@@ -36,7 +36,7 @@ def parse_args():
     args = parser.parse_args()
     return args
     
-def capture_frame_from_webcam(duration=7):
+def capture_frame_from_webcam(duration=5):
     data_transforms = transforms.Compose([
         transforms.Resize((224, 224)),
         transforms.ToTensor(),
@@ -108,21 +108,18 @@ class GenderClassifier:
 
         return predicted_label
     
-    def classify_from_frames(self, frames):
-        predictions = []
-
-        for frame in frames:
-            with torch.no_grad():
-                predictions_frame = self.model(frame.unsqueeze(0))
-            probabilities = torch.nn.functional.softmax(predictions_frame[0], dim=0)
-            predicted_class = torch.argmax(probabilities).item()
-            predicted_label = self.class_names[predicted_class]
-            predictions.append(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
+    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:
@@ -162,16 +159,18 @@ class WomenHairStyleClassifier:
         predicted_label = self.class_names[predicted_class]
 
         return predicted_label
-    def classify_from_frames(self, frames):
-        predictions = []
-
-        for frame in frames:
-            with torch.no_grad():
-                predictions_frame = self.model(frame.unsqueeze(0))
-            probabilities = torch.nn.functional.softmax(predictions_frame[0], dim=0)
-            predicted_class = torch.argmax(probabilities).item()
-            predicted_label = self.class_names[predicted_class]
-            predictions.append(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
@@ -217,21 +216,18 @@ class WomenHairColorClassifier:
 
         return predicted_label
     
-    def classify_from_frames(self, frames):
-        predictions = []
-
-        for frame in frames:
-            with torch.no_grad():
-                predictions_frame = self.model(frame.unsqueeze(0))
-            probabilities = torch.nn.functional.softmax(predictions_frame[0], dim=0)
-            predicted_class = torch.argmax(probabilities).item()
-            predicted_label = self.class_names[predicted_class]
-            predictions.append(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
+    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:
@@ -269,21 +265,18 @@ class BeardClassifier:
         predicted_label = self.class_names[predicted_class]
         return predicted_label
 
-    def classify_from_frames(self, frames):
-        predictions = []
-
-        for frame in frames:
-            with torch.no_grad():
-                predictions_frame = self.model(frame.unsqueeze(0))
-            probabilities = torch.nn.functional.softmax(predictions_frame[0], dim=0)
-            predicted_class = torch.argmax(probabilities).item()
-            predicted_label = self.class_names[predicted_class]
-            predictions.append(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
+    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
@@ -322,21 +315,18 @@ class BeardColorClassifier:
         predicted_label = self.class_names[predicted_class]
         return predicted_label
 
-    def classify_from_frames(self, frames):
-        predictions = []
-
-        for frame in frames:
-            with torch.no_grad():
-                predictions_frame = self.model(frame.unsqueeze(0))
-            probabilities = torch.nn.functional.softmax(predictions_frame[0], dim=0)
-            predicted_class = torch.argmax(probabilities).item()
-            predicted_label = self.class_names[predicted_class]
-            predictions.append(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
+    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
 
 
 
@@ -376,22 +366,18 @@ class HairStyleClassifier:
         predicted_label = self.class_names[predicted_class]
         return predicted_label
 
-    def classify_from_frames(self, frames):
-        predictions = []
-
-        for frame in frames:
-            with torch.no_grad():
-                predictions_frame = self.model(frame.unsqueeze(0))
-            probabilities = torch.nn.functional.softmax(predictions_frame[0], dim=0)
-            predicted_class = torch.argmax(probabilities).item()
-            predicted_label = self.class_names[predicted_class]
-            predictions.append(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
-
+    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):
@@ -428,21 +414,18 @@ class MenHairColorClassifier:
         predicted_label = self.class_names[predicted_class]
         return predicted_label
     
-    def classify_from_frames(self, frames):
-        predictions = []
-
-        for frame in frames:
-            with torch.no_grad():
-                predictions_frame = self.model(frame.unsqueeze(0))
-            probabilities = torch.nn.functional.softmax(predictions_frame[0], dim=0)
-            predicted_class = torch.argmax(probabilities).item()
-            predicted_label = self.class_names[predicted_class]
-            predictions.append(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
+    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
 
 
 
@@ -645,32 +628,120 @@ def generate_funko_figurines(input_image):
 
 def Igenerate_funko_figurines():
     
-    # Capture video from the webcam for 7 seconds
-    captured_duration = 5  # Duration in seconds
+    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
-    #frames = []
-    #start_time = time.time()
-    frames = capture_frame_from_webcam(duration=5)
-    # Continuously capture frames for the specified duration
-    #while (time.time() - start_time) < captured_duration:
-    #    frames.extend(capture_frame_from_webcam())
 
+
+    # 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'])
-    predicted_gender = gender_classifier.classify_from_frames(frames)
 
+    # Detect and classify beard style
+    beard_classifier = BeardClassifier('Data/FunkoSavedModels/FunkoResnet18BeardStyle.pt', ['Bandholz', 'CleanShave', 'FullGoatee', 'Moustache', 'RapIndustryStandards', 'ShortBeard'])
+    
+    # Detect and classify beard color
+    beard_color_classifier = BeardColorClassifier('Data/FunkoSavedModels/FunkoResnet18BeardColor.pt', ['Black', 'DarkBrown', 'Ginger', 'LightBrown', 'SaltAndPepper', 'White'])
 
-    # Process background images and apply beard style and color along with hair style and color
-    final_images = []
+    # Classify hairstyle
+    hair_style_classifier = HairStyleClassifier('Data/FunkoSavedModels/FunkoResnet18HairStyle.pt', ['Afro', 'Bald', 'Puff', 'Spike'])
+    
+    #classify menHairColor
+    menhair_color_classifier = MenHairColorClassifier('Data/FunkoSavedModels/FunkoResnet18MenHairColor.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 background_image_paths in background_image_paths:
-        background_image = Image.open(background_image_paths)
+        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")
@@ -682,21 +753,7 @@ def Igenerate_funko_figurines():
             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/FunkoResnet18BeardStyle.pt', ['Bandholz', 'CleanShave', 'FullGoatee', 'Moustache', 'RapIndustryStandards', 'ShortBeard'])
-            predicted_style_label = beard_classifier.classify_from_frames(frames)
 
-            # Detect and classify beard color
-            beard_color_classifier = BeardColorClassifier('Data/FunkoSavedModels/FunkoResnet18BeardColor.pt', ['Black', 'DarkBrown', 'Ginger', 'LightBrown', 'SaltAndPepper', 'White'])
-            predicted_color_label = beard_color_classifier.classify_from_frames(frames)
-
-            # Classify hairstyle
-            hair_style_classifier = HairStyleClassifier('Data/FunkoSavedModels/FunkoResnet18HairStyle.pt', ['Afro', 'Bald', 'Puff', 'Spike'])
-            predicted_hairStyle_label = hair_style_classifier.classify_from_frames(frames)
-
-            #classify menHairColor
-            menhair_color_classifier = MenHairColorClassifier('Data/FunkoSavedModels/FunkoResnet18MenHairColor.pt', ['Black', 'DarkBrown', 'Ginger', 'LightBrown', 'SaltAndPepper', 'White'])
-            predicted_menhairColor_label = menhair_color_classifier.classify_from_frames(frames)
             if predicted_style_label == 'Bandholz':
                 process_image_Beard(background_image, 320,
                                      f"Data/AdobeColorFunko/Beard/Bandholz/{predicted_color_label}.png",
@@ -730,24 +787,24 @@ def Igenerate_funko_figurines():
             # Add other conditions for different beard styles
 
             # Overlay hairstyle
-            if predicted_hairStyle_label == 'Afro':
+            if predicted_hairstyle_label == 'Afro':
                 process_image_menHair(background_image, 336, 420,
-                                       f"Data/AdobeColorFunko/MenHairstyle/Afro/{predicted_menhairColor_label}.png",
+                                       f"Data/AdobeColorFunko/MenHairstyle/Afro/{predicted_menhaircolor_label}.png",
                                        41, 76)
 
-            if predicted_hairStyle_label == 'Puff':
+            if predicted_hairstyle_label == 'Puff':
                 process_image_menHair(background_image, 305, 420,
-                                       f"Data/AdobeColorFunko/MenHairstyle/Puff/{predicted_menhairColor_label}.png",
+                                       f"Data/AdobeColorFunko/MenHairstyle/Puff/{predicted_menhaircolor_label}.png",
                                        56, 68)
 
-            if predicted_hairStyle_label == 'Spike':
+            if predicted_hairstyle_label == 'Spike':
                 process_image_menHair(background_image, 310, 420,
-                                       f"Data/AdobeColorFunko/MenHairstyle/Spike/{predicted_menhairColor_label}.png",
+                                       f"Data/AdobeColorFunko/MenHairstyle/Spike/{predicted_menhaircolor_label}.png",
                                        52, 70)
 
-            if predicted_hairStyle_label == 'Bald':
+            if predicted_hairstyle_label == 'Bald':
                 process_image_menHair(background_image, 310, 420,
-                                       f"Data/AdobeColorFunko/MenHairstyle/Bald/{predicted_menhairColor_label}.png",
+                                       f"Data/AdobeColorFunko/MenHairstyle/Bald/{predicted_menhaircolor_label}.png",
                                        67, 120)
 
 
@@ -758,29 +815,25 @@ def Igenerate_funko_figurines():
             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_from_frames(frames)
-
-            WomenHairColor_classifier = WomenHairColorClassifier('Data/FunkoSavedModels/WomenHairColor.pt', ['Black', 'Brown', 'Ginger', 'White'])
-            predicted_WomenHairColor = WomenHairColor_classifier.classify_from_frames(frames)
-            if predicted_WomenHairStyle == 'MediumLength':
+            if predicted_women_hairstyle == 'MediumLength':
                 process_image_WomanHair(background_image, 300,460,
-                                     f"Data/AdobeColorFunko/WomenHairstyle/MediumLength/{predicted_WomenHairColor}.png",
+                                     f"Data/AdobeColorFunko/WomenHairstyle/MediumLength/{predicted_women_haircolor}.png",
                                      56, 50)
 
-            if predicted_WomenHairStyle == 'ShortHair':
+            if predicted_women_hairstyle == 'ShortHair':
                 process_image_WomanHair(background_image, 270,460,
-                                     f"Data/AdobeColorFunko/WomenHairstyle/ShortHair/{predicted_WomenHairColor}.png",
+                                     f"Data/AdobeColorFunko/WomenHairstyle/ShortHair/{predicted_women_haircolor}.png",
                                      61, 49)
 
-            if predicted_WomenHairStyle == 'SidePlait':
+            if predicted_women_hairstyle == 'SidePlait':
                 process_image_WomanHair(background_image, 300,450,
-                                     f"Data/AdobeColorFunko/WomenHairstyle/SidePlait/{predicted_WomenHairColor}.png",
+                                     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)
@@ -795,7 +848,7 @@ with gr.Blocks() as demo:
     # Funko POP! Figure Creation
     Enabling Streamlined Automation with Artificial Intelligence 
     """)
-    imageComponent = gr.Image(type="filepath", height=200, width=200)
+    imageComponent = gr.Image(type="filepath", height=250, width=250)
     with gr.Row():
         MyOutputs = [gr.Image(type="pil", label="Generated Image " + str(i + 1), height=300, width=200) for i in range(3)]
     submitButton = gr.Button(value="Submit")