Update app.py

app.py

@@ -1,104 +1,14 @@
 import gradio as gr
-import replicate
-import openai
 import trimesh
 import numpy as np
 from PIL import Image
-import requests
-import io
 import tempfile
 
-# Set API tokens
-os.environ["REPLICATE_API_TOKEN"] = "r8_Pc64F8EPrJ6PiNIIvaBUZcOGmiLC3Jp1gELYB"
-OPENAI_API_KEY = "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"
-
-# Initialize the Replicate client
-rep_client = replicate.Client()
-
-# Material defects structure
-material_defects = {
-    "Steel": ["Rust and Corrosion", "Pitting Corrosion", "Surface Cracks", "Wear Patterns", "Spalling", "Scaling"],
-    "Glass": ["Cracks", "Chips", "Scratches", "Frosting"],
-    "Aluminum": ["Corrosion", "Scratches and Dents", "Anodizing Wear"],
-    "Wood": ["Rot and Decay", "Cracks and Splits", "Weathering"],
-    "Plastics and Polymers": ["Cracking and Crazing", "UV Degradation", "Heat Distortion"],
-    "Rubber": ["Cracking", "Hardening and Brittleness", "Surface Wear"],
-    "Composite Materials": ["Delamination", "Impact Damage", "Fiber Wearing"],
-    "Ceramics": ["Crackling", "Chipping and Pitting", "Glaze Deterioration"]
-}
-
-# Function to ask rail defect question
-def ask_rail_defect_question(question, model_name='ft:gpt-3.5-turbo-0125:personal::99NsSAeQ'):
-    structured_prompt = f"Translate the following user input into a concise, detailed visual description for a 3D model based on this input: '{question}'. Focus only on the defect’s appearance, texture qualities, and visual effects it would have on the material. Start the description directly with no extra words."
-    response = openai.ChatCompletion.create(
-        model=model_name,
-        messages=[
-            {"role": "system", "content": "Provide a concise, detailed visual description of the material's defect texture, focusing on visual and tactile qualities. Do not include any additional context or introductory phrases. Imagine the textures on railway components, but describe only the texture and material."},
-            {"role": "user", "content": structured_prompt}
-        ],
-    )
-    refined_description = response.choices[0].message['content']
-    return refined_description.strip()
-
-# Function to generate images from prompts
-def generate_images(prompt):
-    prediction = rep_client.predictions.create(
-        version="ac732df83cea7fff18b8472768c88ad041fa750ff7682a21affe81863cbe77e",
-        input={"prompt": prompt}
-    )
-    prediction.wait()
-    if prediction.status == "succeeded":
-        image_url = prediction.output[0]
-        response = requests.get(image_url)
-        image = Image.open(io.BytesIO(response.content))
-        return image
-    return "Failed to generate texture image"
-
-# Function to create data URL from PIL image
-def image_to_data_url(pil_image):
-    buffered = io.BytesIO()
-    pil_image.save(buffered, format="JPEG")
-    base64_image = base64.b64encode(buffered.getvalue()).decode('utf-8')
-    return f"data:image/jpeg;base64,{base64_image}"
-
-# Function to inpaint images
-def inpaint_texture(image, prompt):
-    if isinstance(image, np.ndarray):
-        image = Image.fromarray(image)
-    
-    image_data_url = image_to_data_url(image)
-       
-    input = {
-      "image": image_data_url,
-      "prompt": prompt,
-      "scheduler": "K_EULER_ANCESTRAL",
-      "num_outputs": 1,
-      "guidance_scale": 7.5,
-      "num_inference_steps": 100,
-      "image_guidance_scale": 1.5
-    }
-
-    prediction = rep_client.predictions.create(
-        version="30c1d0b916a6f8efce20493f5d61ee27491ab2a60437c13c588468b9810ec23f",
-        input = input
-    )
-    prediction.wait()
-    if prediction.status == "succeeded":
-        image_url = prediction.output[0]
-        response = requests.get(image_url)
-        image = Image.open(io.BytesIO(response.content))
-        return image
-    return "Failed to generate inpainted texture image"
-
-# Function to update defect options
-def update_defect_options(selected_material):
-    return gr.update(value='', choices=material_defects[selected_material])
-
 # Function to visualize texture based on selection criteria
