Update app.py

app.py

@@ -6,18 +6,29 @@ from PIL import Image
 import torch
 from torchvision.transforms import Compose, Resize, ToTensor, Normalize
 from segment_anything import SamPredictor, sam_model_registry
+import requests
+import os
 
-# Set Streamlit configuration
+# Streamlit configuration
 st.set_page_config(page_title="Volume Estimator", layout="wide")
 st.title("Volume Estimation using SAM Segmentation + MiDaS Depth")
 
 # Load SAM and MiDaS models
 @st.cache_resource
 def load_models():
-    sam_checkpoint = "https://drive.google.com/file/d/1pgv2kKrytyY2_uHSypJh5yQz2BRlfjSS/view?usp=drive_link/sam_vit_h_4b8939.pth"
-    sam = sam_model_registry["vit_h"](checkpoint=sam_checkpoint).to("cuda" if torch.cuda.is_available() else "cpu")
+    # Download SAM checkpoint from Hugging Face
+    checkpoint_url = "https://huggingface.co/HCMUE-Research/SAM-vit-h/resolve/main/sam_vit_h_4b8939.pth"
+    checkpoint_path = "sam_vit_h_4b8939.pth"
+    if not os.path.exists(checkpoint_path):
+        with open(checkpoint_path, "wb") as f:
+            f.write(requests.get(checkpoint_url).content)
+
+    # Load SAM
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    sam = sam_model_registry["vit_h"](checkpoint=checkpoint_path).to(device)
     predictor = SamPredictor(sam)
 
+    # Load MiDaS
     midas = torch.hub.load("intel-isl/MiDaS", "DPT_Large")
     midas.eval()
     midas_transform = Compose([
@@ -46,7 +57,6 @@ elif source_option == "Use Webcam":
     if run_camera:
         cap = cv2.VideoCapture(0)
         stframe = st.empty()
-        capture = False
 
         while run_camera and cap.isOpened():
             ret, frame = cap.read()
@@ -57,11 +67,10 @@ elif source_option == "Use Webcam":
 
             if st.button("Capture Frame"):
                 image_pil = Image.fromarray(frame_rgb)
-                run_camera = False
                 cap.release()
                 break
 
-# Continue processing if we have an image
+# Continue only if an image is available
 if image_pil:
     image_np = np.array(image_pil)
     img_h, img_w = image_np.shape[:2]
@@ -85,7 +94,7 @@ if image_pil:
         depth_prediction = midas_model(input_tensor).squeeze().cpu().numpy()
     depth_resized = cv2.resize(depth_prediction, (img_w, img_h))
 
-    # Object volume computation
+    # Compute object volumes
     volume_data = []
     for i, mask in enumerate(masks):
         mask_np = mask
@@ -97,7 +106,6 @@ if image_pil:
         height_cm = height_px * pixel_to_cm_y
 
         depth_masked = depth_resized[mask_np > 0.5]
-
         if depth_masked.size == 0:
             continue
 
@@ -114,7 +122,7 @@ if image_pil:
             "Volume": f"{volume_cm3} cm³"
         })
 
-    # Display volume table
+    # Display volume results
     if volume_data:
         df = pd.DataFrame(volume_data)
         st.markdown("### Object Dimensions and Volume")