better model

app.py

@@ -13,13 +13,23 @@ def load_class_names(file_path):
         class_names = [line.strip() for line in f.readlines()]
     return class_names
 # Function to load the model from a .pkl file
-def load_model(model_path):
+def load_model(model_path, model_type='resnet'):
     # Load the model state dictionary
     model_state_dict = torch.load(model_path, map_location=torch.device('cpu'))
-    
-    # Create an instance of the model
-    model = models.mobilenet_v2(pretrained=False)
-    model.classifier[1] = nn.Linear(model.last_channel, num_classes)  # Adjust to your number of classes
+
+    # Create an instance of the model based on model_type
+    if model_type == 'mobilenet':
+        model = models.mobilenet_v2(pretrained=False)
+        model.classifier[1] = nn.Linear(model.last_channel, num_classes)
+    elif model_type == 'resnet':
+        model = models.resnet50(pretrained=False)
+        model.fc = nn.Linear(model.fc.in_features, num_classes)
+    elif model_type == 'densenet':
+        model = models.densenet121(pretrained=False)
+        model.classifier = nn.Linear(model.classifier.in_features, num_classes)
+    else:
+        raise ValueError(f"Unsupported model type: {model_type}")
+
     # Load the state dictionary into the model
     model.load_state_dict(model_state_dict)
     # Set the model to evaluation mode
@@ -51,8 +61,8 @@ def predict_image(image):
     confidence_score = confidence.item() * 100
     return f"{class_name}: {confidence_score:.2f}%"
 
-# Load your trained model and class names
-model_path = 'mobilenet_model.pkl'
+# Load trained model and class names
+model_path = 'resnet30EpochsPretrainedNFeatureX_model.pkl'
 class_file_path = 'classes.txt'
 class_names = load_class_names(class_file_path)
 num_classes = len(class_names)

densenet30EpochsPretrainedNFeatureX_model.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:66c1769d7f5c4415d6f6f7827b74b8730036271d759f98b0a8eb489849434d11
+size 28813477

mobilenet30EpochsPretrainedNFeatureX_model.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:506ae238566267e277c4d952cacf6da4b87f827e7f30c7302b423809a5c0047c
+size 9612161

resnet30EpochsPretrainedNFeatureX_model.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:edf0def0a33c6cd61025b08ee9a6256238bc26c0871c2488a7519b46c58918b3
+size 95097345