Update app.py

app.py

@@ -232,7 +232,7 @@ class WomenHairColorClassifier:
 # Function to classify beard style
 class BeardClassifier:
     def __init__(self, model_path, class_names):
-        self.model = torch.hub.load('pytorch/vision', 'resnet18', pretrained=False)
+        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)
@@ -282,7 +282,7 @@ class BeardClassifier:
 # Function to classify beard color
 class BeardColorClassifier:
     def __init__(self, model_path, class_names):
-        self.model = torch.hub.load('pytorch/vision', 'resnet18', pretrained=False)
+        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)
@@ -333,7 +333,7 @@ class BeardColorClassifier:
 # Function to classify hairstyle
 class HairStyleClassifier:
     def __init__(self, model_path, class_names):
-        self.model = torch.hub.load('pytorch/vision', 'resnet18', pretrained=False)
+        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)
@@ -381,7 +381,7 @@ class HairStyleClassifier:
 
 class MenHairColorClassifier:
     def __init__(self, model_path, class_names):
-        self.model = torch.hub.load('pytorch/vision', 'resnet18', pretrained=False)
+        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)
@@ -521,19 +521,19 @@ def generate_funko_figurines(input_image):
             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'])
+            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/FunkoResnet18BeardColor.pt', ['Black', 'DarkBrown', 'Ginger', 'LightBrown', 'SaltAndPepper', 'White'])
+            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/FunkoResnet18HairStyle.pt', ['Afro', 'Bald', 'Puff', 'Spike'])
+            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/FunkoResnet18MenHairColor.pt', ['Black', 'DarkBrown', 'Ginger', 'LightBrown', 'SaltAndPepper', 'White'])
+            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':