# Material Transfer System (Style2Shape)

<p align="center">
  <img src="assets/images/chair-000/0.jpg" width="800" alt="Material Transfer Preview">
</p>

## Abstract

This paper presents Style2Shape, a novel framework for generating physically-based rendering (PBR) materials for 3D models from single reference images. Recovering PBR materials from single reference images is inherently ill-posed, with complex interactions between geometry, material properties, and lighting creating fundamental ambiguities—bright surfaces could result from high albedo or strong illumination.

Our key insight is that procedural materials and generated textures are complementary—procedural parameters ensure correct reflectance properties (roughness, metallic, etc.), while generated textures capture arbitrary visual patterns. By optimizing their learnable combination in a physics-based differentiable rendering pipeline, we achieve materials that are both physically plausible and visually faithful to reference images.

The framework operates in three stages: (1) Structure-guided appearance transfer, reformulating appearance transfer as an image-to-image translation problem, synthesizing geometrically aligned supervision images; (2) Hybrid PBR material initialization, retrieving procedural materials based on physical properties and generating complementary textures for visual patterns; (3) Physics-based differentiable optimization, jointly optimizing procedural parameters, generated textures, and their blending in a differentiable rendering framework.

Key contributions include: a unified framework from single image input to production-ready PBR materials; hybrid material representation combining procedural material reflectance correctness and generated texture arbitrary appearance; physics-based optimization framework leveraging differentiable rendering to jointly optimize all material components while maintaining physical plausibility.

## Overview

Style2Shape addresses the fundamental challenge of recovering PBR materials from single reference images. Current methods have significant limitations: text-driven methods suffer from descriptive ambiguity, image-based methods produce non-editable implicit representations, and procedural materials are constrained by predefined parameter spaces unable to capture arbitrary real-world appearances.

Our approach addresses these challenges through hybrid material optimization:

1. **Physical Plausibility**
   - Procedural materials ensure correct reflectance properties (roughness, metallic, etc.)
   - Physics-based differentiable rendering pipeline
   - Maintain physical constraints on material properties

2. **Visual Fidelity**
   - Generated textures capture arbitrary visual patterns
   - Geometrically aligned supervision signal synthesis
   - Faithful reproduction of reference image appearance

3. **Editability**
   - Production-ready PBR material output
   - Controllable adjustment of procedural parameters
   - Learnable combination of hybrid materials

## Environment Setup

### Using Conda
Create and activate the environment using the provided `env.yaml` file:
```bash
conda env create -f env.yaml
conda activate pbrtransfer
```

### Pre-trained Model Downloads

Style2Shape requires the following pre-trained models. Please download and update the paths in the code:

