File size: 4,524 Bytes
f3436bc 7e25f9d fe62c07 f3436bc fe62c07 f3436bc 7e25f9d 32400ac f3436bc 7e25f9d fe62c07 f3436bc 7e25f9d f3436bc fe62c07 f3436bc 7e25f9d fe62c07 f3436bc 7e25f9d f3436bc 7e25f9d f3436bc 7e25f9d f3436bc 7e25f9d f3436bc 7e25f9d f3436bc fe62c07 | 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 | import gradio as gr
import cv2
import numpy as np
import torch
import os
from pathlib import Path
# Ensure directories exist
INPUT_DIR = Path("input")
OUTPUT_DIR = Path("output")
MODELS_DIR = Path("models")
INPUT_DIR.mkdir(exist_ok=True)
OUTPUT_DIR.mkdir(exist_ok=True)
# Function to load pre-trained models
def load_model(model_path, use_cpu=False):
if not model_path.exists():
raise FileNotFoundError(f"Model file not found: {model_path}")
device = "cpu" if use_cpu or not torch.cuda.is_available() else "cuda"
# Load state_dict if the model was saved that way
model_state = torch.load(model_path, map_location=device)
# If a full model object was saved, load it directly
if isinstance(model_state, torch.nn.Module):
model = model_state
else:
# If saved as state_dict, we need a model architecture (Assuming CNN or custom model)
model = torch.nn.Sequential(
torch.nn.Conv2d(3, 64, kernel_size=3, padding=1),
torch.nn.ReLU(),
torch.nn.Conv2d(64, 3, kernel_size=3, padding=1)
)
model.load_state_dict(model_state)
model.to(device)
model.eval()
return model
# Process image and save to output folder
def process_image(input_path, tile_size=512, seamless=False, use_cpu=False):
# Read input image
img = cv2.imread(str(input_path))
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
# Load models
normal_model = load_model(MODELS_DIR / "1x_NormalMapGenerator-CX-Lite_200000_G.pth", use_cpu)
franken_model = load_model(MODELS_DIR / "1x_frankenMapGenerator-CX-Lite_215000_G.pth", use_cpu)
# Convert to tensor
img_tensor = torch.from_numpy(img.transpose(2, 0, 1)).float() / 255.0
img_tensor = img_tensor.unsqueeze(0) # Add batch dimension
device = "cpu" if use_cpu or not torch.cuda.is_available() else "cuda"
img_tensor = img_tensor.to(device)
# Generate maps
with torch.no_grad():
normal_map = normal_model(img_tensor).cpu().numpy().squeeze()
franken_map = franken_model(img_tensor).cpu().numpy().squeeze()
# Post-process maps
normal_map = (normal_map.transpose(1, 2, 0) * 255).clip(0, 255).astype(np.uint8)
disp_map = (franken_map[0] * 255).clip(0, 255).astype(np.uint8)
rough_map = (franken_map[1] * 255).clip(0, 255).astype(np.uint8)
# Convert grayscale to RGB for Gradio display
disp_map = np.stack([disp_map] * 3, axis=-1)
rough_map = np.stack([rough_map] * 3, axis=-1)
# Define output paths
base_name = input_path.stem
normal_path = OUTPUT_DIR / f"{base_name}_normal.png"
disp_path = OUTPUT_DIR / f"{base_name}_displacement.png"
rough_path = OUTPUT_DIR / f"{base_name}_roughness.png"
# Save outputs
cv2.imwrite(str(normal_path), cv2.cvtColor(normal_map, cv2.COLOR_RGB2BGR))
cv2.imwrite(str(disp_path), cv2.cvtColor(disp_map, cv2.COLOR_RGB2BGR))
cv2.imwrite(str(rough_path), cv2.cvtColor(rough_map, cv2.COLOR_RGB2BGR))
return normal_path, disp_path, rough_path
# Gradio function
def generate_maps(input_image, tile_size, seamless, use_cpu):
# Save uploaded image to input folder
input_path = INPUT_DIR / "uploaded_texture.png"
input_img = np.array(input_image)
cv2.imwrite(str(input_path), cv2.cvtColor(input_img, cv2.COLOR_RGB2BGR))
# Process image
normal_path, disp_path, rough_path = process_image(input_path, tile_size, seamless, use_cpu)
# Read outputs for display
normal_map = cv2.cvtColor(cv2.imread(str(normal_path)), cv2.COLOR_BGR2RGB)
disp_map = cv2.cvtColor(cv2.imread(str(disp_path)), cv2.COLOR_BGR2RGB)
rough_map = cv2.cvtColor(cv2.imread(str(rough_path)), cv2.COLOR_BGR2RGB)
return input_image, normal_map, disp_map, rough_map
# Gradio interface
interface = gr.Interface(
fn=generate_maps,
inputs=[
gr.Image(type="pil", label="Diffuse Texture"),
gr.Slider(minimum=256, maximum=1024, step=64, value=512, label="Tile Size"),
gr.Checkbox(label="Seamless", value=False),
gr.Checkbox(label="Use CPU", value=False),
],
outputs=[
gr.Image(type="numpy", label="Input Diffuse Texture"),
gr.Image(type="numpy", label="Normal Map"),
gr.Image(type="numpy", label="Displacement Map"),
gr.Image(type="numpy", label="Roughness Map"),
],
title="Material Map Generator",
description="Upload a diffuse texture to generate Normal, Displacement, and Roughness maps."
)
if __name__ == "__main__":
interface.launch()
|