updates new formats

.gitignore

@@ -1,2 +1,3 @@
 
 /.venv
+*.obj

app.py

@@ -5,8 +5,10 @@ import numpy as np
 from PIL import Image, ImageDraw
 import torch
 from diffusers import StableDiffusionPipeline, StableDiffusionInpaintPipeline
-import os
+import io
 import matplotlib.pyplot as plt
+#import pyrender
+#import scipy
 
 # Load the Stable Diffusion model for text-to-image generation and inpainting
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -16,58 +18,85 @@ pipeline_inpaint = StableDiffusionInpaintPipeline.from_pretrained(
     torch_dtype=torch.float16
 ).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")
+DEFAULT_OBJ_FILE = "female.obj"
+DEFAULT_GLB_FILE = "vroid_girl1.glb"
+DEFAULT_VRM_FILE = "fischl.vrm"
+DEFAULT_VRM_FILE2 = "woman.vrm"
+DEFAULT_VRM_FILE3 = "mona.vrm"
+DEFAULT_TEXTURE = "future.png"
+DEFAULT_TEXTURE2 = "woman1.jpeg"
+DEFAULT_TEXTURE3 = "woman2.jpeg"
+def generate_default_uv(mesh, quality='medium'):
+    """
+    Generate default UV coordinates for a mesh if UV mapping is missing.
+    """
+    if quality == 'low':
+        bounds = mesh.bounds
+        width = bounds[1][0] - bounds[0][0]
+        height = bounds[1][1] - bounds[0][1]
+        uv_coords = np.zeros((len(mesh.vertices), 2))
+        uv_coords[:, 0] = (mesh.vertices[:, 0] - bounds[0][0]) / width
+        uv_coords[:, 1] = (mesh.vertices[:, 1] - bounds[0][1]) / height
+    
+    elif quality == 'medium':
+        height_range = mesh.vertices[:, 2].max() - mesh.vertices[:, 2].min()
+        radius = np.sqrt(mesh.vertices[:, 0]**2 + mesh.vertices[:, 1]**2)
+        uv_coords = np.zeros((len(mesh.vertices), 2))
+        uv_coords[:, 0] = np.arctan2(mesh.vertices[:, 1], mesh.vertices[:, 0]) / (2 * np.pi) + 0.5
+        uv_coords[:, 1] = (mesh.vertices[:, 2] - mesh.vertices[:, 2].min()) / height_range
+    
+    elif quality == 'high':
+        radius = np.sqrt(np.sum(mesh.vertices**2, axis=1))
+        uv_coords = np.zeros((len(mesh.vertices), 2))
+        uv_coords[:, 0] = np.arctan2(mesh.vertices[:, 1], mesh.vertices[:, 0]) / (2 * np.pi) + 0.5
+        uv_coords[:, 1] = np.arccos(mesh.vertices[:, 2] / radius) / np.pi
+    else:
+        raise ValueError("Invalid quality parameter. Choose from 'low', 'medium', or 'high'.")
 
+    return uv_coords
 
-def apply_texture(mesh, texture_file, uv_scale):
+def apply_texture(mesh, texture_image, uv_scale, uv_quality='medium'):
     """
-    Applies the texture to the mesh with UV scaling to make triangles/rectangles smaller or larger.
+    Applies the texture to the mesh with UV scaling.
     """
-    texture_image = Image.open(texture_file)
-    uv_coords = mesh.visual.uv
+    if not hasattr(mesh.visual, 'uv') or mesh.visual.uv is None:
+        # If the mesh does not have UV coordinates, generate them
+        print("No UV coordinates found; generating default UV mapping.")
+        uv_coords = generate_default_uv(mesh, quality=uv_quality)
+    else:
+        uv_coords = mesh.visual.uv
 
-    # Apply scaling to UV coordinates to make the mapping finer or coarser
+    # Ensure UV coordinates exist
+    if uv_coords is None:
+        raise ValueError("UV coordinates are missing from the mesh.")
+
+    # Apply UV scaling and ensure it is within valid range
     uv_coords = np.clip(uv_coords * uv_scale, 0, 1)
 
-    # Get the size of the texture image
     img_width, img_height = texture_image.size
     texture_array = np.array(texture_image)
 