#### 1. RmBG 2.0 Model
**Purpose**: Automatic background removal for reference images
**Download**: [RmBG 2.0 Model](https://github.com/danielgatis/rembg/releases/download/v2.0.50/u2net.onnx)
```bash
mkdir -p models/rmbg
wget https://github.com/danielgatis/rembg/releases/download/v2.0.50/u2net.onnx -O models/rmbg/u2net.onnx

# Or manually download and place in models/rmbg/u2net.onnx
```

**Path Configuration**:
Update the model path in `rmbg.py`:
```python
# In rmbg.py, update the model path
MODEL_PATH = "models/rmbg/u2net.onnx"  # Update this path
```

#### 2. RGB-X Decomposition Model
**Purpose**: Intrinsic decomposition for albedo, normal, roughness, and metallic maps
**Download**: [Stable Diffusion RGB-X Model](https://huggingface.co/stabilityai/stable-diffusion-2-1)
```bash
pip install diffusers transformers

# Models will be automatically downloaded on first use
# Or manually download from Hugging Face
```

**Path Configuration**:
Update the model path in `intrinsic_decompose.py`:
```python
# In intrinsic_decompose.py, update the model path
MODEL_NAME = "stabilityai/stable-diffusion-2-1"  # Update this path if needed
```

#### 3. SBS Material Library
**Purpose**: Procedural material library
**Download**: [DiffMat SBS Material Library](https://github.com/mit-gfx/diffmat/tree/master/test/sbs/match_v1)
```bash
git clone https://github.com/mit-gfx/diffmat.git
cd diffmat/test/sbs/match_v1

# Copy SBS materials to your material library
cp -r * /path/to/your/material/library/

# Or download specific materials as needed
```

**Path Configuration**:
Update the material library path in `config.py`:
```python
# In config.py, update the material library path
MATERIAL_ROOT_DIR = "/path/to/your/material/library"  # Update this path
```

**More SBS Materials**: For additional SBS materials, please refer to the [DiffMat SBS Material Library](https://github.com/mit-gfx/diffmat/tree/master/test/sbs/match_v1), which provides a comprehensive collection of procedural materials suitable for various use cases.

#### 4. Model Directory Structure
After downloading, ensure the following directory structure:
```
ttt/
├── models/
│   ├── rmbg/u2net.onnx
│   └── rgbx/stable-diffusion-2-1/
├── material_library/
│   ├── metal/
│   ├── fabric/
│   └── wood/
└── material_transfer/
```

#### 5. Path Configuration
Update model paths in the respective files:
```python
# rmbg.py
MODEL_PATH = "models/rmbg/u2net.onnx"

# intrinsic_decompose.py  
MODEL_NAME = "stabilityai/stable-diffusion-2-1"

# config.py
MATERIAL_ROOT_DIR = "/path/to/your/material/library"
```

## Directory Structure

### Key Directories and Files
- **./inverse_rendering.py**: Main material transfer pipeline
  - `run_inverse_rendering()`: Main function for processing material optimization examples
  - `MaterialOptimizer`: Core optimization class managing scene setup and rendering
  - Implements complete inverse rendering loop
  - Manages material parameter updates and loss function evaluation
  - Handles seamless texture generation and material matching

- **./inverse_rendering_helper.py**: Core inverse rendering implementation
  - `inverse_rendering_func()`: Main optimization function supporting DrJit
  - `MaterialOptimizer`: Advanced optimizer with multi-stage functionality
  - Particle to material parameter transfer
  - Mesh-based material property updates and constitutive model evaluation

- **./intrinsic_decompose.py**: Material decomposition module
  - `intrinsic_decomp()`: RGB to material property decomposition
  - Stable Diffusion AOV generation for albedo, normal, roughness, metallic
  - Supports multiple decomposition samples and quality control

- **./config/**: Configuration files
  - `config.py`: Path configuration for material library, 3D models, and transfer data
  - Each experiment defines scene parameters, material properties, and optimization settings

- **./utils/**: Utility functions
  - `utils.py`: Image processing and material processing tools
  - `alignment.py`: Affine transformation alignment from image to mask

- **./split_mesh.py**: Mesh preprocessing module
  - `split_mesh()`: Main function for mesh segmentation and component separation
  - `identify_components()`: Component identification based on connectivity analysis
  - `assign_materials()`: Material assignment for separated components
  - Handles decomposition of complex 3D models for individual material transfer

- **./render_mesh.py**: Mesh rendering module
  - `render_mesh()`: Main rendering function for mesh components
  - `generate_normal_map()`: Normal map generation for each component
  - `generate_masks()`: Mask generation for material transfer regions
  - `extract_textures()`: Texture extraction and processing
  - Generates required data format for TRANSFER_DATA_DIR

- **./rmbg.py**: Background removal module
  - `remove_background()`: Main function for automatic background removal
  - Uses RmBG 2.0 pre-trained model (u2net.onnx)
  - Supports multiple image formats and batch processing
  - Integrates with reference image preprocessing pipeline

- **./intrinsic_decompose.py**: Material decomposition module
  - `intrinsic_decomp()`: RGB to material property decomposition
  - Uses Stable Diffusion RGB-X model for AOV generation
  - Generates albedo, normal, roughness, and metallic maps
  - Supports multiple decomposition samples and quality control

- **./pair_materials.py**: Material matching module
  - `pair_materials()`: Main function for material matching and synthesis
  - Uses DiffMat SBS material library for procedural material generation
  - Supports similarity matching based on texture descriptors
  - Integrates with material optimization pipeline

## Configuration Parameters

### Material Transfer Environment
- `resolution`: Rendering resolution (default: 512x512)
- `total_iter`: Maximum optimization iterations
- `eval_step`: Evaluation frequency during optimization
- `loss_type`: Loss function type (mae, mse)
- `use_tv_loss`: Enable total variation regularization
- `tv_loss_weight`: TV loss weight coefficient
- Camera parameters:
  - `camera_position`: Camera position in 3D space
  - `camera_target`: Camera target point
  - `camera_up`: Camera up vector
- `device`: Computing device (cuda, cpu)

### Material Properties
- `albedo`: Diffuse base color texture
- `normal`: Surface normal map
- `roughness`: Surface roughness map
- `metallic`: Metallic property map
- `material_category`: Material type classification
- `material_name`: Specific material identifier

### Optimization Settings
- `learning_rate`: Gradient descent learning rate
- `max_steps`: Maximum optimization steps
- `fw_spp`: Forward rendering samples per pixel
- `bw_spp`: Backward rendering samples per pixel

### Preprocessing Settings
- `min_component_size`: Minimum vertex count per component for mesh segmentation
- `connectivity_threshold`: Connectivity threshold for component separation
- `material_preservation`: Preserve original materials during segmentation
- `uv_generation`: Generate UV coordinates for separated components
- `normal_map_generation`: Enable normal map generation
- `mask_generation`: Enable component mask generation
- `texture_extraction`: Enable texture data extraction
- `best_view_detection`: Enable best view detection

### Model Paths
- `RMBG_MODEL_PATH`: RmBG 2.0 model path (u2net.onnx)
- `RGBX_MODEL_PATH`: RGB-X decomposition model path
- `STABLE_DIFFUSION_MODEL`: Stable Diffusion model name or path
- `MODEL_CACHE_DIR`: Directory for caching downloaded models
- `MATERIAL_ROOT_DIR`: SBS material library path from DiffMat
- `SBS_MATERIAL_CATEGORIES`: Available material categories (metal, fabric, wood, etc.)

## Style2Shape Extension

Style2Shape addresses the fundamental challenge of recovering PBR materials from single images through hybrid material optimization. Our framework consists of three key stages:

### Stage 1: Structure-guided Appearance Transfer
- Reformulates appearance transfer as an image-to-image translation problem
- Synthesizes geometrically aligned supervision images, maintaining reference appearance while matching target model geometry
- Resolves domain gaps between arbitrary reference images and target 3D geometry

### Stage 2: Hybrid PBR Material Initialization
- Retrieves procedural materials based on physical properties
- Generates complementary textures for visual patterns
- Combines procedural materials and generated textures through learnable blending weights

### Stage 3: Physics-based Differentiable Optimization
- Jointly optimizes procedural parameters, generated textures, and their blending in a differentiable rendering framework
- Ensures convergence to physically correct and visually accurate materials
- Maintains physical plausibility while achieving visual fidelity

### Technical Advantages
- **Unified Framework**: From single image input to production-ready PBR materials
- **Hybrid Representation**: Combines procedural material reflectance correctness and generated texture arbitrary appearance
- **Physics-based Optimization**: Leverages differentiable rendering to jointly optimize all material components

## Running Material Transfer

### Method 1: Automatic Seamless Texture Generation

Style2Shape supports automatic seamless texture generation without manual upload. The system automatically calls APIs to generate seamless textures based on specified target objects.
However, you need to configure the relevant api_key in the files.

#### Configuration Parameter Description

In the `examples` configuration:
- `model_name`: Model name (e.g., "chair", "bag")
- `model_index`: Model index (e.g., "000", "005")
- `gt_idx`: Index of GPT-generated image to use (corresponding to 0.jpg, 1.jpg, etc.)

In `target_obj_lists`:
- Specify the model component indices for which you want to generate seamless textures
- Format as a 2D list, where each sublist corresponds to target objects for one experiment

#### Basic Usage Example

```python
from inverse_rendering import run_inverse_rendering

# Define experiment configuration
examples = [
    {
        "model_name": "bag",
        "model_index": "005", 
        "gt_idx": 2  # Use GPT-generated image index 2
    }
]

# Specify target object list (only generate seamless textures for these objects)
target_obj_lists = [
    [0, 1]  # Generate seamless textures for objects 0 and 1
]

# Run automatic generation mode
run_inverse_rendering(
    examples=examples,
    auto_generate=True,  # Enable automatic generation
    target_obj_lists=target_obj_lists
)
```

#### Command Line Execution

```bash
python inverse_rendering.py
```

Set in `inverse_rendering.py`:
```python
# Configuration options
AUTO_GENERATE = True  # Set to True for automatic generation

# User-specified target object list
target_obj_lists = [
    [0, 1]  # Experiment 0: Generate seamless textures for objects 0 and 1
]
```

#### Automatic Generation Process

1. **Material Region Cropping**: System automatically crops material regions from final_render based on masks
2. **API Call**: Calls GPT-Image-1 API to generate seamless textures
3. **Texture Saving**: Saves generated textures in seamless folder
4. **Material Transfer**: Automatically applies to material transfer process

#### API Configuration

Configure your API key in `api_generate.py`:
```python
client = OpenAI(
    api_key="your-api-key-here", 
    base_url="https://api.laozhang.ai/v1"
)
```

### Method 2: Manual Upload Configuration

If you already have seamless textures, you can choose to upload manually:

1. **GPT-generated reference images**: Place in `assets/images/{model_name}-{model_index}/` directory, named as `0.jpg`, `1.jpg`, etc.
2. **Seamless textures**: In asset/images/{model_name}-{model}-{mode_index}/ directory, named as: seamless_1, seamless_2, etc.

#### Manual Configuration Example

```python
# Seamless texture configuration
seamless_texture_idx_list = [
    [1, 3],  # Apply seamless textures to components 1 and 3
]

seamless_texture_list = [
    [5, 5],  # Both components 1 and 3 use seamless_5 texture
]

# Run manual mode
run_inverse_rendering(
    examples=examples,
    seamless_texture_idx_list=seamless_texture_idx_list,
    seamless_texture_list=seamless_texture_list,
    auto_generate=False  # Use manual specification
)
```

### Configuration Parameter Description

#### Automatic Generation Mode Parameters

In the `examples` configuration:
- `model_name`: Model name (e.g., "chair", "bag")
- `model_index`: Model index (e.g., "000", "005")
- `gt_idx`: Index of GPT-generated image to use (corresponding to 0.jpg, 1.jpg, etc.)

In `target_obj_lists`:
- Specify the model component indices for which you want to generate seamless textures
- Format as a 2D list, where each sublist corresponds to target objects for one experiment
- Example: `[[0, 1], [0], [0, 1, 2]]` means three experiments process objects [0,1], [0], [0,1,2] respectively

#### Manual Mode Parameters

In `seamless_texture_idx_list`:
- Specify the model component indices to which you want to apply seamless textures

In `seamless_texture_list`:
- Specify the seamless image indices used by the corresponding model components

### Basic Command Structure
```bash
python inverse_rendering.py \
  --examples config/experiments.json \
  --auto_generate \
  --target_obj_lists [[0,1],[0],[0,1,2]]
```

### Detailed Examples: Automatic Seamless Texture Generation

#### Example 1: Single Experiment Automatic Generation
```python
from inverse_rendering import run_inverse_rendering

# Define experiment configuration
examples = [
    {
        "model_name": "bag",
        "model_index": "005", 
        "gt_idx": 2  # Use GPT-generated image index 2
    }
]

# Specify target object list (only generate seamless textures for these objects)
target_obj_lists = [
    [0, 1]  # Generate seamless textures for objects 0 and 1
]

# Run automatic generation mode
run_inverse_rendering(
    examples=examples,
    auto_generate=True,  # Enable automatic generation
    target_obj_lists=target_obj_lists
)
```

#### Example 2: Multi-experiment Batch Processing
```python
# Define multiple experiment configurations
examples = [
    {
        "model_name": "bag",
        "model_index": "005", 
        "gt_idx": 2
    },
    {
        "model_name": "chair",
        "model_index": "000", 
        "gt_idx": 1
    },
    {
        "model_name": "bed",
        "model_index": "009", 
        "gt_idx": 3
    }
]

# Specify different target objects for each experiment
target_obj_lists = [
    [0, 1],    # Experiment 0: Generate seamless textures for objects 0 and 1
    [0],       # Experiment 1: Only generate seamless texture for object 0
    [0, 1, 2]  # Experiment 2: Generate seamless textures for objects 0, 1, 2
]

# Run batch automatic generation
run_inverse_rendering(
    examples=examples,
    auto_generate=True,
    target_obj_lists=target_obj_lists
)
```

#### Example 3: Manual Mode (Existing Seamless Textures)
```python
# Seamless texture configuration
seamless_texture_idx_list = [
    [1, 3],  # Apply seamless textures to components 1 and 3
]

seamless_texture_list = [
    [5, 5],  # Both components 1 and 3 use seamless_5 texture
]

# Run manual mode
run_inverse_rendering(
    examples=examples,
    seamless_texture_idx_list=seamless_texture_idx_list,
    seamless_texture_list=seamless_texture_list,
    auto_generate=False  # Use manual specification
)
```

### Material Transfer Results
The system generates optimized material textures and final rendered images in the output directory.

## Data Preprocessing

### Overview
Style2Shape transforms raw 3D models and reference images into optimization-ready material systems through structure-guided appearance transfer and hybrid material initialization.

### Key Steps
1. **Geometrically Aligned Supervision Synthesis**
   - Use `utils.py` tools to load mesh geometry from OBJ files
   - Generate geometrically aligned supervision images, maintaining reference appearance
   - Resolve domain gaps between reference images and target 3D geometry

2. **Hybrid Material Initialization**
   - Retrieve procedural materials based on physical properties
   - Generate complementary textures for visual patterns
   - Combine procedural materials and generated textures through learnable blending weights

3. **Differentiable Rendering Preparation**
   - Initialize differentiable rendering pipeline
   - Set up joint optimization of procedural parameters and generated textures
   - Configure physical constraints and visual fidelity objectives

4. **Optimization Framework Configuration**
   - Initialize camera parameters and lighting settings
   - Pre-compute material optimization constraints
   - Set up joint optimization objectives for procedural parameters and generated textures

## Mesh Preprocessing Pipeline

### Overview
Before material transfer, each 3D model must go through a preprocessing pipeline to generate the data format required by `TRANSFER_DATA_DIR`. This pipeline consists of two main steps: mesh segmentation and mesh rendering.

### Step 1: Mesh Segmentation (`split_mesh.py`)
**Purpose**: Segment complex 3D models into separate components for individual material assignment.

**Input**: Raw 3D model files (OBJ format)
**Output**: Separated mesh components with individual material assignments

**Key Functions**:
- `split_mesh()`: Main function for mesh segmentation
- `identify_components()`: Component identification based on connectivity
- `assign_materials()`: Material assignment for separated components

**Usage Example**:
```python
from split_mesh import split_mesh

# Segment complex chair model into components
split_mesh(
    input_path="assets/model/chair/000/obj_0.obj",
    output_dir="assets/images/chair-000/",
    model_name="chair",
    model_index="000"
)
```

### Step 2: Mesh Rendering (`render_mesh.py`)
**Purpose**: Generate normal maps, masks, and texture data for each separated component.

**Input**: Separated mesh components from Step 1
**Output**: Normal maps, masks, and texture data in `TRANSFER_DATA_DIR` format

**Key Functions**:
- `render_mesh()`: Main rendering function for mesh components
- `generate_normal_map()`: Normal map generation for each component
- `generate_masks()`: Mask generation for material transfer regions
- `extract_textures()`: Texture extraction and processing

**Usage Example**:
```python
from render_mesh import render_mesh

# Render separated mesh components
render_mesh(
    mesh_dir="assets/images/chair-000/",
    output_dir="assets/images/chair-000/",
    resolution=512,
    generate_normal=True,
    generate_masks=True
)
```

### Required Output Format for TRANSFER_DATA_DIR

The preprocessing pipeline must generate the following structure in `TRANSFER_DATA_DIR`:

```
assets/images/{model_name}-{model_index}/
├── 0.jpg                    # GPT-generated reference image 0
├── 1.jpg                    # GPT-generated reference image 1
├── 2.jpg                    # GPT-generated reference image 2
├── 3.jpg                    # GPT-generated reference image 3
├── normal.png               # Normal map of the entire model
├── mask.png                 # Overall mask of the model
├── mask_0.png              # Mask of component 0
├── mask_1.png              # Mask of component 1
├── mask_2.png              # Mask of component 2
├── mask_3.png              # Mask of component 3
├── best_view.json          # Best view parameters
├── seamless_0             # Seamless texture 0
├── seamless_1/             # Seamless texture 1
   
```

### Complete Preprocessing Workflow

```python
import os
from split_mesh import split_mesh
from render_mesh import render_mesh

def preprocess_model(model_name, model_index, input_model_path):
    """
    Complete preprocessing workflow for 3D models.
    
    Parameters:
        model_name: Model name (e.g., "chair")
        model_index: Model index (e.g., "000")
        input_model_path: Path to input OBJ file
    """
    
    # Step 1: Segment mesh into components
    output_dir = f"assets/images/{model_name}-{model_index}/"
    os.makedirs(output_dir, exist_ok=True)
    
    split_mesh(
        input_path=input_model_path,
        output_dir=output_dir,
        model_name=model_name,
        model_index=model_index
    )
    
    # Step 2: Render separated components
    render_mesh(
        mesh_dir=output_dir,
        output_dir=output_dir,
        resolution=512,
        generate_normal=True,
        generate_masks=True,
        extract_textures=True
    )
    
    print(f"Preprocessing completed: {model_name}-{model_index}")

# Usage example
preprocess_model("chair", "000", "assets/model/chair/000/obj_0.obj")
```

### Integration with Material Transfer

After preprocessing, the data generated in `TRANSFER_DATA_DIR` can be directly used by the material transfer system:

```python
from inverse_rendering import run_inverse_rendering

# Preprocessed data can now be used for material transfer
examples = [
    {
        "model_name": "chair",
        "model_index": "000",
        "gt_idx": 3  # Reference image index
    }
]

run_inverse_rendering(
    examples=examples,
    auto_generate=True,
    target_obj_lists=[[0, 1, 2]]  # Target components from preprocessing
)
```
# Citation

If you use Style2Shape in your research, please cite:
```bibtex
@article{style2shape2024,
  title={Style2Shape: Material Transfer System Based on Inverse Rendering},
  author={Your Name and Collaborators},
  journal={arXiv preprint arXiv:XXXX.XXXXX},
  year={2024}
}
```
