Create App. Py

App. Py

@@ -0,0 +1,79 @@
+import streamlit as st
+import numpy as np
+import torch
+import torchvision.transforms as T
+from PIL import Image
+import matplotlib.pyplot as plt
+import open3d as o3d
+
+# Load MiDaS model
+@st.cache_resource
+def load_model():
+    model_type = "DPT_Large"  # Use DPT_Large for higher accuracy
+    model = torch.hub.load("intel-isl/MiDaS", model_type)
+    model.eval()
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
+
+    midas_transforms = torch.hub.load("intel-isl/MiDaS", "transforms")
+    transform = midas_transforms.default_transform
+    return model, transform, device
+
+model, transform, device = load_model()
+
+# Streamlit app UI
+st.title("2D to 3D Image Converter")
+st.write("Upload a 2D image to generate its 3D depth map and point cloud.")
+
+# File uploader
+uploaded_file = st.file_uploader("Upload an Image", type=["jpg", "jpeg", "png"])
+if uploaded_file:
+    # Load the image
+    input_image = Image.open(uploaded_file).convert("RGB")
+    st.image(input_image, caption="Uploaded Image", use_column_width=True)
+
+    # Preprocess the image
+    input_batch = transform(input_image).unsqueeze(0).to(device)
+
+    # Predict depth map
+    st.write("Generating Depth Map...")
+    with torch.no_grad():
+        prediction = model(input_batch)
+        depth_map = torch.nn.functional.interpolate(
+            prediction.unsqueeze(1),
+            size=input_image.size[::-1],
+            mode="bicubic",
+            align_corners=False,
+        ).squeeze().cpu().numpy()
+
+    # Display depth map
+    st.write("Depth Map:")
+    fig, ax = plt.subplots()
+    ax.imshow(depth_map, cmap="plasma")
+    ax.axis("off")
+    st.pyplot(fig)
+
+    # Generate 3D point cloud
+    st.write("Generating 3D Point Cloud...")
+    h, w = depth_map.shape
+    xx, yy = np.meshgrid(np.arange(w), np.arange(h))
+    points = np.stack((xx, yy, depth_map), axis=-1).reshape(-1, 3)
+
+    # Normalize points
+    points -= points.mean(axis=0)
+    points /= points.max(axis=0)
+
+    # Create Open3D point cloud
+    pcd = o3d.geometry.PointCloud()
+    pcd.points = o3d.utility.Vector3dVector(points)
+
+    # Save point cloud for download
+    output_file = "point_cloud.ply"
+    o3d.io.write_point_cloud(output_file, pcd)
+    st.write("3D Point Cloud Generated!")
+    st.download_button(
+        label="Download 3D Point Cloud (.ply)",
+        data=open(output_file, "rb").read(),
+        file_name="3d_point_cloud.ply",
+        mime="application/octet-stream",
+    )