-    # Prepare to store the colors per face
     face_colors = []
 
     for face in mesh.faces:
         uv_face = uv_coords[face]
         pixel_coords = np.round(uv_face * np.array([img_width - 1, img_height - 1])).astype(int)
 
-        # Ensure the UV coordinates are within the bounds of the texture image
         valid_coords = np.all((pixel_coords[:, 0] >= 0) & (pixel_coords[:, 0] < img_width) &
                               (pixel_coords[:, 1] >= 0) & (pixel_coords[:, 1] < img_height))
 
         if valid_coords:
-            # Get the average color for the face from the corresponding UV points in the texture
             face_color = np.mean(texture_array[pixel_coords[:, 1], pixel_coords[:, 0]], axis=0)
-            face_colors.append(face_color / 255.0)  # Normalize to [0, 1]
+            face_colors.append(face_color / 255.0)
         else:
-            # Assign a default color (e.g., gray) if UV coordinates are not valid
-            face_colors.append([0.5, 0.5, 0.5])  # Default to gray if no texture is applied
+            face_colors.append([0.5, 0.5, 0.5])
 
-    # Ensure no face is left unpainted
     face_colors = np.array(face_colors)
     if len(face_colors) < len(mesh.faces):
         face_colors = np.pad(face_colors, ((0, len(mesh.faces) - len(face_colors)), (0, 0)), 'constant', constant_values=0.5)
 
     return face_colors
 
-
 def load_glb_file(filename):
     trimesh_scene = trimesh.load(filename)
     if isinstance(trimesh_scene, trimesh.Scene):
