Update app.py

app.py

@@ -16,7 +16,6 @@ import pandas as pd
 from PIL import Image
 import torch
 from torchvision.transforms import Compose, Resize, ToTensor, Normalize
-from segment_anything import SamPredictor, sam_model_registry
 
 # Set Streamlit configuration
 st.set_page_config(page_title="Volume Estimator", layout="wide")
@@ -25,8 +24,8 @@ st.title("📦 Volume Estimation using SAM Segmentation + MiDaS Depth")
 # Load SAM and MiDaS models
 @st.cache_resource
 def load_models():
-    sam_checkpoint = "sam_vit_b_01ec64.pth"
-    sam = sam_model_registry["vit_h"](checkpoint=sam_checkpoint).to("cuda" if torch.cuda.is_available() else "cpu")
+    sam_checkpoint = "sam_vit_b_01ec64.pth"  # Make sure this file is present in the same folder
+    sam = sam_model_registry["vit_b"](checkpoint=sam_checkpoint).to("cuda" if torch.cuda.is_available() else "cpu")
     predictor = SamPredictor(sam)
 
     midas = torch.hub.load("intel-isl/MiDaS", "DPT_Large")