Spaces:
Running
Running
File size: 6,302 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 |
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
import matplotlib.pyplot as plt
# 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)
# Get the current directory
current_dir = os.getcwd()
# Default object file path
DEFAULT_OBJ_FILE = os.path.join(current_dir, "female.obj")
# Temporary texture file path
TEMP_TEXTURE_FILE = os.path.join(current_dir, "generated_texture.png")
# File path to save the 2D image
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):
texture_image = Image.open(texture_file)
uv_coords = mesh.visual.uv
uv_coords = np.clip(uv_coords, 0, 1)
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
])
texture_colors = texture_colors / 255.0
return texture_colors
def display_3d_object(obj_file, texture_file, light_intensity, ambient_intensity, color):
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.")
if texture_file:
colors = apply_texture(mesh, texture_file)
else:
colors = color
ambient_intensity = max(0, min(ambient_intensity, 1))
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):
trimesh_scene = trimesh.load(filename)
if isinstance(trimesh_scene, trimesh.Scene):
mesh = trimesh_scene.dump(concatenate=True)
else:
mesh = trimesh_scene
return mesh
def extract_uv_texture(obj_file):
mesh = trimesh.load(obj_file)
if mesh.visual.uv is None or len(mesh.visual.uv) == 0:
raise ValueError("The mesh does not have UV mapping information.")
texture_image = None
if mesh.visual.material.image is not None:
texture_image = mesh.visual.material.image
else:
texture_size = 1024
texture_image = Image.new('RGB', (texture_size, texture_size), color=(255, 255, 255))
return texture_image
def generate_clothing_image(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")
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):
uv_texture_image = extract_uv_texture(file)
uv_texture_image.save(OUTPUT_IMAGE_FILE)
return uv_texture_image
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)
|