Update app.py
Browse files
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("
|
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/
|
28 |
-
sam = sam_model_registry["
|
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("
|
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("###
|
131 |
st.dataframe(df)
|
132 |
|
133 |
csv = df.to_csv(index=False).encode('utf-8')
|
134 |
-
st.download_button("
|
135 |
else:
|
136 |
-
st.warning("
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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.")
|