lakshmi082024 commited on
Commit
8567cb5
·
verified ·
1 Parent(s): 23a11ce

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +8 -18
app.py CHANGED
@@ -1,14 +1,3 @@
1
- import os
2
- import sys
3
- import subprocess
4
-
5
- # ⬇️ Install `segment_anything` if not already installed
6
- try:
7
- from segment_anything import SamPredictor, sam_model_registry
8
- except ImportError:
9
- subprocess.check_call([sys.executable, "-m", "pip", "install", "git+https://github.com/facebookresearch/segment-anything.git"])
10
- from segment_anything import SamPredictor, sam_model_registry
11
-
12
  import streamlit as st
13
  import cv2
14
  import numpy as np
@@ -16,16 +5,17 @@ import pandas as pd
16
  from PIL import Image
17
  import torch
18
  from torchvision.transforms import Compose, Resize, ToTensor, Normalize
 
19
 
20
  # Set Streamlit configuration
21
  st.set_page_config(page_title="Volume Estimator", layout="wide")
22
- st.title("📦 Volume Estimation using SAM Segmentation + MiDaS Depth")
23
 
24
  # Load SAM and MiDaS models
25
  @st.cache_resource
26
  def load_models():
27
- sam_checkpoint = "https://drive.google.com/file/d/1pgv2kKrytyY2_uHSypJh5yQz2BRlfjSS/view?usp=drive_link/sam_vit_b_01ec64.pth" # Make sure this file is present in the same folder
28
- sam = sam_model_registry["vit_b"](checkpoint=sam_checkpoint).to("cuda" if torch.cuda.is_available() else "cpu")
29
  predictor = SamPredictor(sam)
30
 
31
  midas = torch.hub.load("intel-isl/MiDaS", "DPT_Large")
@@ -65,7 +55,7 @@ elif source_option == "Use Webcam":
65
  frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
66
  stframe.image(frame_rgb, caption="Live Camera Feed", channels="RGB")
67
 
68
- if st.button("📸 Capture Frame"):
69
  image_pil = Image.fromarray(frame_rgb)
70
  run_camera = False
71
  cap.release()
@@ -127,10 +117,10 @@ if image_pil:
127
  # Display volume table
128
  if volume_data:
129
  df = pd.DataFrame(volume_data)
130
- st.markdown("### 📊 Object Dimensions and Volume")
131
  st.dataframe(df)
132
 
133
  csv = df.to_csv(index=False).encode('utf-8')
134
- st.download_button("📂 Download Volume Table as CSV", csv, "object_volumes_with_units.csv", "text/csv")
135
  else:
136
- st.warning("🚫 No objects were segmented.")
 
 
 
 
 
 
 
 
 
 
 
 
1
  import streamlit as st
2
  import cv2
3
  import numpy as np
 
5
  from PIL import Image
6
  import torch
7
  from torchvision.transforms import Compose, Resize, ToTensor, Normalize
8
+ from segment_anything import SamPredictor, sam_model_registry
9
 
10
  # Set Streamlit configuration
11
  st.set_page_config(page_title="Volume Estimator", layout="wide")
12
+ st.title("Volume Estimation using SAM Segmentation + MiDaS Depth")
13
 
14
  # Load SAM and MiDaS models
15
  @st.cache_resource
16
  def load_models():
17
+ sam_checkpoint = "https://drive.google.com/file/d/1pgv2kKrytyY2_uHSypJh5yQz2BRlfjSS/view?usp=drive_link/sam_vit_h_4b8939.pth"
18
+ sam = sam_model_registry["vit_h"](checkpoint=sam_checkpoint).to("cuda" if torch.cuda.is_available() else "cpu")
19
  predictor = SamPredictor(sam)
20
 
21
  midas = torch.hub.load("intel-isl/MiDaS", "DPT_Large")
 
55
  frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
56
  stframe.image(frame_rgb, caption="Live Camera Feed", channels="RGB")
57
 
58
+ if st.button("Capture Frame"):
59
  image_pil = Image.fromarray(frame_rgb)
60
  run_camera = False
61
  cap.release()
 
117
  # Display volume table
118
  if volume_data:
119
  df = pd.DataFrame(volume_data)
120
+ st.markdown("### Object Dimensions and Volume")
121
  st.dataframe(df)
122
 
123
  csv = df.to_csv(index=False).encode('utf-8')
124
+ st.download_button("Download Volume Table as CSV", csv, "object_volumes_with_units.csv", "text/csv")
125
  else:
126
+ st.warning("No objects were segmented.")