Update app.py

app.py

@@ -65,7 +65,22 @@ class GenderClassifier:
         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
 
     
 class WomenHairStyleClassifier:
@@ -105,6 +120,22 @@ 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)
+
+        # 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):
@@ -144,6 +175,22 @@ 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
+
 # Function to classify beard style
 class BeardClassifier:
     def __init__(self, model_path, class_names):
@@ -180,6 +227,22 @@ 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
+
 
 # Function to classify beard color
 class BeardColorClassifier:
@@ -217,6 +280,22 @@ 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
+
 
 
 # Function to classify hairstyle
@@ -255,6 +334,21 @@ 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
 
 
 class MenHairColorClassifier:
@@ -292,299 +386,21 @@ class MenHairColorClassifier:
         predicted_label = self.class_names[predicted_class]
         return predicted_label
     
-class IGenderClassifier:
-    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 classify_from_frames(self, frames):
+        predictions = []
 
-    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')))
+        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_gender(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 IWomenHairStyleClassifier:
-    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, 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 IWomenHairColorClassifier:
-    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, 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 IBeardClassifier:
-    def __init__(self, model_path, class_names):
-        self.model = torch.hub.load('pytorch/vision', 'resnet18', 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, 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 IBeardColorClassifier:
-    def __init__(self, model_path, class_names):
-        self.model = torch.hub.load('pytorch/vision', 'resnet18', 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, 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 IHairStyleClassifier:
-    def __init__(self, model_path, class_names):
-        self.model = torch.hub.load('pytorch/vision', 'resnet18', 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, 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 to classify Mens haircolor
-class IMenHairColorClassifier:
-    def __init__(self, model_path, class_names):
-        self.model = torch.hub.load('pytorch/vision', 'resnet18', 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, 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
 
 
 
@@ -842,8 +658,8 @@ def Igenerate_funko_figurines():
     #    frames.extend(capture_frame_from_webcam())
 
     # Detect and classify gender
-    gender_classifier = IGenderClassifier('Data/FunkoSavedModels/Gender.pt', ['Female', 'Male'])
-    predicted_gender = gender_classifier.classify_gender(frames)
+    gender_classifier = GenderClassifier('Data/FunkoSavedModels/Gender.pt', ['Female', 'Male'])
+    predicted_gender = gender_classifier.classify_from_frames(frames)
 
 
     # Process background images and apply beard style and color along with hair style and color
@@ -869,19 +685,19 @@ def Igenerate_funko_figurines():
             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_beard(frames)
+            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_beard_color(frames)
+            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_hair(frames)
+            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_menHair_color(frames)
+            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",
@@ -944,10 +760,10 @@ def Igenerate_funko_figurines():
             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(frames)
+            predicted_WomenHairStyle = WomenHairStyle_classifier.classify_from_frames(frames)
 
             WomenHairColor_classifier = WomenHairColorClassifier('Data/FunkoSavedModels/WomenHairColor.pt', ['Black', 'Brown', 'Ginger', 'White'])
-            predicted_WomenHairColor = WomenHairColor_classifier.classify_hairColor(frames)
+            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_WomenHairColor}.png",