Update app.py

app.py

@@ -828,143 +828,34 @@ def generate_funko_figurines(input_image):
 
     return final_images
 
-def Igenerate_funko_figurines(image_input):
+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
+    captured_duration = 5  # Duration in seconds
 
-    print("Capturing frames")
-    frames = capture_frame_from_webcam(duration=7)
-    print("Frames captured")
+    # 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 = IWomenHairStyleClassifier('Data/FunkoSavedModels/WomenHairStyle.pt', ['MediumLength', 'ShortHair', 'SidePlait'])
-    
-    # Classify women hair color
-    women_hair_color_classifier = IWomenHairColorClassifier('Data/FunkoSavedModels/WomenHairColor.pt', ['Black', 'Brown', 'Ginger', 'White'])
-    
     # Detect and classify gender
-    gender_classifier = IGenderClassifier('Data/FunkoSavedModels/Gender.pt', ['Female', 'Male'])
-
-    # Detect and classify beard style
-    beard_classifier = IBeardClassifier('Data/FunkoSavedModels/FunkoResnet18BeardStyle.pt', ['Bandholz', 'CleanShave', 'FullGoatee', 'Moustache', 'RapIndustryStandards', 'ShortBeard'])
-    
-    # Detect and classify beard color
-    beard_color_classifier = IBeardColorClassifier('Data/FunkoSavedModels/FunkoResnet18BeardColor.pt', ['Black', 'DarkBrown', 'Ginger', 'LightBrown', 'SaltAndPepper', 'White'])
-
-    # Classify hairstyle
-    hair_style_classifier = IHairStyleClassifier('Data/FunkoSavedModels/FunkoResnet18HairStyle.pt', ['Afro', 'Bald', 'Puff', 'Spike'])
-    
-    #classify menHairColor
-    menhair_color_classifier = IMenHairColorClassifier('Data/FunkoSavedModels/FunkoResnet18MenHairColor.pt', ['Black', 'DarkBrown', 'Ginger', 'LightBrown', 'SaltAndPepper', 'White'])
-    
-    def predict_male_features_from_frames(frame):
-        return  [
-            beard_classifier.classify_beard(image=frame,image_type=False),
-            beard_color_classifier.classify_beard_color(image=frame,image_type=False),
-            hair_style_classifier.classify_hair(image=frame,image_type=False),
-            menhair_color_classifier.classify_menHair_color(image=frame,image_type=False)
-        ]
-    
-    def predict_female_features_from_frames(frame):
-        return  [
-            women_hairstyle_classifier.classify_hairStyle(image=frame,image_type=False),
-            women_hair_color_classifier.classify_hairColor(image=frame,image_type=False),
-        ]
-    
-    def predict_gender_from_frames(frame):
-        return gender_classifier.classify_gender(image=frame,image_type=False)
-    
-
-    if image_input == True:
-        
-        predicted_women_hairstyle = women_hairstyle_classifier.classify_hairStyle(image="<input as image>",image_type=image_input)
-        # Predict women hair color
-        predicted_women_haircolor = women_hair_color_classifier.classify_hairColor(image="<input as image>",image_type=image_input)
-        # Predict Gender
-        predicted_gender = gender_classifier.classify_gender(image="<input as image>",image_type=image_input)
-        # Predict beard style
-        predicted_style_label = beard_classifier.classify_beard(image="<input as image>",image_type=image_input)
-        # Predict beard color
-        predicted_color_label = beard_color_classifier.classify_beard_color(image="<input as image>",image_type=image_input)
-        # Classify hairstyle
-        predicted_hairstyle_label = hair_style_classifier.classify_hair(image="<input as image>",image_type=image_input)
-        #classify menHairColor
-        predicted_menhaircolor_label = menhair_color_classifier.classify_menHair_color(image="<input as image>",image_type=image_input)
-        
-        if predicted_gender == 'Male':
-            background_image_paths = male_background_image_paths
-        if predicted_gender == 'Female':
-            background_image_paths = female_background_image_paths
-
-    else:
-        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')
+    gender_classifier = GenderClassifier('Data/FunkoSavedModels/Gender.pt', ['Female', 'Male'])
+    predicted_gender = gender_classifier.classify_from_frames(frames)
 
-        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)
+    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")
@@ -976,7 +867,21 @@ def Igenerate_funko_figurines(image_input):
             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",
@@ -1010,24 +915,24 @@ def Igenerate_funko_figurines(image_input):
             # 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)
 
 
@@ -1038,19 +943,24 @@ def Igenerate_funko_figurines(image_input):
             x_coordinate = 90
             y_coordinate = 50
             dummy_eye(background_image, x, y, placeholder_image_path, x_coordinate, y_coordinate)
-            if predicted_women_hairstyle == 'MediumLength':
+            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':
                 process_image_WomanHair(background_image, 300,460,
-                                     f"Data/AdobeColorFunko/WomenHairstyle/MediumLength/{predicted_women_haircolor}.png",
+                                     f"Data/AdobeColorFunko/WomenHairstyle/MediumLength/{predicted_WomenHairColor}.png",
                                      56, 50)
 
-            if predicted_women_hairstyle == 'ShortHair':
+            if predicted_WomenHairStyle == 'ShortHair':
                 process_image_WomanHair(background_image, 270,460,
-                                     f"Data/AdobeColorFunko/WomenHairstyle/ShortHair/{predicted_women_haircolor}.png",
+                                     f"Data/AdobeColorFunko/WomenHairstyle/ShortHair/{predicted_WomenHairColor}.png",
                                      61, 49)
 
-            if predicted_women_hairstyle == 'SidePlait':
+            if predicted_WomenHairStyle == 'SidePlait':
                 process_image_WomanHair(background_image, 300,450,
-                                     f"Data/AdobeColorFunko/WomenHairstyle/SidePlait/{predicted_women_haircolor}.png",
+                                     f"Data/AdobeColorFunko/WomenHairstyle/SidePlait/{predicted_WomenHairColor}.png",
                                      54, 56)
 
 
@@ -1064,7 +974,6 @@ def Igenerate_funko_figurines(image_input):
 
 
 
-
 with gr.Blocks() as demo:
     gr.Markdown(
     """
@@ -1079,10 +988,5 @@ with gr.Blocks() as demo:
     RecordButton = gr.Button(value="Generate My Custom Funko POP")
     RecordButton.click(Igenerate_funko_figurines, outputs=MyOutputs)
 
-
-
-
-
-
 if __name__ == "__main__":
     demo.launch()