Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| # 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 | |
| from PIL import Image | |
| from io import BytesIO | |
| import numpy as np | |
| # Set the title of the app | |
| st.title('Image to 3D Asset Generator') | |
| # Upload image file | |
| uploaded_file = st.file_uploader("Upload an Image", type=["png", "jpg", "jpeg"]) | |
| # Shape options | |
| shape_options = [ | |
| "Box", "Sphere", "Cone", "Cylinder", "Torus", "Icosahedron", "Octahedron", "Dodecahedron", | |
| "Grid", "Plane" | |
| ] | |
| # Select the shape | |
| shape = st.selectbox("Choose a 3D Shape", shape_options) | |
| # Transformation Options | |
| scale_x = st.slider("Scale X", min_value=0.1, max_value=5.0, value=1.0) | |
| scale_y = st.slider("Scale Y", min_value=0.1, max_value=5.0, value=1.0) | |
| scale_z = st.slider("Scale Z", min_value=0.1, max_value=5.0, value=1.0) | |
| rotation_x = st.slider("Rotate X (degrees)", min_value=0, max_value=360, value=0) | |
| rotation_y = st.slider("Rotate Y (degrees)", min_value=0, max_value=360, value=0) | |
| rotation_z = st.slider("Rotate Z (degrees)", min_value=0, max_value=360, value=0) | |
| position_x = st.slider("Position X", min_value=-5.0, max_value=5.0, value=0.0) | |
| position_y = st.slider("Position Y", min_value=-5.0, max_value=5.0, value=0.0) | |
| position_z = st.slider("Position Z", min_value=-5.0, max_value=5.0, value=0.0) | |
| # Color and Material options | |
| color = st.color_picker("Choose a Color", "#00ff00") | |
| file_format = st.selectbox("Choose Export Format", ["GLB", "OBJ", "STL"]) | |
| # Process the image and generate 3D asset | |
| if uploaded_file is not None: | |
| try: | |
| # 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": | |
| mesh = trimesh.creation.uv_sphere(radius=1.0) | |
| elif shape == "Cone": | |
| mesh = trimesh.creation.cone(radius=1.0, height=1.0) | |
| elif shape == "Cylinder": | |
| mesh = trimesh.creation.cylinder(radius=1.0, height=1.0) | |
| elif shape == "Torus": | |
| mesh = trimesh.creation.torus(radius=1.0, tube_radius=0.5) | |
| elif shape == "Icosahedron": | |
| mesh = trimesh.creation.icosahedron() | |
| elif shape == "Octahedron": | |
| mesh = trimesh.creation.octahedron() | |
| elif shape == "Dodecahedron": | |
| mesh = trimesh.creation.dodecahedron() | |
| elif shape == "Grid": | |
| mesh = trimesh.creation.box(extents=[5, 5, 0.1]) | |
| elif shape == "Plane": | |
| mesh = trimesh.creation.box(extents=[5, 5, 0.01]) | |
| else: | |
| st.error("Unsupported shape selected.") | |
| if mesh: | |
| # Apply transformations | |
| mesh.apply_scale([scale_x, scale_y, scale_z]) | |
| rotation_matrix = transformations.euler_matrix( | |
| np.radians(rotation_x), np.radians(rotation_y), np.radians(rotation_z) | |
| ) | |
| mesh.apply_transform(rotation_matrix) | |
| mesh.apply_translation([position_x, position_y, position_z]) | |
| # Apply color | |
| material_color = np.array( | |
| [int(color[1:3], 16) / 255, int(color[3:5], 16) / 255, int(color[5:7], 16) / 255, 1.0] | |
| ) | |
| mesh.visual.face_colors = (material_color * 255).astype(np.uint8) | |
| # Export the mesh | |
| with io.BytesIO() as buffer: | |
| mesh.export(buffer, file_type=file_format.lower()) | |
| buffer.seek(0) | |
| st.download_button( | |
| label=f"Download 3D Asset ({file_format})", | |
| data=buffer, | |
| file_name=f"output_model.{file_format.lower()}", | |
| mime=f"model/{file_format.lower()}" | |
| ) | |
| except Exception as e: | |
| st.error(f"An error occurred while generating the 3D asset: {e}") |