Add support for UNet++ with EfficientNet-B3 and implement model caching

app.py

@@ -217,6 +217,11 @@ MODEL_OPTIONS = {
         "path": "./model/unet_fibril_seg_model.pth",
         "encoder": "resnet34",
         "architecture": "Unet"
+    },
+    "UNet++ (efficientnet-b3)": {
+        "path": "./model/encoder_efficientnet-b3_decoder_UnetPlusPlus_fibril_seg_model.pth",
+        "encoder": "efficientnet-b3",
+        "architecture": "UnetPlusPlus"
     }
 }
 
@@ -231,7 +236,39 @@ def get_transform(size):
 transform = get_transform(512)
 
 # ─── Model Loader ──────────────────────────────────────────
+# def load_model(model_name):
+#     config = MODEL_OPTIONS[model_name]
+#     if config["architecture"] == "UnetPlusPlus":
+#         model = smp.UnetPlusPlus(
+#             encoder_name=config["encoder"],
+#             encoder_weights="imagenet",
+#             decoder_channels=(256, 128, 64, 32, 16),
+#             in_channels=1,
+#             classes=1,
+#             activation=None
+#         )
+#     elif config["architecture"] == "Unet":
+#         model = smp.Unet(
+#             encoder_name=config["encoder"],
+#             encoder_weights="imagenet",
+#             decoder_channels=(256, 128, 64, 32, 16),
+#             in_channels=1,
+#             classes=1,
+#             activation=None
+#         )
+#     else:
+#         raise ValueError(f"Unsupported architecture: {config['architecture']}")
+
+#     model.load_state_dict(torch.load(config["path"], map_location=device))
+#     model.eval()
+#     return model.to(device)
+
+model_cache = {}
+
 def load_model(model_name):
+    if model_name in model_cache:
+        return model_cache[model_name]
+
     config = MODEL_OPTIONS[model_name]
     if config["architecture"] == "UnetPlusPlus":
         model = smp.UnetPlusPlus(
@@ -256,7 +293,8 @@ def load_model(model_name):
 
     model.load_state_dict(torch.load(config["path"], map_location=device))
     model.eval()
-    return model.to(device)
+    model_cache[model_name] = model.to(device)
+    return model_cache[model_name]
 
 # ─── Prediction Function ───────────────────────────────────
 def predict(image, model_name):

model/encoder_efficientnet-b3_decoder_UnetPlusPlus_fibril_seg_model.pth

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