Spaces:
Running
Running
File size: 5,617 Bytes
0a89b21 027d724 0a89b21 027d724 0a89b21 027d724 0a89b21 027d724 0a89b21 027d724 0a89b21 027d724 0a89b21 027d724 0a89b21 027d724 0a89b21 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 |
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'))
#Cleaning Temp file
if os.path.exists(TEMP_TEXTURE_FILE):
os.remove(TEMP_TEXTURE_FILE)
print(f"Deleted existing file: {TEMP_TEXTURE_FILE}")
else:
print(f"File not found: {TEMP_TEXTURE_FILE}")
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 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")
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)
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)
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)
|