Update app.py

app.py

@@ -1,26 +1,25 @@
-import streamlit as st
-import cv2
+import gradio as gr
+import torch
 import numpy as np
-import pandas as pd
+import cv2
 from PIL import Image
-import torch
+import pandas as pd
 from torchvision.transforms import Compose, Resize, ToTensor, Normalize
-#from segment_anything import SamPredictor, sam_model_registry
-from segment_anything import sam_model_registry, SamPredictor
-
-# Set Streamlit configuration
-st.set_page_config(page_title="Volume Estimator", layout="wide")
-st.title("📦 Volume Estimation using SAM Segmentation + MiDaS Depth")
+from segment_anything import SamPredictor, sam_model_registry
+import os
 
 # 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")
+    sam_checkpoint = "sam_vit_h_4b8939.pth"
+    if not os.path.exists(sam_checkpoint):
+        raise FileNotFoundError("Please upload the SAM checkpoint file to the working directory.")
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    sam = sam_model_registry["vit_h"](checkpoint=sam_checkpoint).to(device)
     predictor = SamPredictor(sam)
 
     midas = torch.hub.load("intel-isl/MiDaS", "DPT_Large")
-    midas.eval()
+    midas.eval().to(device)
     midas_transform = Compose([
         Resize(384),
         ToTensor(),
@@ -30,45 +29,12 @@ def load_models():
 
 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:
+# Processing function
+def process_image(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-world reference dimensions (adjust as needed)
     real_image_width_cm = 100
     real_image_height_cm = 75
     assumed_max_depth_cm = 100
@@ -76,12 +42,12 @@ if image_pil:
     pixel_to_cm_x = real_image_width_cm / img_w
     pixel_to_cm_y = real_image_height_cm / img_h
 
-    # SAM Segmentation
+    # SAM segmentation
     predictor.set_image(image_np)
     masks, _, _ = predictor.predict(multimask_output=False)
 
-    # MiDaS Depth Estimation
-    input_tensor = midas_transform(image_pil).unsqueeze(0)
+    # MiDaS depth estimation
+    input_tensor = midas_transform(image_pil).unsqueeze(0).to(next(midas_model.parameters()).device)
     with torch.no_grad():
         depth_prediction = midas_model(input_tensor).squeeze().cpu().numpy()
     depth_resized = cv2.resize(depth_prediction, (img_w, img_h))
@@ -89,39 +55,43 @@ if image_pil:
     # 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))
+        x, y, w, h = cv2.boundingRect(mask.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]
-
+        depth_masked = depth_resized[mask > 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.")
+        volume_data.append([
+            f"Object #{i+1}",
+            round(depth_cm, 2),
+            round(width_cm, 2),
+            round(height_cm, 2),
+            volume_cm3
+        ])
+
+    if not volume_data:
+        return image_pil, "No objects segmented."
+
+    df = pd.DataFrame(volume_data, columns=["Object", "Length (Depth) cm", "Breadth (Width) cm", "Height cm", "Volume cm³"])
+    return image_pil, df
+
+# Gradio Interface
+with gr.Blocks() as demo:
+    gr.Markdown("# 📦 Volume Estimation using SAM + MiDaS")
+    with gr.Row():
+        image_input = gr.Image(type="pil", label="Upload Image")
+        run_btn = gr.Button("Estimate Volume")
+    with gr.Row():
+        output_image = gr.Image(label="Original Image")
+        volume_table = gr.Dataframe(headers=["Object", "Length (Depth) cm", "Breadth (Width) cm", "Height cm", "Volume cm³"])
+    run_btn.click(fn=process_image, inputs=image_input, outputs=[output_image, volume_table])
+
+demo.launch()