Spaces:
Running
Running
File size: 7,129 Bytes
027d724 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 |
import gradio as gr
import plotly.graph_objs as go
import trimesh
import numpy as np
from PIL import Image
import torch
from diffusers import StableDiffusionPipeline
import os
from pythreejs import *
from IPython.display import display
# Load the Stable Diffusion model for text-to-image generation
device = "cuda" if torch.cuda.is_available() else "cpu"
pipeline = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4").to(device)
import os
# Get the current directory
current_dir = os.getcwd()
# Default object file path (replaced with current directory)
DEFAULT_OBJ_FILE = os.path.join(current_dir, "female.obj")
# Temporary texture file path (replaced with current directory)
TEMP_TEXTURE_FILE = os.path.join(current_dir, "generated_texture.png")
# File path to save the 2D image (replaced with current directory)
OUTPUT_IMAGE_FILE = os.path.join(current_dir, "output_image.png")
DEFAULT_GLB_FILE= os.path.join(current_dir, "vroid_girl1.glb")
def apply_texture(mesh, texture_file):
# Load the texture image
texture_image = Image.open(texture_file)
# Get texture coordinates and faces from the mesh
uv_coords = mesh.visual.uv
# Ensure UV coordinates are within the valid range
uv_coords = np.clip(uv_coords, 0, 1)
# Extract texture colors for each vertex
texture_colors = np.array([
texture_image.getpixel((
int(u * (texture_image.width - 1)),
int(v * (texture_image.height - 1))
)) for u, v in uv_coords
])
# Normalize colors to be between 0 and 1
texture_colors = texture_colors / 255.0
return texture_colors
def display_3d_object(obj_file, texture_file, light_intensity, ambient_intensity, color):
# Determine the file type and load accordingly
file_extension = obj_file.split('.')[-1].lower()
if (file_extension == 'obj'):
mesh = trimesh.load(obj_file)
elif file_extension == 'glb':
mesh = load_glb_file(obj_file)
else:
raise ValueError("Unsupported file format. Please upload a .obj or .glb file.")
# Apply texture if available
if texture_file:
colors = apply_texture(mesh, texture_file)
else:
colors = color
# Clamp the ambient_intensity to be between 0 and 1
ambient_intensity = max(0, min(ambient_intensity, 1))
# Create a 3D plot using Plotly
fig = go.Figure(data=[
go.Mesh3d(
x=mesh.vertices[:, 0],
y=mesh.vertices[:, 1],
z=mesh.vertices[:, 2],
i=mesh.faces[:, 0],
j=mesh.faces[:, 1],
k=mesh.faces[:, 2],
facecolor=colors if texture_file else None,
color=color if not texture_file else None,
opacity=0.50,
lighting=dict(
ambient=ambient_intensity,
diffuse=light_intensity,
specular=0.5,
roughness=0.1,
fresnel=0.2
),
lightposition=dict(
x=100,
y=200,
z=300
)
)
])
fig.update_layout(scene=dict(aspectmode='data'))
return fig
def load_glb_file(filename):
# Load GLB file using trimesh
trimesh_scene = trimesh.load(filename)
# Extract the first mesh from the scene
if isinstance(trimesh_scene, trimesh.Scene):
mesh = trimesh_scene.dump(concatenate=True)
else:
mesh = trimesh_scene
return mesh
def extract_uv_texture(obj_file):
# Load the 3D model
mesh = trimesh.load(obj_file)
# Ensure the mesh has visual/texture information
if mesh.visual.uv is None or len(mesh.visual.uv) == 0:
raise ValueError("The mesh does not have UV mapping information.")
# Extract the texture map from the mesh, if it exists
texture_image = None
if mesh.visual.material.image is not None:
texture_image = mesh.visual.material.image
else:
# If no texture is found, create a blank white texture
texture_size = 1024
texture_image = Image.new('RGB', (texture_size, texture_size), color=(255, 255, 255))
return texture_image
def generate_clothing_image(prompt):
"""Generate an image of clothing based on the text prompt."""
image = pipeline(prompt).images[0]
image.save(TEMP_TEXTURE_FILE)
return TEMP_TEXTURE_FILE, image
def update_texture_display(prompt, texture_file):
if prompt:
texture_path, image = generate_clothing_image(prompt)
return image
elif texture_file:
return Image.open(texture_file)
return None
with gr.Blocks() as demo:
gr.Markdown("## 3D Object Viewer with Custom Texture, Color, and Adjustable Lighting")
with gr.Row():
with gr.Column(scale=1):
gr.Markdown("### Texture Options")
prompt_input = gr.Textbox(label="Enter a Prompt to Generate Texture", placeholder="Type a prompt...")
generate_button = gr.Button("Generate Texture")
texture_file = gr.File(label="Upload Texture file (PNG or JPG, optional)", type="filepath")
texture_preview = gr.Image(label="Texture Preview", visible=True)
gr.Markdown("### Lighting & Color Settings")
light_intensity_slider = gr.Slider(minimum=0, maximum=2, step=0.1, value=0.8, label="Light Intensity")
ambient_intensity_slider = gr.Slider(minimum=0, maximum=1, step=0.1, value=0.5, label="Ambient Intensity") # Updated the max value to 1
color_picker = gr.ColorPicker(value="#D3D3D3", label="Object Color")
submit_button = gr.Button("Submit")
obj_file = gr.File(label="Upload OBJ or GLB file", value=DEFAULT_OBJ_FILE, type='filepath')
with gr.Column(scale=2):
display = gr.Plot(label="3D Viewer")
extract_button = gr.Button("Extract UV Texture")
output_image = gr.Image(label="Extracted UV Texture", visible=True)
def update_display(file, texture, light_intensity, ambient_intensity, color):
texture_to_use = TEMP_TEXTURE_FILE if os.path.exists(TEMP_TEXTURE_FILE) else texture
return display_3d_object(file, texture_to_use, light_intensity, ambient_intensity, color)
def extract_and_display_uv_texture(file):
return extract_uv_texture(file)
submit_button.click(fn=update_display, inputs=[obj_file, texture_file, light_intensity_slider, ambient_intensity_slider, color_picker], outputs=display)
generate_button.click(fn=update_texture_display, inputs=[prompt_input, texture_file], outputs=texture_preview)
texture_file.change(fn=update_texture_display, inputs=[prompt_input, texture_file], outputs=texture_preview)
extract_button.click(fn=extract_and_display_uv_texture, inputs=[obj_file], outputs=output_image)
demo.load(fn=update_display, inputs=[obj_file, texture_file, light_intensity_slider, ambient_intensity_slider, color_picker], outputs=display)
gr.Examples(
examples=[[DEFAULT_OBJ_FILE, None],[DEFAULT_GLB_FILE, None]],
inputs=[obj_file, texture_file],
label="Example Files"
)
demo.launch(debug=True)
|