SAM model

app.py

@@ -1,50 +1,61 @@
+import torch
+from transformers import SamConfig, SamProcessor, SamModel
 from shiny import App, render, ui
-from shiny.ui import output_image, input_file, panel_main, panel_sidebar, layout_sidebar
-import matplotlib.pyplot as plt
 import numpy as np
-import torch  # Assuming the model is a PyTorch model
 from PIL import Image
 import io
+import matplotlib.pyplot as plt
+
+# Load the model configuration
+model_config = SamConfig.from_pretrained("facebook/sam-vit-base")
+processor = SamProcessor.from_pretrained("facebook/sam-vit-base")
+
+# Create an instance of the model architecture with the loaded configuration
+my_model = SamModel(config=model_config)
+
+# Update the model by loading the weights from a saved file
+my_model.load_state_dict(torch.load("models/model_checkpoint_trained_on_train.pth"))
 
-# Load your pre-trained model
-model = torch.load('models/model_checkpoint_trained_on_train.pth')
-model.eval()
+# Set model to evaluation mode
+my_model.eval()
 
 def predict(image_file):
-    # Open image
+    # Open and preprocess the image
     image = Image.open(image_file)
-    # Process image as per model requirements (resize, normalize, etc.)
-    # Convert to tensor, e.g., if model expects tensor input
-    image_tensor = torch.Tensor(np.array(image) / 255).unsqueeze(0)  # adjust processing as needed
+    image_array = np.array(image)
+
+    # Process the image
+    inputs = processor(images=image_array, return_tensors="pt")
+
+    # Make a prediction with the model
     with torch.no_grad():
-        output = model(image_tensor)
-    # Post-process output, e.g., binarize, segment
-    segmentation = (output > 0.5).squeeze().numpy()  # adjust as needed
+        outputs = my_model(**inputs)
+
+    # Extract the mask or segmentation map
+    segmentation = outputs[0].squeeze().numpy()  # Adjust to extract necessary outputs
+
     return segmentation
 
 def server(input, output, session):
     @output
     @render.image
     def segmented_image():
-        # Check if a file has been uploaded
-        if not input.image_file:
+        if not input.image_file():
             return None
 
-        # Predict segmentation
+        # Get the uploaded file
         uploaded_file = input.image_file()[0]
         segmentation = predict(uploaded_file)
 
-        # Create figure to show segmentation
+        # Visualize the segmentation
         plt.imshow(segmentation, cmap='gray')
         plt.axis('off')
 
-        # Save to buffer and return
         buf = io.BytesIO()
         plt.savefig(buf, format='png')
         buf.seek(0)
         return buf
 
-print("Creating app UI")
 app_ui = ui.page_fluid(
     layout_sidebar(
         panel_sidebar(
@@ -55,6 +66,8 @@ app_ui = ui.page_fluid(
         ),
     )
 )
-print("Creating app")
+
 app = App(app_ui, server)
 
+if __name__ == "__main__":
+    app.run()

requirements.txt

@@ -3,3 +3,9 @@ torch
 numpy
 pillow
 matplotlib
+segment-geospatial 
+groundingdino-py 
+leafmap 
+localtileserver 
+datasets 
+transformers