@@ -81,112 +110,115 @@ def generate_clothing_image(prompt, num_inference_steps):
     Generates the clothing texture based on the provided prompt and number of inference steps.
     """
     image = pipeline(prompt, num_inference_steps=num_inference_steps).images[0]
-    image.save(TEMP_TEXTURE_FILE)
-    return TEMP_TEXTURE_FILE, image
+    return image
 
-def generate_uv_specific_texture(prompt, uv_map_file, num_inference_steps=50):
-    """
-    Generates a texture for the 3D model using a given prompt and UV map.
-    
-    Args:
-        prompt (str): The prompt for the diffusion model to generate the texture.
-        uv_map_file (str): Path to the UV map file.
-        num_inference_steps (int): The number of iterations/steps for the diffusion process.
-        
-    Returns:
-        (str, PIL.Image): The path to the generated texture file and the generated texture image.
-    """
-    # Load UV map as a mask
-    uv_map = Image.open(uv_map_file)
-    
-    # Generate texture based on UV map and the provided prompt
-    image = pipeline_inpaint(
-        prompt=prompt, 
-        image=uv_map, 
-        mask_image=uv_map, 
-        num_inference_steps=num_inference_steps  # Set custom number of inference steps
-    ).images[0]
-    
-    # Save the generated texture
-    image.save(TEMP_TEXTURE_FILE)
-    
-    return TEMP_TEXTURE_FILE, image
+def load_vrm_file(filename):
+    try:
+        vrm_data = trimesh.load(filename, file_type='glb')
+        if isinstance(vrm_data, trimesh.Scene):
+            mesh = vrm_data.dump(concatenate=True)
+        else:
+            mesh = vrm_data
+    except Exception as e:
+        raise ValueError(f"Failed to load VRM file: {e}")
 
-def display_3d_object(obj_file, texture_file, light_intensity, ambient_intensity, color, uv_scale, transparency):
-    """
-    Displays the 3D object with applied texture or color, with support for UV scaling and transparency.
-    """
+    return mesh
+
+def display_3d_object(obj_file, texture_image, light_intensity, ambient_intensity, color, uv_scale, transparency, uv_quality=None):
     file_extension = obj_file.split('.')[-1].lower()
 
-    # Load mesh
-    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 file_extension == 'vrm':
+        mesh = load_vrm_file(obj_file)
+        try:
+            if texture_image:
+                face_colors = apply_texture(mesh, texture_image, uv_scale, uv_quality)
+            else:
+                face_colors = np.array([color] * len(mesh.faces))
+        except ValueError as e:
+            face_colors = np.array([color] * len(mesh.faces))
+
+        vertices = mesh.vertices
+        faces = mesh.faces
+
+        fig = go.Figure(data=[
+            go.Mesh3d(
+                x=vertices[:, 0],
+                y=vertices[:, 1],
+                z=vertices[:, 2],
+                i=faces[:, 0],
+                j=faces[:, 1],
+                k=faces[:, 2],
+                facecolor=face_colors,
+                opacity=transparency,
+                lighting=dict(
+                    ambient=ambient_intensity,
+                    diffuse=light_intensity,
+                    specular=0.8,
+                    roughness=0.3,
+                    fresnel=0.1
+                ),
+                lightposition=dict(
+                    x=100,
+                    y=200,
+                    z=300
+                )
+            )
+        ])
+        fig.update_layout(scene=dict(aspectmode='data'))
+        return fig
 
-    # Apply texture or color
-    if texture_file:
-        face_colors = apply_texture(mesh, texture_file, uv_scale)
     else:
-        face_colors = np.array([color] * len(mesh.faces))  # Use a single color for all faces if no texture
-
-    # Define lighting settings
-    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=face_colors if texture_file else None,
-            color=color if not texture_file else None,
-            opacity=transparency,  # Adjustable transparency
-            lighting=dict(
-                ambient=ambient_intensity,
-                diffuse=light_intensity,
-                specular=0.8,  # Fine-tuned specular to avoid excessive shininess
-                roughness=0.3,
-                fresnel=0.1
-            ),
-            lightposition=dict(
-                x=100,
-                y=200,
-                z=300
+        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, .glb, or .vrm file.")
+
+        if texture_image:
+            face_colors = apply_texture(mesh, texture_image, uv_scale)
+        else:
+            face_colors = np.array([color] * len(mesh.faces))
+
+        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=face_colors,
+                opacity=transparency,
+                lighting=dict(
+                    ambient=ambient_intensity,
+                    diffuse=light_intensity,
+                    specular=0.8,
+                    roughness=0.3,
+                    fresnel=0.1
+                ),
+                lightposition=dict(
+                    x=100,
+                    y=200,
+                    z=300
+                )
             )
-        )
-    ])
-    fig.update_layout(scene=dict(aspectmode='data'))
+        ])
+        fig.update_layout(scene=dict(aspectmode='data'))
 
-    return fig
+        return fig
 
 def clear_texture():
-    """
-    Clears the texture preview and removes the texture file, allowing for a reset.
-    """
-    if os.path.exists(TEMP_TEXTURE_FILE):
-        os.remove(TEMP_TEXTURE_FILE)
     return None
 
 def restore_original(obj_file):
-    """
-    Restores the original 3D object without any applied texture.
-    """
-    return display_3d_object(obj_file, None, 0.8, 0.5, "#D3D3D3", 1.0, 1.0)  # Default settings for restoration
+    return display_3d_object(obj_file, None, 0.8, 0.5, "#D3D3D3", 1.0, 1.0)
 
 def update_texture_display(prompt, texture_file, num_inference_steps):
-    """
-    Update the texture display either by generating a texture from the prompt
-    or by displaying an uploaded texture.
-    """
     if prompt:
-        # Generate new texture based on the prompt using the Stable Diffusion pipeline
-        texture_path, image = generate_clothing_image(prompt, num_inference_steps)
+        image = generate_clothing_image(prompt, num_inference_steps)
         return image
     elif texture_file:
-        # Display the uploaded texture file
         return Image.open(texture_file)
     return None
 
@@ -203,35 +235,65 @@ with gr.Blocks() as demo:
             texture_preview = gr.Image(label="Texture Preview", visible=True)
 
             gr.Markdown("### Mapping, Lighting & Color Settings")
-            uv_scale_slider = gr.Slider(minimum=0.1, maximum=5, step=0.1, value=1.0, label="UV Mapping Scale (Make smaller/bigger)")
+            uv_scale_slider = gr.Slider(minimum=0.1, maximum=5, step=0.1, value=1.0, label="UV Mapping Scale")
+            uv_quality_dropdown = gr.Dropdown(label="UV Quality (for VRM files)", choices=['low', 'medium', 'high'], value='medium')
             light_intensity_slider = gr.Slider(minimum=0, maximum=1, 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")
-            transparency_slider = gr.Slider(minimum=0.1, maximum=1.0, step=0.1, value=1.0, label="Transparency (1.0 is fully opaque)")
+            transparency_slider = gr.Slider(minimum=0.1, maximum=1.0, step=0.1, value=1.0, label="Transparency")
             color_picker = gr.ColorPicker(value="#D3D3D3", label="Object Color")
             submit_button = gr.Button("Submit")
             restore_button = gr.Button("Restore")
             clear_button = gr.Button("Clear")
-            obj_file = gr.File(label="Upload OBJ or GLB file", value=DEFAULT_OBJ_FILE, type='filepath')
+            obj_file = gr.File(label="Upload OBJ, GLB, or VRM file", value=DEFAULT_OBJ_FILE, type='filepath')
 
         with gr.Column(scale=2):
             display = gr.Plot(label="3D Viewer")
 
-    def update_display(file, texture, uv_scale, light_intensity, ambient_intensity, transparency, color, num_inference_steps):
-        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, uv_scale, transparency)
+    def update_display(file, texture, uv_scale, uv_quality, light_intensity, ambient_intensity, transparency, color, num_inference_steps):
+        file_extension = file.split('.')[-1].lower()
+        texture_image = None
+
+        if texture:
+            texture_image = Image.open(texture)
 
-    submit_button.click(fn=update_display, inputs=[obj_file, texture_file, uv_scale_slider, light_intensity_slider, ambient_intensity_slider, transparency_slider, color_picker, num_inference_steps_slider], outputs=display)
+        if file_extension == 'vrm':
+            return display_3d_object(file, texture_image, light_intensity, ambient_intensity, color, uv_scale, transparency, uv_quality)
+        else:
+            return display_3d_object(file, texture_image, light_intensity, ambient_intensity, color, uv_scale, transparency)
+
+    def toggle_uv_quality_dropdown(file):
+        if file is None:
+            return gr.update(visible=False)
+
+        file_extension = file.split('.')[-1].lower()
+        return gr.update(visible=(file_extension == 'vrm'))
+
+    submit_button.click(
+        fn=update_display,
+        inputs=[obj_file, texture_file, uv_scale_slider, uv_quality_dropdown, light_intensity_slider, ambient_intensity_slider, transparency_slider, color_picker, num_inference_steps_slider],
+        outputs=display
+    )
+
+    obj_file.change(fn=toggle_uv_quality_dropdown, inputs=[obj_file], outputs=uv_quality_dropdown)
     generate_button.click(fn=update_texture_display, inputs=[prompt_input, texture_file, num_inference_steps_slider], outputs=texture_preview)
     restore_button.click(fn=restore_original, inputs=[obj_file], outputs=display)
     clear_button.click(fn=clear_texture, outputs=texture_preview)
     texture_file.change(fn=update_texture_display, inputs=[prompt_input, texture_file, num_inference_steps_slider], outputs=texture_preview)
 
-    demo.load(fn=update_display, inputs=[obj_file, texture_file, uv_scale_slider, light_intensity_slider, ambient_intensity_slider, transparency_slider, color_picker, num_inference_steps_slider], outputs=display)
+    demo.load(fn=update_display, inputs=[obj_file, texture_file, uv_scale_slider, uv_quality_dropdown, light_intensity_slider, ambient_intensity_slider, transparency_slider, color_picker, num_inference_steps_slider], outputs=display)
+
+
+
 
     gr.Examples(
-        examples=[[DEFAULT_OBJ_FILE, None],[DEFAULT_GLB_FILE, None]],
+        examples=[[DEFAULT_VRM_FILE, DEFAULT_TEXTURE],
+                  [DEFAULT_OBJ_FILE, None],
+                  [DEFAULT_GLB_FILE, None],
+                  [DEFAULT_VRM_FILE2, DEFAULT_TEXTURE2],
+                  [DEFAULT_VRM_FILE3, DEFAULT_TEXTURE3]],
         inputs=[obj_file, texture_file],
         label="Example Files"
     )
 
 demo.launch(debug=True)
+