Update app.py

app.py

@@ -1,129 +1,126 @@
-import streamlit as st
-st.set_page_config(page_title="Object Volume Estimator", layout="wide")
-
-import cv2
-import torch
-import numpy as np
-from PIL import Image
-import pandas as pd
-from ultralytics import YOLO
-from torchvision.transforms import Compose, Resize, ToTensor, Normalize
-
-# Load models
-@st.cache_resource
-def load_models():
-    yolo = YOLO("yolov8n.pt")
-    midas = torch.hub.load("intel-isl/MiDaS", "DPT_Large")
-    midas.eval()
-    transform = Compose([
-        Resize(384),
-        ToTensor(),
-        Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
-    ])
-    return yolo, midas, transform
-
-yolo_model, midas_model, midas_transform = load_models()
-
-# --- Streamlit App ---
-st.title("Object Dimension & Volume Estimator")
-
-# 📸 Image Source Selection
-option = st.radio("Choose Image Source:", ("📤 Upload Image", "📸 Use Camera"))
-
-# Mode Selection
-mode = st.selectbox("Select Image Type:", ["2D RGB Image", "RGB + Depth Image"])
-
-image_pil = None
-
-if option == "📤 Upload Image":
-    uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
-    if uploaded_file:
-        image_pil = Image.open(uploaded_file).convert("RGB")
-
-elif option == "📸 Use Camera":
-    camera_input = st.camera_input("Take a picture")
-    if camera_input:
-        image_pil = Image.open(camera_input).convert("RGB")
-
-# Upload depth image if RGB + Depth selected
-depth_map = None
-if mode == "RGB + Depth Image" and image_pil is not None:
-    depth_file = st.file_uploader("Upload corresponding Depth Map (grayscale image)", type=["png", "jpg", "jpeg", "tiff"])
-    if depth_file:
-        depth_map = np.array(Image.open(depth_file).convert("L"))
-
-# Proceed if we have an image
-if image_pil is not None:
-    # Convert to OpenCV format
-    image_cv = np.array(image_pil)
-    image_cv = cv2.cvtColor(image_cv, cv2.COLOR_RGB2BGR)
-    img_h, img_w = image_cv.shape[:2]
-    st.image(image_pil, caption="Input Image", use_container_width=True)
-
-    # Store detected object data
-    object_data = []
-
-    # 🔍 YOLO Object Detection
-    results = yolo_model(image_cv)
-
-    # 📏 Depth Estimation
-    if mode == "2D RGB Image":
-        input_tensor = midas_transform(image_pil).unsqueeze(0)
-        with torch.no_grad():
-            estimated_depth = midas_model(input_tensor).squeeze().cpu().numpy()
-        depth_map_resized = cv2.resize(estimated_depth, (img_w, img_h))
-    else:
-        if depth_map is not None:
-            depth_map_resized = cv2.resize(depth_map, (img_w, img_h))
-        else:
-            st.error("Please upload a corresponding depth map for 3D images.")
-            st.stop()
-
-    # 🎯 Process each detection
-    for r in results:
-        for box in r.boxes:
-            x1, y1, x2, y2 = map(int, box.xyxy[0])
-            width = x2 - x1
-            height = y2 - y1
-            region = depth_map_resized[y1:y2, x1:x2]
-            if region.size == 0:
-                continue
-            depth = np.mean(region)
-            volume = round(depth * width * height, 2)
-
-            # Draw bounding box & label
-            cv2.rectangle(image_cv, (x1, y1), (x2, y2), (0, 255, 0), 2)
-            label = f"L×B×H: {depth:.2f}×{width}×{height} | V: {volume}"
-            cv2.putText(image_cv, label, (x1, y1 - 10),
-                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 0), 2)
-
-            # Store data for CSV
-            object_data.append({
-                "Length": round(depth, 2),
-                "Breadth": int(width),
-                "Height": int(height),
-                "Volume": volume
-            })
-
-            # Show object details
-            st.markdown(f"**🧠 Object Detected:**")
-            st.write(f"📏 Length × Breadth × Height: {depth:.2f} × {width} × {height}")
-            st.write(f"📦 Estimated Volume: {volume} (relative units³)")
-
-    # Show annotated image
-    result_img = cv2.cvtColor(image_cv, cv2.COLOR_BGR2RGB)
-    st.image(result_img, caption="Detected Dimensions", use_container_width=True)
-
-    # Show table & CSV download if data exists
-    if object_data:
-        df = pd.DataFrame(object_data)
-        st.markdown("### 📋 Detected Objects Table")
-        st.dataframe(df)
-
-        csv = df.to_csv(index=False).encode('utf-8')
-        st.download_button(
-            label="📥 Download Results as CSV",
-            data=csv,
-            file_name='object_dimensions_volume.csv',
-            mime='text/csv',
-        )
+import streamlit as st
+import cv2
+import numpy as np
+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")
+st.title("📦 Volume Estimation using SAM Segmentation + MiDaS Depth")
+
+# Load SAM and MiDaS models
+@st.cache_resource
+def load_models():
+    sam_checkpoint = "C:/Users/Administrator/Desktop/streamlit_tl/models/sam_vit_h_4b8939.pth"
+    sam = sam_model_registry["vit_h"](checkpoint=sam_checkpoint).to("cuda" if torch.cuda.is_available() else "cpu")
+    predictor = SamPredictor(sam)
+
+    midas = torch.hub.load("intel-isl/MiDaS", "DPT_Large")
+    midas.eval()
+    midas_transform = Compose([
+        Resize(384),
+        ToTensor(),
+        Normalize(mean=[0.5]*3, std=[0.5]*3)
+    ])
+    return predictor, midas, midas_transform
+
+predictor, midas_model, midas_transform = load_models()
+
+# Input source selection
+source_option = st.radio("Select input source", ("Upload Image", "Use Webcam"))
+
+uploaded_file = None
+image_pil = None
+
+if source_option == "Upload Image":
+    uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
+    if uploaded_file:
+        image_pil = Image.open(uploaded_file).convert("RGB")
+
+elif source_option == "Use Webcam":
+    run_camera = st.checkbox("Start Camera")
+
+    if run_camera:
+        cap = cv2.VideoCapture(0)
+        stframe = st.empty()
+        capture = False
+
+        while run_camera and cap.isOpened():
+            ret, frame = cap.read()
+            if not ret:
+                break
+            frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+            stframe.image(frame_rgb, caption="Live Camera Feed", channels="RGB")
+
+            if st.button("📸 Capture Frame"):
+                image_pil = Image.fromarray(frame_rgb)
+                run_camera = False
+                cap.release()
+                break
+
+# Continue processing if we have an image
+if image_pil:
+    image_np = np.array(image_pil)
+    img_h, img_w = image_np.shape[:2]
+    st.image(image_pil, caption="Selected Image", use_container_width=True)
+
+    # Real-world reference dimensions
+    real_image_width_cm = 100
+    real_image_height_cm = 75
+    assumed_max_depth_cm = 100
+
+    pixel_to_cm_x = real_image_width_cm / img_w
+    pixel_to_cm_y = real_image_height_cm / img_h
+
+    # SAM Segmentation
+    predictor.set_image(image_np)
+    masks, _, _ = predictor.predict(multimask_output=False)
+
+    # MiDaS Depth Estimation
+    input_tensor = midas_transform(image_pil).unsqueeze(0)
+    with torch.no_grad():
+        depth_prediction = midas_model(input_tensor).squeeze().cpu().numpy()
+    depth_resized = cv2.resize(depth_prediction, (img_w, img_h))
+
+    # Object volume computation
+    volume_data = []
+    for i, mask in enumerate(masks):
+        mask_np = mask
+        x, y, w, h = cv2.boundingRect(mask_np.astype(np.uint8))
+        width_px = w
+        height_px = h
+
+        width_cm = width_px * pixel_to_cm_x
+        height_cm = height_px * pixel_to_cm_y
+
+        depth_masked = depth_resized[mask_np > 0.5]
+
+        if depth_masked.size == 0:
+            continue
+
+        normalized_depth = (depth_masked - np.min(depth_resized)) / (np.max(depth_resized) - np.min(depth_resized) + 1e-6)
+        depth_cm = np.mean(normalized_depth) * assumed_max_depth_cm
+
+        volume_cm3 = round(depth_cm * width_cm * height_cm, 2)
+
+        volume_data.append({
+            "Object": f"Object #{i+1}",
+            "Length (Depth)": f"{round(depth_cm, 2)} cm",
+            "Breadth (Width)": f"{round(width_cm, 2)} cm",
+            "Height": f"{round(height_cm, 2)} cm",
+            "Volume": f"{volume_cm3} cm³"
+        })
+
+    # Display volume table
+    if volume_data:
+        df = pd.DataFrame(volume_data)
+        st.markdown("### 📊 Object Dimensions and Volume")
+        st.dataframe(df)
+
+        csv = df.to_csv(index=False).encode('utf-8')
+        st.download_button("📂 Download Volume Table as CSV", csv, "object_volumes_with_units.csv", "text/csv")
+    else:
+        st.warning("🚫 No objects were segmented.")