Create app.py

app.py

@@ -0,0 +1,131 @@
+from fastapi import FastAPI, HTTPException
+from fastapi.responses import JSONResponse
+import cv2
+import numpy as np
+from PIL import Image
+import noise
+import io
+import base64
+from pydantic import BaseModel
+
+app = FastAPI(title="Advanced Material Map Generator API")
+
+# Request model for input
+class MapRequest(BaseModel):
+    image_base64: str
+    normal_strength: float = 1.0
+    normal_blur: int = 5
+    normal_bilateral: bool = False
+    normal_color: float = 0.3
+    disp_contrast: float = 1.0
+    disp_noise: bool = False
+    disp_noise_scale: float = 0.1
+    disp_edge: float = 1.0
+    rough_invert: bool = True
+    rough_sharpness: float = 1.0
+    rough_detail: float = 0.5
+    rough_freq: float = 0.5
+
+def generate_normal_map(image: np.ndarray, strength: float, blur_size: int, use_bilateral: bool, color_influence: float) -> Image.Image:
+    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
+    if use_bilateral:
+        gray = cv2.bilateralFilter(gray, 9, 75, 75)
+    else:
+        gray = cv2.GaussianBlur(gray, (blur_size, blur_size), 0)
+    
+    levels = 3
+    normal_map = np.zeros((gray.shape[0], gray.shape[1], 3), dtype=np.float32)
+    for i in range(levels):
+        scale = 1 / (2 ** i)
+        resized = cv2.resize(gray, None, fx=scale, fy=scale, interpolation=cv2.INTER_AREA)
+        sobel_x = cv2.Scharr(resized, cv2.CV_64F, 1, 0)
+        sobel_y = cv2.Scharr(resized, cv2.CV_64F, 0, 1)
+        sobel_x = cv2.resize(sobel_x, (gray.shape[1], gray.shape[0]), interpolation=cv2.INTER_LINEAR)
+        sobel_y = cv2.resize(sobel_y, (gray.shape[1], gray.shape[0]), interpolation=cv2.INTER_LINEAR)
+        normal_map[..., 0] += sobel_x * (1.0 / levels)
+        normal_map[..., 1] += sobel_y * (1.0 / levels)
+    
+    normal_map[..., 0] = cv2.normalize(normal_map[..., 0], None, -strength, strength, cv2.NORM_MINMAX)
+    normal_map[..., 1] = cv2.normalize(normal_map[..., 1], None, -strength, strength, cv2.NORM_MINMAX)
+    normal_map[..., 2] = 1.0
+    
+    color_factor = color_influence * strength
+    normal_map[..., 0] += (image[..., 0] / 255.0 - 0.5) * color_factor
+    normal_map[..., 1] += (image[..., 1] / 255.0 - 0.5) * color_factor
+    
+    norm = np.linalg.norm(normal_map, axis=2, keepdims=True)
+    normal_map = np.divide(normal_map, norm, out=np.zeros_like(normal_map), where=norm != 0)
+    normal_map = (normal_map + 1) * 127.5
+    normal_map = np.clip(normal_map, 0, 255).astype(np.uint8)
+    return Image.fromarray(normal_map)
+
+def generate_displacement_map(image: np.ndarray, contrast: float, add_noise: bool, noise_scale: float, edge_boost: float) -> Image.Image:
+    img = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
+    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
+    img = clahe.apply(img)
+    img = cv2.convertScaleAbs(img, alpha=contrast, beta=0)
+    laplacian = cv2.Laplacian(img, cv2.CV_64F)
+    laplacian = cv2.convertScaleAbs(laplacian, alpha=edge_boost, beta=0)
+    img = cv2.addWeighted(img, 1.0, laplacian, 0.5 * edge_boost, 0)
+    if add_noise:
+        height, width = img.shape
+        noise_map = np.zeros((height, width), dtype=np.float32)
+        for y in range(height):
+            for x in range(width):
+                noise_map[y, x] = noise.pnoise2(x / 50.0, y / 50.0, octaves=6) * noise_scale * 255
+        img = cv2.add(img, noise_map.astype(np.uint8))
+    return Image.fromarray(img)
+
+def generate_roughness_map(image: np.ndarray, invert: bool, sharpness: float, detail_boost: float, frequency_weight: float) -> Image.Image:
+    img = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
+    low_freq = cv2.bilateralFilter(img, 9, 75, 75)
+    high_freq = cv2.subtract(img, low_freq)
+    img = cv2.addWeighted(low_freq, 1.0 - frequency_weight, high_freq, frequency_weight, 0)
+    if invert:
+        img = 255 - img
+    blurred = cv2.GaussianBlur(img, (5, 5), 0)
+    img = cv2.addWeighted(img, 1.0 + sharpness, blurred, -sharpness, 0)
+    img = cv2.addWeighted(img, 1.0 + detail_boost, blurred, -detail_boost, 0)
+    return Image.fromarray(img)
+
+def image_to_base64(img: Image.Image) -> str:
+    buffered = io.BytesIO()
+    img.save(buffered, format="PNG")
+    return base64.b64encode(buffered.getvalue()).decode("utf-8")
+
+@app.post("/generate_maps/")
+async def generate_maps(request: MapRequest):
+    try:
+        # Decode base64 image
+        image_bytes = base64.b64decode(request.image_base64)
+        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
+        img_array = np.array(image)
+
+        # Generate maps
+        normal_map = generate_normal_map(
+            img_array, request.normal_strength, request.normal_blur, 
+            request.normal_bilateral, request.normal_color
+        )
+        displacement_map = generate_displacement_map(
+            img_array, request.disp_contrast, request.disp_noise, 
+            request.disp_noise_scale, request.disp_edge
+        )
+        roughness_map = generate_roughness_map(
+            img_array, request.rough_invert, request.rough_sharpness, 
+            request.rough_detail, request.rough_freq
+        )
+
+        # Convert to base64
+        normal_base64 = image_to_base64(normal_map)
+        displacement_base64 = image_to_base64(displacement_map)
+        roughness_base64 = image_to_base64(roughness_map)
+
+        return JSONResponse(content={
+            "status": "success",
+            "normal_map": normal_base64,
+            "displacement_map": displacement_base64,
+            "roughness_map": roughness_base64
+        })
+
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))