Update app.py

app.py

@@ -1,5 +1,11 @@
 import streamlit as st
-from rembg import remove
+
+# Attempt to import required libraries and handle missing dependencies
+try:
+    from rembg import remove
+except ModuleNotFoundError:
+    st.error("The `rembg` module requires `onnxruntime`. Please install it using `pip install onnxruntime` or `pip install onnxruntime-gpu` if you have a GPU.")
+
 import io
 import trimesh
 from trimesh import transformations
@@ -39,22 +45,24 @@ position_z = st.slider("Position Z", min_value=-5.0, max_value=5.0, value=0.0)
 color = st.color_picker("Choose a Color", "#00ff00")
 file_format = st.selectbox("Choose Export Format", ["GLB", "OBJ", "STL"])
 
-# When the file is uploaded
+# Process the image and generate 3D asset
 if uploaded_file is not None:
-    # Display the image
-    st.image(uploaded_file, caption="Uploaded Image", use_column_width=True)
-    
-    # Process the image for background removal (using rembg)
-    input_data = uploaded_file.read()
-    output_data = remove(input_data)
-
-    # Display the processed image (background removed)
-    output_image = Image.open(BytesIO(output_data))
-    st.image(output_image, caption="Processed Image (Background Removed)", use_column_width=True)
-    
-    # Create the selected shape and apply transformations
     try:
-        # Create the mesh
+        # Display the image
+        st.image(uploaded_file, caption="Uploaded Image", use_column_width=True)
+        
+        # Process the image for background removal (if `rembg` is available)
+        if "remove" in globals():
+            input_data = uploaded_file.read()
+            output_data = remove(input_data)
+            # Display the processed image (background removed)
+            output_image = Image.open(BytesIO(output_data))
+            st.image(output_image, caption="Processed Image (Background Removed)", use_column_width=True)
+        else:
+            st.warning("Image processing with background removal is unavailable. Please install `onnxruntime` and restart the app.")
+        
+        # Create the selected shape and apply transformations
+        mesh = None
         if shape == "Box":
             mesh = trimesh.creation.box(extents=[1, 1, 1])
         elif shape == "Sphere":
@@ -77,19 +85,14 @@ if uploaded_file is not None:
             mesh = trimesh.creation.box(extents=[5, 5, 0.01])
         else:
             st.error("Unsupported shape selected.")
-            mesh = None
-        
+
         if mesh:
-            # Apply scaling
+            # Apply transformations
             mesh.apply_scale([scale_x, scale_y, scale_z])
-            
-            # Apply rotation
             rotation_matrix = transformations.euler_matrix(
                 np.radians(rotation_x), np.radians(rotation_y), np.radians(rotation_z)
             )
             mesh.apply_transform(rotation_matrix)
-            
-            # Apply translation
             mesh.apply_translation([position_x, position_y, position_z])
 
             # Apply color