Update app.py

app.py

@@ -5,57 +5,95 @@ import torch
 import os
 from pathlib import Path
 
-# Placeholder for the actual model loading and processing logic
-# Replace this with the actual code from generate.py
+# 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)
+
+# Load pre-trained models
 def load_model(model_path, use_cpu=False):
-    # Example: Load your ESRGAN model here
-    # This is a placeholder; adapt it based on the actual model loading in generate.py
-    model = torch.load(model_path)
+    model = torch.load(model_path, map_location="cpu" if use_cpu else None)
     if not use_cpu and torch.cuda.is_available():
         model = model.cuda()
     model.eval()
     return model
 
-def process_image(input_image, tile_size=512, seamless=False, use_cpu=False):
-    # Convert Gradio input (PIL image) to OpenCV format
-    img = np.array(input_image)
-    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+# 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 / "NormalMapGenerator-CX-Lite_200000_G.pth", use_cpu)
+    franken_model = load_model(MODELS_DIR / "frankenMapGenerator-CX-Lite_215000_G.pth", use_cpu)
 
-    # Load models (adjust paths to match your uploaded model files)
-    normal_model = load_model("models/normal_model.pth", use_cpu)
-    disp_model = load_model("models/displacement_model.pth", use_cpu)
-    rough_model = load_model("models/roughness_model.pth", use_cpu)
+    # Convert to tensor
+    img_tensor = torch.from_numpy(img.transpose(2, 0, 1)).float() / 255.0
+    if not use_cpu and torch.cuda.is_available():
+        img_tensor = img_tensor.cuda()
 
-    # Placeholder processing logic (replace with actual generate.py logic)
-    # For example, apply the model to the image
+    # Generate maps
     with torch.no_grad():
-        # Convert image to tensor
-        img_tensor = torch.from_numpy(img.transpose(2, 0, 1)).float() / 255.0
-        if not use_cpu and torch.cuda.is_available():
-            img_tensor = img_tensor.cuda()
+        # Normal map
+        normal_map = normal_model(img_tensor.unsqueeze(0)).cpu().numpy().squeeze()
+        # Franken map (contains Displacement and Roughness)
+        franken_map = franken_model(img_tensor.unsqueeze(0)).cpu().numpy().squeeze()
 
-        # Generate maps (simplified example)
-        normal_map = normal_model(img_tensor.unsqueeze(0)).cpu().numpy().squeeze().transpose(1, 2, 0)
-        disp_map = disp_model(img_tensor.unsqueeze(0)).cpu().numpy().squeeze().transpose(1, 2, 0)
-        rough_map = rough_model(img_tensor.unsqueeze(0)).cpu().numpy().squeeze().transpose(1, 2, 0)
+    # Post-process maps
+    # Normal map (RGB)
+    if normal_map.ndim == 3:
+        normal_map = normal_map.transpose(1, 2, 0)
+    else:
+        normal_map = np.stack([normal_map] * 3, axis=-1)  # Convert grayscale to RGB if needed
+    normal_map = (normal_map * 255).clip(0, 255).astype(np.uint8)
 
-    # Convert back to uint8 for display
-    normal_map = (normal_map * 255).astype(np.uint8)
-    disp_map = (disp_map * 255).astype(np.uint8)
-    rough_map = (rough_map * 255).astype(np.uint8)
+    # Franken map: Extract Displacement (red) and Roughness (green)
+    if franken_map.ndim == 3:
+        franken_map = franken_map.transpose(1, 2, 0)
+    # Displacement map (red channel, grayscale)
+    disp_map = franken_map[:, :, 0]  # Red channel
+    disp_map = (disp_map * 255).clip(0, 255).astype(np.uint8)
+    disp_map = np.stack([disp_map] * 3, axis=-1)  # Convert to RGB for Gradio display
 
-    # Convert to RGB for Gradio output
-    normal_map = cv2.cvtColor(normal_map, cv2.COLOR_BGR2RGB)
-    disp_map = cv2.cvtColor(disp_map, cv2.COLOR_BGR2RGB)
-    rough_map = cv2.cvtColor(rough_map, cv2.COLOR_BGR2RGB)
+    # Roughness map (green channel, grayscale)
+    rough_map = franken_map[:, :, 1]  # Green channel
+    rough_map = (rough_map * 255).clip(0, 255).astype(np.uint8)
+    rough_map = np.stack([rough_map] * 3, axis=-1)  # Convert to RGB for Gradio display
 
-    return normal_map, disp_map, rough_map
+    # 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"
 
-# Gradio interface
+    # 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):
-    normal_map, disp_map, rough_map = process_image(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=[
@@ -71,7 +109,7 @@ interface = gr.Interface(
         gr.Image(type="numpy", label="Roughness Map"),
     ],
     title="Material Map Generator",
-    description="Upload a diffuse texture to generate AI-generated Normal, Displacement, and Roughness maps."
+    description="Upload a diffuse texture to generate Normal, Displacement, and Roughness maps."
 )
 
 if __name__ == "__main__":