app.py

import gradio as gr
import replicate
import openai
import trimesh
import numpy as np
from PIL import Image
import requests
import io
import tempfile
import os

# Set API tokens
os.environ["REPLICATE_API_TOKEN"] = "r8_Pc64F8EPrJ6PiNIIvaBUZcOGmiLC3Jp1gELYB"
# Initialize the Replicate client
rep_client = replicate.Client()

# Set your OpenAI API key
OPENAI_API_KEY = "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"
openai.api_key = OPENAI_API_KEY

# Initialize the Replicate client
rep_client = replicate.Client()

#--------------------------------2D Defect Simulator----------------------

predefined_defects = [
    "Missing bolts on railway track",
    "Cracks on railway track",
    "Overgrown vegetation near railway track",
    "Broken railings on railway bridge",
    "Debris on railway track",
    "Damaged railway platform"
]

# 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):
    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='gpt-3.5-turbo',
        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="ac732df83cea7fff18b8472768c88ad041fa750ff7682a21affe81863cbe77e4",
        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))
        image.save("defect.png")
        return image
    return None

# 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))
        image.save("defect.png")
        return image
    return None

# 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):
    # Load the original mesh
    mesh = trimesh.load('train.glb', force='mesh')
    rust_texture = Image.open('defect.png').convert('RGB')

    # Predefined sections
    if predefined_section == 'right compartments':
        selected_indices = np.where(mesh.vertices[:, 0] > (train_bounds[0][0] + train_bounds[1][0]) / 2)[0]
    elif predefined_section == 'left compartments':
        selected_indices = np.where(mesh.vertices[:, 0] <= (train_bounds[0][0] + train_bounds[1][0]) / 2)[0]
    elif predefined_section == 'freight_body':
        selected_indices = np.where((mesh.vertices[:, 0] >= train_bounds[0][0]) & (mesh.vertices[:, 0] <= train_bounds[1][0]) &
                                    (mesh.vertices[:, 2] <= (train_bounds[0][2] + train_bounds[1][2]) / 2))[0]
    elif predefined_section == 'custom':
        # Use custom sliders for custom section
        selected_indices = np.where((mesh.vertices[:, 0] >= x_min) & (mesh.vertices[:, 0] <= x_max) &
                                    (mesh.vertices[:, 1] >= y_min) & (mesh.vertices[:, 1] <= y_max) &
                                    (mesh.vertices[:, 2] >= z_min) & (mesh.vertices[:, 2] <= z_max))[0]
    else:
        selected_indices = np.array([])

    # Initialize UV mapping
    uv = np.random.rand(len(mesh.vertices), 2)
    new_uv = np.zeros_like(uv)
    new_uv[selected_indices, :] = uv[selected_indices, :]

    # Create material and apply the new texture
    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
    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.glb')
    textured_mesh.export(temp_file.name, file_type='glb')
    temp_file.close()
    return temp_file.name

# Load train model to establish bounding box ranges
train_model = trimesh.load('train.glb', force='mesh')
train_bounds = train_model.bounds

# Get slider ranges based on train model bounds
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 UI with a single window, dynamic updates, and real-time changes
with gr.Blocks() as app:
    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)

    gr.Markdown("### 3D Defect Simulator Tabs")
    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")
            model_output_predefined = gr.Model3D(label="3D Model with Applied 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]
            )

            visualize_button_predefined = gr.Button("Visualize 3D Model")
            visualize_button_predefined.click(
                fn=update_model,
                inputs=[section_dropdown, x_min_slider, x_max_slider, y_min_slider, y_max_slider, z_min_slider, z_max_slider],
                outputs=[model_output_predefined]
            )

        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")
            model_output_custom = gr.Model3D(label="3D Model with Applied 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]
            )

            visualize_button_custom = gr.Button("Visualize 3D Model")
            visualize_button_custom.click(
                fn=update_model,
                inputs=[section_dropdown, x_min_slider, x_max_slider, y_min_slider, y_max_slider, z_min_slider, z_max_slider],
                outputs=[model_output_custom]
            )

        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")
            model_output_inpaint = gr.Model3D(label="3D Model with Applied 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]
            )

            visualize_button_inpaint = gr.Button("Visualize 3D Model")
            visualize_button_inpaint.click(
                fn=update_model,
                inputs=[section_dropdown, x_min_slider, x_max_slider, y_min_slider, y_max_slider, z_min_slider, z_max_slider],
                outputs=[model_output_inpaint]
            )

        with gr.Tab("2D Defect Simulator"):
            with gr.Tabs():
                with gr.Tab("Current Defects"):
                    with gr.Row():
                        prompt_input = gr.Dropdown(choices=predefined_defects, label="Select a prompt")
                        number_input_dropdown = gr.Number(label="Number of images to generate", value=1, minimum=1, maximum=10)
                        submit_button_dropdown = gr.Button("Generate")
                    image_outputs_dropdown = gr.Gallery()

                    def on_submit_click_dropdown(prompt, number_of_images):
                        images = process_railway_defects(prompt, number_of_images)
                        return images

                    submit_button_dropdown.click(
                        fn=on_submit_click_dropdown,
                        inputs=[prompt_input, number_input_dropdown],
                        outputs=image_outputs_dropdown
                    )

                with gr.Tab("Custom Defect"):
                    with gr.Row():
                        custom_prompt_input = gr.Textbox(label="Custom Defect")
                        number_input_custom = gr.Number(label="Number of images to generate", value=1, minimum=1, maximum=10)
                        submit_button_custom = gr.Button("Generate")
                    image_outputs_custom = gr.Gallery()

                    def on_submit_click_custom(custom_prompt, number_of_images):
                        images = process_railway_defects(custom_prompt, number_of_images)
                        return images

                    submit_button_custom.click(
                        fn=on_submit_click_custom,
                        inputs=[custom_prompt_input, number_input_custom],
                        outputs=image_outputs_custom
                    )

app.launch()