Update app.py

app.py

@@ -124,7 +124,6 @@
 # if __name__ == "__main__":
 #     demo.launch(share=True, height=900)
 
-
 import gradio as gr
 import rasterio
 import os
@@ -135,10 +134,10 @@ from SemanticModel.model_core import SegmentationModel
 from SemanticModel.prediction import PredictionPipeline
 
 from segmentation_models_pytorch.decoders.unet.model import Unet
-
-# Add this before loading your model
 from segmentation_models_pytorch.encoders.timm_regnet import RegNetEncoder
-torch.serialization.add_safe_globals([Unet, RegNetEncoder])
+
+# Previously we added safe globals globally, now we'll use the context manager approach.
+# torch.serialization.add_safe_globals([Unet, RegNetEncoder])
 
 COLORS = [
    (0, 0, 0),      # Background (black) 
@@ -180,12 +179,14 @@ def predict_and_show(input_file):
             img = src.read()
             img = np.moveaxis(img, 0, 2)[:,:,:3]
             img_normalized = np.clip(img/img.max(), 0, 1)
-        model = SegmentationModel(
-           classes=['bg', 'cacao', 'matarraton', 'abarco'],
-           architecture='unet',
-           encoder='timm-regnety_120',
-            weights='models/CacaoShadeClassification.pth'
-       )
+        # Use the safe_globals context manager when loading the model
+        with torch.serialization.safe_globals([RegNetEncoder]):
+            model = SegmentationModel(
+               classes=['bg', 'cacao', 'matarraton', 'abarco'],
+               architecture='unet',
+               encoder='timm-regnety_120',
+               weights='models/CacaoShadeClassification.pth'
+           )
         predictor = PredictionPipeline(model)
         pred, _ = predictor.predict_raster(input_file.name, tile_size=1024)
         colored_pred = style_prediction(pred)
@@ -273,3 +274,4 @@ with gr.Blocks(theme=gr.themes.Default()) as demo:
 if __name__ == "__main__":
     demo.launch(share=True, height=900)
 
+