-def visualize_dynamic_texture(predefined_section, x_min, x_max, y_min, y_max, z_min, z_max, custom_texture_path):
+def visualize_dynamic_texture(predefined_section, x_min, x_max, y_min, y_max, z_min, z_max):
     # Load the original mesh
     mesh = trimesh.load('train.glb', force='mesh')
-    custom_texture = Image.open(custom_texture_path).convert('RGB')
+    rust_texture = Image.open('rust_steel.png').convert('RGB')
 
     # Predefined sections
     if predefined_section == 'right compartments':
@@ -122,8 +32,8 @@ def visualize_dynamic_texture(predefined_section, x_min, x_max, y_min, y_max, z_
     new_uv[selected_indices, :] = uv[selected_indices, :]
 
     # Create material and apply the new texture
-    material = trimesh.visual.texture.SimpleMaterial(image=custom_texture)
-    color_visuals = trimesh.visual.TextureVisuals(uv=new_uv, image=custom_texture, material=material)
+    material = trimesh.visual.texture.SimpleMaterial(image=rust_texture)
+    color_visuals = trimesh.visual.TextureVisuals(uv=new_uv, image=rust_texture, material=material)
     textured_mesh = trimesh.Trimesh(vertices=mesh.vertices, faces=mesh.faces, visual=color_visuals, validate=True, process=False)
 
     # Save the mesh to a temporary file
@@ -141,91 +51,39 @@ x_min_range, x_max_range = train_bounds[0][0], train_bounds[1][0]
 y_min_range, y_max_range = train_bounds[0][1], train_bounds[1][1]
 z_min_range, z_max_range = train_bounds[0][2], train_bounds[1][2]
 
-# Gradio app interface setup
+# Gradio UI with a single window, dynamic updates, and real-time changes
 with gr.Blocks() as app:
-    with gr.Tabs():
-        with gr.Tab("3D Defect Simulator"):
-            with gr.Tabs():
-                with gr.Tab("Predefined Defect Texture"):
-                    with gr.Row():
-                        material_input = gr.Dropdown(choices=list(material_defects.keys()), label="Select Material")
-                        defect_input = gr.Dropdown(choices=[], label="Select Defect Type")
-                        generate_button = gr.Button("Generate Texture")
-                    image_output = gr.Image(label="Generated Texture")
-
-                    material_input.change(fn=update_defect_options, inputs=[material_input], outputs=[defect_input])
-                    generate_button.click(
-                        fn=lambda material, defect: generate_images(ask_rail_defect_question(f"Describe the texture of {defect} on {material}")),
-                        inputs=[material_input, defect_input],
-                        outputs=[image_output]
-                    )
-
-                with gr.Tab("Custom Defect Texture"):
-                    with gr.Row():
-                        custom_prompt_input = gr.Textbox(label="Enter Custom Prompt for Texture", placeholder="Describe any texture detail you need.")
-                        refine_button = gr.Button("Refine Prompt")
-                    refined_prompt_output = gr.Textbox(label="Refined Prompt", placeholder="This will show the refined prompt.")
-
-                    with gr.Row():
-                        generate_button = gr.Button("Generate Texture")
-                    custom_image_output = gr.Image(label="Generated Texture")
-
-                    # Refine the input prompt
-                    refine_button.click(
-                        fn=lambda prompt: ask_rail_defect_question(prompt),
-                        inputs=[custom_prompt_input],
-                        outputs=[refined_prompt_output]
-                    )
-
-                    # Use the refined prompt to generate the texture image
-                    generate_button.click(
-                        fn=lambda prompt: generate_images(prompt),
-                        inputs=[refined_prompt_output],
-                        outputs=[custom_image_output]
-                    )
-
-                with gr.Tab("Inpaint Defect Texture"):
-                    with gr.Row():
-                        image_input = gr.Image(label="Upload Image for Inpainting")
-                        inpaint_prompt_input = gr.Textbox(label="Enter Prompt for Texture Inpainting")
-                        inpaint_button = gr.Button("Generate Inpainted Texture")
-                    inpaint_image_output = gr.Image(label="Generated Inpainted Texture")
-
-                    # Use the images and prompt to generate the inpainted texture image
-                    inpaint_button.click(
-                        fn=lambda img, prompt: inpaint_texture(img, prompt),
-                        inputs=[image_input, inpaint_prompt_input],
-                        outputs=[inpaint_image_output]
-                    )
-
-                with gr.Tab("Apply Custom Texture"):
-                    with gr.Row():
-                        predefined_section = gr.Radio(choices=['right compartments', 'left compartments', 'freight_body', 'custom'], label="Select Section", value='custom')
-                        custom_texture_path = gr.Textbox(label="Path to Custom Texture Image")
-
-                    with gr.Row(visible=True) as custom_controls:
-                        x_min_slider = gr.Slider(minimum=x_min_range, maximum=x_max_range, step=0.01, label="X Min", value=x_min_range)
-                        x_max_slider = gr.Slider(minimum=x_min_range, maximum=x_max_range, step=0.01, label="X Max", value=x_max_range)
-
-                        y_min_slider = gr.Slider(minimum=y_min_range, maximum=y_max_range, step=0.01, label="Y Min", value=y_min_range)
-                        y_max_slider = gr.Slider(minimum=y_min_range, maximum=y_max_range, step=0.01, label="Y Max", value=y_max_range)
-
-                        z_min_slider = gr.Slider(minimum=z_min_range, maximum=z_max_range, step=0.01, label="Z Min", value=z_min_range)
-                        z_max_slider = gr.Slider(minimum=z_min_range, maximum=z_max_range, step=0.01, label="Z Max", value=z_max_range)
-
-                    # Toggle visibility of custom controls
-                    def toggle_custom_controls(predefined_section):
-                        return gr.update(visible=(predefined_section == 'custom'))
-
-                    predefined_section.change(fn=toggle_custom_controls, inputs=[predefined_section], outputs=custom_controls)
-
-                    # Update model dynamically
-                    def update_model(predefined_section, x_min, x_max, y_min, y_max, z_min, z_max, custom_texture_path):
-                        return visualize_dynamic_texture(predefined_section, x_min, x_max, y_min, y_max, z_min, z_max, custom_texture_path)
-
-                    # Add event listeners for real-time updates when sliders or dropdown change
-                    inputs = [predefined_section, x_min_slider, x_max_slider, y_min_slider, y_max_slider, z_min_slider, z_max_slider, custom_texture_path]
-                    for input_component in inputs:
-                        input_component.change(fn=update_model, inputs=inputs, outputs=modified_model)
-
-app.launch()
+    gr.Markdown("### 3D Model Texture Application with Predefined and Custom Sections")
+    original_model = gr.Model3D('train.glb', label="Original Model")
+    modified_model = gr.Model3D(label="Model with Applied Texture")
+
+    # Dropdown for predefined and custom selection
+    section_dropdown = gr.Radio(choices=['right compartments', 'left compartments', 'freight_body', 'custom'], label="Select Section", value='custom')
+
+    # Custom sliders for the bounding box selection
+    with gr.Row(visible=True) as custom_controls:
+        x_min_slider = gr.Slider(minimum=x_min_range, maximum=x_max_range, step=0.01, label="X Min", value=x_min_range)
+        x_max_slider = gr.Slider(minimum=x_min_range, maximum=x_max_range, step=0.01, label="X Max", value=x_max_range)
+
+        y_min_slider = gr.Slider(minimum=y_min_range, maximum=y_max_range, step=0.01, label="Y Min", value=y_min_range)
+        y_max_slider = gr.Slider(minimum=y_min_range, maximum=y_max_range, step=0.01, label="Y Max", value=y_max_range)
+
+        z_min_slider = gr.Slider(minimum=z_min_range, maximum=z_max_range, step=0.01, label="Z Min", value=z_min_range)
+        z_max_slider = gr.Slider(minimum=z_min_range, maximum=z_max_range, step=0.01, label="Z Max", value=z_max_range)
+
+    # Toggle visibility of custom controls
+    def toggle_custom_controls(predefined_section):
+        return gr.update(visible=(predefined_section == 'custom'))
+
+    section_dropdown.change(fn=toggle_custom_controls, inputs=section_dropdown, outputs=custom_controls)
+
+    # Update model dynamically
+    def update_model(predefined_section, x_min, x_max, y_min, y_max, z_min, z_max):
+        return visualize_dynamic_texture(predefined_section, x_min, x_max, y_min, y_max, z_min, z_max)
+
+    # Add event listeners for real-time updates when sliders or dropdown change
+    inputs = [section_dropdown, x_min_slider, x_max_slider, y_min_slider, y_max_slider, z_min_slider, z_max_slider]
+    for input_component in inputs:
+        input_component.change(fn=update_model, inputs=inputs, outputs=modified_model)
+
+app.launch()