new models

app.py

@@ -19,7 +19,7 @@ def load_class_info(file_path):
     return class_info
 
 # Function to load the model from a .pkl file
-def load_model(model_path, model_type='densnet121'):
+def load_model(model_path, model_type):
     # Load the model state dictionary
     model_state_dict = torch.load(model_path, map_location=torch.device('cpu'))
 
@@ -50,7 +50,15 @@ val_transform = transforms.Compose([
 ])
 
 # Define the prediction function
-def predict_image(image):
+def predict_image(image, model_choice):
+    global model, current_model
+
+    # Load the selected model if it's not already loaded
+    if model_choice != current_model:
+        model_path = model_paths[model_choice]
+        model = load_model(model_path, model_choice)
+        current_model = model_choice
+    
     # Convert the NumPy array to a PIL Image
     if isinstance(image, np.ndarray):
         image = Image.fromarray(image.astype('uint8'), 'RGB')
@@ -76,19 +84,29 @@ def predict_image(image):
     
     return result, html_result
 
-# Load trained model and class names
-model_path = 'densenet121_15EpochsPretrainedNoExtractionNoLR_model'
+# Load class names and class info
 class_file_path = 'classes.txt'
 class_info_path = 'classinfo.txt'
 class_names = load_class_names(class_file_path)
 class_info = load_class_info(class_info_path)
 num_classes = len(class_names)
-model = load_model(model_path)
+
+# Define model paths
+model_paths = {
+    'densenet121': 'densenet121_15EpochsPretrainedNoExtractionNoLR_model.pkl',
+    'resnet18': 'resnet18_25EpochsPretrainedExtractionNoLR_model.pkl',
+    'mobilenetv2': 'mobilenetv2_25EpochsPretrainedExtractionNoLR_model.pkl'
+}
+
+# Set default model
+current_model = 'densenet121'
+model = load_model(model_paths[current_model], current_model)
 
 # Create the Gradio interface
 iface = gr.Interface(
     fn=predict_image, 
-    inputs=gr.Image(height=500),
+    inputs=[gr.Image(height=500),
+            gr.Dropdown(choices=["densenet121", "resnet18", "mobilenetv2"], value="densenet121", label="Select Model")],
     outputs=[gr.Label(num_top_classes=1), gr.HTML()],
     title="Image Classification",
     description="Upload an image to get the predicted label",

mobilenetv2_25EpochsPretrainedExtractionNoLR_model.pkl

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

resnet18_25EpochsPretrainedExtractionNoLR_model.pkl

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