import altair as alt
import numpy as np
import pandas as pd
import streamlit as st

"""
# Welcome to Streamlit!

Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
forums](https://discuss.streamlit.io).

In the meantime, below is an example of what you can do with just a few lines of code:
"""

num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
num_turns = st.slider("Number of turns in spiral", 1, 300, 31)

indices = np.linspace(0, 1, num_points)
theta = 2 * np.pi * num_turns * indices
radius = indices

x = radius * np.cos(theta)
y = radius * np.sin(theta)

df = pd.DataFrame({
    "x": x,
    "y": y,
    "idx": indices,
    "rand": np.random.randn(num_points),
})

st.altair_chart(alt.Chart(df, height=700, width=700)
    .mark_point(filled=True)
    .encode(
        x=alt.X("x", axis=None),
        y=alt.Y("y", axis=None),
        color=alt.Color("idx", legend=None, scale=alt.Scale()),
        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
    ))


import streamlit as st
import cv2
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.")