Update app.py

app.py

@@ -23,38 +23,37 @@ if uploaded_file is not None:
     image = Image.open(uploaded_file)
     st.image(image, caption="Uploaded Image", use_column_width=True)
 
-    # Process image with DepthPro for depth estimation
-    inputs = image_processor(images=image, return_tensors="pt").to(device)
-    with torch.no_grad():
-        outputs = model(**inputs)
-
-    # Post-process depth output
-    post_processed_output = image_processor.post_process_depth_estimation(
-        outputs, target_sizes=[(image.height, image.width)],
-    )
-
-    depth = post_processed_output[0]["predicted_depth"]
-    depth = (depth - depth.min()) / (depth.max() - depth.min())
-    depth = depth * 255.
-    depth = depth.detach().cpu().numpy()
-    depth_image = Image.fromarray(depth.astype("uint8"))
-
-    st.subheader("Depth Map")
-    st.image(depth_image, caption="Estimated Depth Map", use_column_width=True)
-
-    # Colorize the depth map to make it more visible
-    colormap = depth_image.convert("RGB")
-    st.subheader("Colorized Depth Map")
-    st.image(colormap, caption="Colorized Depth Map", use_column_width=True)
-
-    # Option to save depth image
-    if st.button('Save Depth Image'):
-        depth_image.save('depth_image.png')
-        st.success("Depth image saved successfully!")
-
-    # Option for interactive painting (Placeholder)
-    st.subheader("Interactive Depth-based Painting (Demo Placeholder)")
-    st.write("This feature will allow users to paint on surfaces based on depth. For now, we can show the depth and its effects.")
-
-# Placeholder for future interactive painting functionality.
-# This could be extended with AR-based libraries or Unity integration in a full-scale app.
+    # Add Generate button
+    if st.button("Generate"):
+        # Process image with DepthPro for depth estimation
+        inputs = image_processor(images=image, return_tensors="pt").to(device)
+        with torch.no_grad():
+            outputs = model(**inputs)
+
+        # Post-process depth output
+        post_processed_output = image_processor.post_process_depth_estimation(
+            outputs, target_sizes=[(image.height, image.width)],
+        )
+
+        depth = post_processed_output[0]["predicted_depth"]
+        depth = (depth - depth.min()) / (depth.max() - depth.min())
+        depth = depth * 255.
+        depth = depth.detach().cpu().numpy()
+        depth_image = Image.fromarray(depth.astype("uint8"))
+
+        st.subheader("Depth Map")
+        st.image(depth_image, caption="Estimated Depth Map", use_column_width=True)
+
+        # Colorize the depth map to make it more visible
+        colormap = depth_image.convert("RGB")
+        st.subheader("Colorized Depth Map")
+        st.image(colormap, caption="Colorized Depth Map", use_column_width=True)
+
+        # Option to save depth image
+        if st.button('Save Depth Image'):
+            depth_image.save('depth_image.png')
+            st.success("Depth image saved successfully!")
+
+        # Option for interactive painting (Placeholder)
+        st.subheader("Interactive Depth-based Painting (Demo Placeholder)")
+        st.write("This feature will allow users to paint on surfaces based on depth. For now, we can show the depth and its effects.")