Spaces:
Sleeping
Sleeping
| import argparse | |
| import gradio as gr | |
| import json | |
| import os | |
| import tempfile | |
| from typing import Any | |
| import cv2 | |
| import numpy as np | |
| import torch | |
| from PIL import Image, ImageOps | |
| from sklearn.decomposition import PCA | |
| from stl import mesh | |
| from transformers import AutoModelForDepthEstimation, AutoProcessor | |
| try: | |
| import pillow_heif | |
| pillow_heif.register_heif_opener() | |
| HEIF_SUPPORT = True | |
| except Exception as heif_error: # pragma: no cover - surfaced in UI if HEIC open fails | |
| pillow_heif = None | |
| HEIF_SUPPORT = False | |
| HEIF_IMPORT_ERROR = heif_error | |
| else: | |
| HEIF_IMPORT_ERROR = None | |
| MIN_RESOLUTION_LIMIT = 10_000 | |
| DEFAULT_MAX_RESOLUTION = int(os.getenv("DEFAULT_MAX_RESOLUTION", "1500000")) | |
| MAX_RESOLUTION_LIMIT = int(os.getenv("MAX_RESOLUTION_LIMIT", "20000000")) | |
| def parse_launch_args() -> argparse.Namespace: | |
| parser = argparse.ArgumentParser(description="Image to 3D Relief Gradio app") | |
| parser.add_argument( | |
| "--max-resolution-pixels", | |
| type=int, | |
| default=None, | |
| help=( | |
| "Maximum number of pixels processed for depth estimation. " | |
| "Default comes from DEFAULT_MAX_RESOLUTION or 1,500,000." | |
| ), | |
| ) | |
| parser.add_argument( | |
| "--server-name", | |
| default=os.getenv("GRADIO_SERVER_NAME", "0.0.0.0"), | |
| help="Host/interface for Gradio to bind.", | |
| ) | |
| parser.add_argument( | |
| "--server-port", | |
| type=int, | |
| default=int(os.getenv("PORT", os.getenv("GRADIO_SERVER_PORT", "7860"))), | |
| help="Port for Gradio to bind.", | |
| ) | |
| args, _ = parser.parse_known_args() | |
| return args | |
| LAUNCH_ARGS = parse_launch_args() | |
| REQUESTED_MAX_RESOLUTION = LAUNCH_ARGS.max_resolution_pixels or DEFAULT_MAX_RESOLUTION | |
| CONFIGURED_MAX_RESOLUTION = int(max( | |
| MIN_RESOLUTION_LIMIT, | |
| min(MAX_RESOLUTION_LIMIT, REQUESTED_MAX_RESOLUTION), | |
| )) | |
| MODEL_ID = "depth-anything/Depth-Anything-V2-Large-hf" | |
| DEFAULT_TEXTURE_STRENGTH = 0.1 | |
| DEFAULT_TEXTURE_SMOOTHING = 5 | |
| DEFAULT_MIN_Z = 0.5 | |
| DEFAULT_MAX_Z = 5.0 | |
| DEFAULT_CURVE_POINTS = [ | |
| {"x": 0.0, "y": 0.0}, | |
| {"x": 1.0, "y": 1.0}, | |
| ] | |
| ACCEPTED_IMAGE_TYPES = [ | |
| ".png", | |
| ".jpg", | |
| ".jpeg", | |
| ".webp", | |
| ".bmp", | |
| ".tif", | |
| ".tiff", | |
| ".heic", | |
| ".heif", | |
| ] | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| processor = AutoProcessor.from_pretrained(MODEL_ID, use_fast=True) | |
| model = AutoModelForDepthEstimation.from_pretrained(MODEL_ID).to(device) | |
| model.eval() | |
| print(f"Model loaded successfully on {device}.") | |
| APP_CSS = """ | |
| #curve-points-json, #depth-histogram-json, #auto-preview-depth-btn { | |
| display: none !important; | |
| } | |
| .depth-editor-card { | |
| border: 1px solid var(--border-color-primary); | |
| border-radius: 12px; | |
| padding: 14px; | |
| background: var(--background-fill-secondary); | |
| } | |
| .depth-editor-toolbar { | |
| display: flex; | |
| flex-wrap: wrap; | |
| gap: 8px; | |
| margin: 8px 0 10px; | |
| } | |
| .depth-editor-toolbar button { | |
| border: 1px solid var(--border-color-primary); | |
| border-radius: 8px; | |
| padding: 6px 10px; | |
| background: var(--button-secondary-background-fill); | |
| color: var(--body-text-color); | |
| cursor: pointer; | |
| } | |
| .depth-editor-toolbar button:hover { | |
| background: var(--button-secondary-background-fill-hover); | |
| } | |
| #depth-curve-canvas { | |
| width: 100%; | |
| max-width: 720px; | |
| height: 360px; | |
| display: block; | |
| touch-action: none; | |
| border: 1px solid var(--border-color-primary); | |
| border-radius: 8px; | |
| background: #111827; | |
| } | |
| .depth-editor-help { | |
| margin-top: 8px; | |
| color: var(--body-text-color-subdued); | |
| font-size: 0.9rem; | |
| line-height: 1.35; | |
| } | |
| .depth-editor-readout { | |
| min-height: 1.2rem; | |
| margin-top: 6px; | |
| color: var(--body-text-color-subdued); | |
| font-family: var(--font-mono); | |
| font-size: 0.85rem; | |
| } | |
| """ | |
| CURVE_EDITOR_HEAD = """ | |
| <script> | |
| (() => { | |
| const DEFAULT_POINTS = [{x: 0, y: 0}, {x: 1, y: 1}]; | |
| const PRESETS = { | |
| linear: DEFAULT_POINTS, | |
| subject: [{x: 0, y: 0}, {x: 0.22, y: 0.13}, {x: 0.50, y: 0.50}, {x: 0.78, y: 0.88}, {x: 1, y: 1}], | |
| compressHigh: [{x: 0, y: 0}, {x: 0.45, y: 0.52}, {x: 0.75, y: 0.74}, {x: 1, y: 0.86}], | |
| compressLow: [{x: 0, y: 0.14}, {x: 0.25, y: 0.26}, {x: 0.55, y: 0.50}, {x: 1, y: 1}], | |
| foregroundPop: [{x: 0, y: 0}, {x: 0.35, y: 0.22}, {x: 0.62, y: 0.70}, {x: 1, y: 1}], | |
| invert: [{x: 0, y: 1}, {x: 1, y: 0}] | |
| }; | |
| const clamp = (v, lo = 0, hi = 1) => Math.max(lo, Math.min(hi, v)); | |
| function getTextarea(elemId) { | |
| const root = document.getElementById(elemId); | |
| if (!root) return null; | |
| return root.querySelector('textarea, input'); | |
| } | |
| function setTextareaValue(elemId, value) { | |
| const el = getTextarea(elemId); | |
| if (!el) return false; | |
| const proto = el.tagName === 'TEXTAREA' ? window.HTMLTextAreaElement.prototype : window.HTMLInputElement.prototype; | |
| const setter = Object.getOwnPropertyDescriptor(proto, 'value')?.set; | |
| if (setter) setter.call(el, value); | |
| else el.value = value; | |
| el.dispatchEvent(new Event('input', { bubbles: true })); | |
| el.dispatchEvent(new Event('change', { bubbles: true })); | |
| return true; | |
| } | |
| function readTextareaValue(elemId) { | |
| const el = getTextarea(elemId); | |
| return el ? el.value : ''; | |
| } | |
| function clickGradioButton(elemId) { | |
| const root = document.getElementById(elemId); | |
| const button = root?.querySelector('button'); | |
| if (button) button.click(); | |
| } | |
| function normalizePointMode(rawMode) { | |
| return String(rawMode || 'smooth').toLowerCase() === 'corner' ? 'corner' : 'smooth'; | |
| } | |
| function normalizePoints(points) { | |
| if (!Array.isArray(points)) return normalizePoints(DEFAULT_POINTS); | |
| const normalized = []; | |
| for (const point of points) { | |
| const x = typeof point?.x === 'number' ? point.x : Array.isArray(point) ? point[0] : NaN; | |
| const y = typeof point?.y === 'number' ? point.y : Array.isArray(point) ? point[1] : NaN; | |
| const mode = Array.isArray(point) ? 'smooth' : normalizePointMode(point?.mode); | |
| if (Number.isFinite(x) && Number.isFinite(y)) normalized.push({ x: clamp(x), y: clamp(y), mode }); | |
| } | |
| if (!normalized.length) normalized.push(...DEFAULT_POINTS.map(point => ({ ...point, mode: 'smooth' }))); | |
| normalized.sort((a, b) => a.x - b.x); | |
| normalized[0].x = 0; | |
| normalized[0].mode = 'smooth'; | |
| normalized[normalized.length - 1].x = 1; | |
| normalized[normalized.length - 1].mode = 'smooth'; | |
| return normalized; | |
| } | |
| function parseJson(value, fallback) { | |
| try { return JSON.parse(value); } catch (_) { return fallback; } | |
| } | |
| function initCurveEditor() { | |
| const canvas = document.getElementById('depth-curve-canvas'); | |
| if (!canvas || canvas.dataset.depthCurveReady === '1') return; | |
| canvas.dataset.depthCurveReady = '1'; | |
| const ctx = canvas.getContext('2d'); | |
| const readout = document.getElementById('depth-curve-readout'); | |
| const state = { | |
| points: normalizePoints(DEFAULT_POINTS), | |
| histogram: [], | |
| dragging: -1, | |
| selected: -1, | |
| pressedSelected: -1, | |
| pointerMoved: false, | |
| pointerStart: null, | |
| lastCurveJson: '', | |
| lastHistogramJson: '', | |
| hasDepth: false, | |
| refreshTimer: null | |
| }; | |
| function resizeCanvas() { | |
| const rect = canvas.getBoundingClientRect(); | |
| const dpr = window.devicePixelRatio || 1; | |
| const width = Math.max(320, Math.round(rect.width * dpr)); | |
| const height = Math.max(240, Math.round(rect.height * dpr)); | |
| if (canvas.width !== width || canvas.height !== height) { | |
| canvas.width = width; | |
| canvas.height = height; | |
| return true; | |
| } | |
| return false; | |
| } | |
| function graphRect() { | |
| const pad = 34 * (window.devicePixelRatio || 1); | |
| return { x: pad, y: pad * 0.65, w: canvas.width - pad * 1.45, h: canvas.height - pad * 1.5 }; | |
| } | |
| function toCanvas(point) { | |
| const r = graphRect(); | |
| return { x: r.x + point.x * r.w, y: r.y + (1 - point.y) * r.h }; | |
| } | |
| function fromCanvas(x, y) { | |
| const r = graphRect(); | |
| return { x: clamp((x - r.x) / r.w), y: clamp(1 - ((y - r.y) / r.h)) }; | |
| } | |
| function drawGrid(r) { | |
| ctx.save(); | |
| ctx.strokeStyle = 'rgba(255,255,255,0.13)'; | |
| ctx.lineWidth = 1; | |
| for (let i = 0; i <= 4; i++) { | |
| const x = r.x + (r.w * i / 4); | |
| const y = r.y + (r.h * i / 4); | |
| ctx.beginPath(); ctx.moveTo(x, r.y); ctx.lineTo(x, r.y + r.h); ctx.stroke(); | |
| ctx.beginPath(); ctx.moveTo(r.x, y); ctx.lineTo(r.x + r.w, y); ctx.stroke(); | |
| } | |
| ctx.strokeStyle = 'rgba(255,255,255,0.5)'; | |
| ctx.strokeRect(r.x, r.y, r.w, r.h); | |
| ctx.fillStyle = 'rgba(255,255,255,0.72)'; | |
| ctx.font = `${12 * (window.devicePixelRatio || 1)}px ui-monospace, SFMono-Regular, Menlo, monospace`; | |
| ctx.fillText('input height value →', r.x + r.w - 145 * (window.devicePixelRatio || 1), r.y + r.h + 24 * (window.devicePixelRatio || 1)); | |
| ctx.save(); | |
| ctx.translate(r.x - 24 * (window.devicePixelRatio || 1), r.y + 95 * (window.devicePixelRatio || 1)); | |
| ctx.rotate(-Math.PI / 2); | |
| ctx.fillText('output height', 0, 0); | |
| ctx.restore(); | |
| ctx.restore(); | |
| } | |
| function drawHistogram(r) { | |
| if (!state.histogram?.length) return; | |
| const maxBin = Math.max(...state.histogram, 1); | |
| const binW = r.w / state.histogram.length; | |
| ctx.save(); | |
| ctx.fillStyle = 'rgba(96, 165, 250, 0.35)'; | |
| state.histogram.forEach((count, i) => { | |
| const h = (count / maxBin) * r.h; | |
| ctx.fillRect(r.x + i * binW, r.y + r.h - h, Math.max(1, binW), h); | |
| }); | |
| ctx.restore(); | |
| } | |
| function getPchipEndpointSlope(h0, h1, d0, d1) { | |
| const slope = ((2 * h0 + h1) * d0 - h0 * d1) / Math.max(h0 + h1, 1e-6); | |
| if (slope === 0 || d0 === 0 || Math.sign(slope) !== Math.sign(d0)) return 0; | |
| if (Math.sign(d0) !== Math.sign(d1) && Math.abs(slope) > Math.abs(3 * d0)) return 3 * d0; | |
| return slope; | |
| } | |
| function getBezierCurveData(points) { | |
| if (points.length < 2) return null; | |
| const xs = points.map(point => point.x); | |
| const ys = points.map(point => point.y); | |
| const intervalCount = points.length - 1; | |
| const h = new Array(intervalCount); | |
| const secants = new Array(intervalCount); | |
| for (let i = 0; i < intervalCount; i++) { | |
| h[i] = Math.max(xs[i + 1] - xs[i], 1e-6); | |
| secants[i] = (ys[i + 1] - ys[i]) / h[i]; | |
| } | |
| const shared = new Array(points.length).fill(0); | |
| if (points.length === 2) { | |
| shared[0] = secants[0]; | |
| shared[1] = secants[0]; | |
| } else { | |
| shared[0] = getPchipEndpointSlope(h[0], h[1], secants[0], secants[1]); | |
| shared[points.length - 1] = getPchipEndpointSlope(h[h.length - 1], h[h.length - 2], secants[secants.length - 1], secants[secants.length - 2]); | |
| for (let i = 1; i < points.length - 1; i++) { | |
| const leftDelta = secants[i - 1]; | |
| const rightDelta = secants[i]; | |
| if (leftDelta === 0 || rightDelta === 0 || Math.sign(leftDelta) !== Math.sign(rightDelta)) { | |
| shared[i] = 0; | |
| } else { | |
| const leftH = h[i - 1]; | |
| const rightH = h[i]; | |
| const w1 = 2 * rightH + leftH; | |
| const w2 = rightH + 2 * leftH; | |
| shared[i] = (w1 + w2) / ((w1 / leftDelta) + (w2 / rightDelta)); | |
| } | |
| } | |
| } | |
| const leftSlopes = [...shared]; | |
| const rightSlopes = [...shared]; | |
| for (let i = 1; i < points.length - 1; i++) { | |
| if (points[i].mode === 'corner') { | |
| leftSlopes[i] = secants[i - 1]; | |
| rightSlopes[i] = secants[i]; | |
| } | |
| } | |
| return { xs, ys, h, leftSlopes, rightSlopes }; | |
| } | |
| function getSegmentControls(segmentIndex, curveData) { | |
| const x0 = curveData.xs[segmentIndex]; | |
| const y0 = curveData.ys[segmentIndex]; | |
| const x1 = curveData.xs[segmentIndex + 1]; | |
| const y1 = curveData.ys[segmentIndex + 1]; | |
| const width = curveData.h[segmentIndex]; | |
| return { | |
| cp1: { x: x0 + width / 3, y: y0 + curveData.rightSlopes[segmentIndex] * width / 3 }, | |
| cp2: { x: x1 - width / 3, y: y1 - curveData.leftSlopes[segmentIndex + 1] * width / 3 }, | |
| p3: { x: x1, y: y1 } | |
| }; | |
| } | |
| function drawPointMarker(point, index) { | |
| const c = toCanvas(point); | |
| const radius = 7 * (window.devicePixelRatio || 1); | |
| const isSelected = index === state.selected; | |
| const isEndpoint = index === 0 || index === state.points.length - 1; | |
| const isCorner = !isEndpoint && point.mode === 'corner'; | |
| ctx.save(); | |
| ctx.strokeStyle = isSelected ? '#fb923c' : '#38bdf8'; | |
| ctx.lineWidth = 2.5 * (window.devicePixelRatio || 1); | |
| if (isCorner) { | |
| ctx.beginPath(); | |
| ctx.moveTo(c.x, c.y - radius); | |
| ctx.lineTo(c.x + radius, c.y); | |
| ctx.lineTo(c.x, c.y + radius); | |
| ctx.lineTo(c.x - radius, c.y); | |
| ctx.closePath(); | |
| ctx.fillStyle = '#111827'; | |
| ctx.fill(); | |
| ctx.stroke(); | |
| } else { | |
| ctx.beginPath(); | |
| ctx.arc(c.x, c.y, radius, 0, Math.PI * 2); | |
| ctx.fillStyle = isSelected ? '#fb923c' : '#38bdf8'; | |
| ctx.fill(); | |
| ctx.strokeStyle = '#0f172a'; | |
| ctx.lineWidth = 2 * (window.devicePixelRatio || 1); | |
| ctx.stroke(); | |
| } | |
| ctx.restore(); | |
| } | |
| function drawCurve() { | |
| const curveData = getBezierCurveData(state.points); | |
| if (!curveData) return; | |
| ctx.save(); | |
| ctx.strokeStyle = '#f8fafc'; | |
| ctx.lineWidth = 3 * (window.devicePixelRatio || 1); | |
| ctx.beginPath(); | |
| const start = toCanvas(state.points[0]); | |
| ctx.moveTo(start.x, start.y); | |
| for (let i = 0; i < state.points.length - 1; i++) { | |
| const controls = getSegmentControls(i, curveData); | |
| const cp1 = toCanvas(controls.cp1); | |
| const cp2 = toCanvas(controls.cp2); | |
| const p3 = toCanvas(controls.p3); | |
| ctx.bezierCurveTo(cp1.x, cp1.y, cp2.x, cp2.y, p3.x, p3.y); | |
| } | |
| ctx.stroke(); | |
| state.points.forEach((point, index) => drawPointMarker(point, index)); | |
| ctx.restore(); | |
| } | |
| function draw() { | |
| resizeCanvas(); | |
| const r = graphRect(); | |
| ctx.clearRect(0, 0, canvas.width, canvas.height); | |
| ctx.fillStyle = '#111827'; | |
| ctx.fillRect(0, 0, canvas.width, canvas.height); | |
| drawHistogram(r); | |
| drawGrid(r); | |
| drawCurve(); | |
| if (readout) { | |
| const selected = state.selected >= 0 ? state.points[state.selected] : null; | |
| readout.textContent = selected | |
| ? `Selected ${state.selected === 0 || state.selected === state.points.length - 1 ? 'endpoint' : selected.mode} point: input ${selected.x.toFixed(3)} → output ${selected.y.toFixed(3)}` | |
| : `${state.points.length} control points · cubic Bézier`; | |
| } | |
| } | |
| function pushCurve(autoPreview = true) { | |
| state.points = normalizePoints(state.points); | |
| const payload = JSON.stringify(state.points.map(p => ({ x: +p.x.toFixed(5), y: +p.y.toFixed(5), mode: p.mode }))); | |
| state.lastCurveJson = payload; | |
| setTextareaValue('curve-points-json', payload); | |
| draw(); | |
| if (autoPreview && state.hasDepth) { | |
| window.clearTimeout(state.refreshTimer); | |
| state.refreshTimer = window.setTimeout(() => clickGradioButton('auto-preview-depth-btn'), 350); | |
| } | |
| } | |
| function nearestPoint(x, y) { | |
| const threshold = 15 * (window.devicePixelRatio || 1); | |
| let best = -1; | |
| let bestDistance = Infinity; | |
| state.points.forEach((point, index) => { | |
| const c = toCanvas(point); | |
| const distance = Math.hypot(c.x - x, c.y - y); | |
| if (distance < threshold && distance < bestDistance) { | |
| best = index; | |
| bestDistance = distance; | |
| } | |
| }); | |
| return best; | |
| } | |
| function pointerPosition(event) { | |
| const rect = canvas.getBoundingClientRect(); | |
| const dpr = window.devicePixelRatio || 1; | |
| return { x: (event.clientX - rect.left) * dpr, y: (event.clientY - rect.top) * dpr }; | |
| } | |
| function toggleSelectedPointMode() { | |
| if (state.selected <= 0 || state.selected >= state.points.length - 1) return; | |
| state.points[state.selected].mode = state.points[state.selected].mode === 'corner' ? 'smooth' : 'corner'; | |
| pushCurve(true); | |
| } | |
| canvas.addEventListener('pointerdown', (event) => { | |
| event.preventDefault(); | |
| canvas.setPointerCapture(event.pointerId); | |
| const pos = pointerPosition(event); | |
| state.pointerStart = pos; | |
| state.pointerMoved = false; | |
| state.pressedSelected = -1; | |
| const nearest = nearestPoint(pos.x, pos.y); | |
| if (nearest >= 0) { | |
| if (state.selected === nearest) state.pressedSelected = nearest; | |
| state.dragging = nearest; | |
| state.selected = nearest; | |
| } else { | |
| const point = { ...fromCanvas(pos.x, pos.y), mode: 'smooth' }; | |
| state.points.push(point); | |
| state.points.sort((a, b) => a.x - b.x); | |
| state.selected = state.points.findIndex(p => p === point); | |
| state.dragging = state.selected; | |
| pushCurve(false); | |
| } | |
| draw(); | |
| }); | |
| canvas.addEventListener('pointermove', (event) => { | |
| if (state.dragging < 0) return; | |
| event.preventDefault(); | |
| const pos = pointerPosition(event); | |
| if (state.pointerStart && Math.hypot(pos.x - state.pointerStart.x, pos.y - state.pointerStart.y) > 3 * (window.devicePixelRatio || 1)) { | |
| state.pointerMoved = true; | |
| } | |
| const point = fromCanvas(pos.x, pos.y); | |
| const i = state.dragging; | |
| const left = i === 0 ? 0 : state.points[i - 1].x + 0.002; | |
| const right = i === state.points.length - 1 ? 1 : state.points[i + 1].x - 0.002; | |
| state.points[i].x = i === 0 ? 0 : i === state.points.length - 1 ? 1 : clamp(point.x, left, right); | |
| state.points[i].y = point.y; | |
| pushCurve(true); | |
| }); | |
| canvas.addEventListener('pointerup', (event) => { | |
| const shouldToggle = state.dragging === state.pressedSelected && !state.pointerMoved; | |
| if (shouldToggle) toggleSelectedPointMode(); | |
| state.dragging = -1; | |
| state.pressedSelected = -1; | |
| state.pointerMoved = false; | |
| state.pointerStart = null; | |
| try { canvas.releasePointerCapture(event.pointerId); } catch (_) {} | |
| draw(); | |
| }); | |
| canvas.addEventListener('dblclick', (event) => { | |
| event.preventDefault(); | |
| const pos = pointerPosition(event); | |
| const nearest = nearestPoint(pos.x, pos.y); | |
| if (nearest > 0 && nearest < state.points.length - 1) { | |
| state.points.splice(nearest, 1); | |
| state.selected = -1; | |
| pushCurve(true); | |
| } | |
| }); | |
| window.addEventListener('keydown', (event) => { | |
| if ((event.key === 'Delete' || event.key === 'Backspace') && state.selected > 0 && state.selected < state.points.length - 1) { | |
| state.points.splice(state.selected, 1); | |
| state.selected = -1; | |
| pushCurve(true); | |
| } | |
| }); | |
| document.addEventListener('click', (event) => { | |
| const button = event.target.closest('[data-depth-preset]'); | |
| if (!button) return; | |
| const preset = button.dataset.depthPreset; | |
| state.points = normalizePoints(PRESETS[preset] || DEFAULT_POINTS); | |
| state.selected = -1; | |
| pushCurve(true); | |
| }); | |
| window.addEventListener('resize', draw); | |
| window.setInterval(() => { | |
| const curveJson = readTextareaValue('curve-points-json'); | |
| if (curveJson && curveJson !== state.lastCurveJson) { | |
| state.points = normalizePoints(parseJson(curveJson, DEFAULT_POINTS)); | |
| state.lastCurveJson = curveJson; | |
| draw(); | |
| } | |
| const histogramJson = readTextareaValue('depth-histogram-json'); | |
| if (histogramJson && histogramJson !== state.lastHistogramJson) { | |
| const parsed = parseJson(histogramJson, []); | |
| state.histogram = Array.isArray(parsed) ? parsed : []; | |
| state.hasDepth = state.histogram.some(count => Number(count) > 0); | |
| state.lastHistogramJson = histogramJson; | |
| draw(); | |
| } | |
| }, 500); | |
| pushCurve(false); | |
| draw(); | |
| } | |
| function startCurveEditorWatcher() { | |
| initCurveEditor(); | |
| window.setInterval(initCurveEditor, 500); | |
| } | |
| if (document.readyState === 'loading') { | |
| document.addEventListener('DOMContentLoaded', startCurveEditorWatcher); | |
| } else { | |
| startCurveEditorWatcher(); | |
| } | |
| window.addEventListener('load', initCurveEditor); | |
| })(); | |
| </script> | |
| """ | |
| CURVE_EDITOR_HTML = """ | |
| <div class="depth-editor-card"> | |
| <strong>Depth histogram + curves editor</strong> | |
| <div class="depth-editor-toolbar"> | |
| <button type="button" data-depth-preset="linear">Linear</button> | |
| <button type="button" data-depth-preset="subject">Midtone contrast</button> | |
| <button type="button" data-depth-preset="compressHigh">Compress high values</button> | |
| <button type="button" data-depth-preset="compressLow">Compress low values</button> | |
| <button type="button" data-depth-preset="foregroundPop">Foreground pop</button> | |
| <button type="button" data-depth-preset="invert">Invert curve</button> | |
| </div> | |
| <canvas id="depth-curve-canvas" width="720" height="360"></canvas> | |
| <div id="depth-curve-readout" class="depth-editor-readout"></div> | |
| <div class="depth-editor-help"> | |
| Drag control points to remap the normalized height interpretation before STL export. The curve uses piecewise cubic Bézier segments. | |
| Click the curve to add a point. Click an already selected point to toggle smooth/corner continuity. Corner points draw as diamonds. Double-click a non-endpoint, or select it and press Delete, to remove it. | |
| The blue histogram shows the current smoothed base-depth distribution. | |
| </div> | |
| </div> | |
| """ | |
| def clamp(value: float, minimum: float, maximum: float) -> float: | |
| return max(minimum, min(maximum, value)) | |
| def normalize_array(values: np.ndarray) -> np.ndarray: | |
| values = values.astype(np.float32, copy=False) | |
| v_min = float(np.nanmin(values)) | |
| v_max = float(np.nanmax(values)) | |
| if not np.isfinite(v_min) or not np.isfinite(v_max) or v_max <= v_min: | |
| return np.zeros_like(values, dtype=np.float32) | |
| return ((values - v_min) / (v_max - v_min)).astype(np.float32) | |
| def sanitize_max_pixels(max_pixels: Any) -> int: | |
| try: | |
| requested = int(float(max_pixels)) | |
| except (TypeError, ValueError): | |
| requested = DEFAULT_MAX_RESOLUTION | |
| return int(clamp(requested, MIN_RESOLUTION_LIMIT, MAX_RESOLUTION_LIMIT)) | |
| def coerce_float(value: Any, default: float) -> float: | |
| try: | |
| if value is None: | |
| return float(default) | |
| return float(value) | |
| except (TypeError, ValueError): | |
| return float(default) | |
| def coerce_z_bounds(min_z: Any, max_z: Any) -> tuple[float, float]: | |
| min_z_value = coerce_float(min_z, DEFAULT_MIN_Z) | |
| max_z_value = coerce_float(max_z, DEFAULT_MAX_Z) | |
| return min_z_value, max_z_value | |
| def get_lanczos_resample_filter() -> int: | |
| return getattr(getattr(Image, "Resampling", Image), "LANCZOS") | |
| def load_and_resize_image(input_filepath: str, max_pixels: int) -> tuple[np.ndarray, str]: | |
| if input_filepath is None: | |
| raise gr.Error("Please upload an image.") | |
| filepath = str(input_filepath) | |
| filepath_lower = filepath.lower() | |
| try: | |
| image_pil = Image.open(filepath) | |
| image_pil = ImageOps.exif_transpose(image_pil).convert("RGB") | |
| except Exception as exc: | |
| heif_hint = "" | |
| if filepath_lower.endswith((".heic", ".heif")): | |
| if HEIF_SUPPORT: | |
| heif_hint = " The file looks like HEIC/HEIF; pillow-heif is installed but could not decode this file." | |
| else: | |
| heif_hint = f" HEIC/HEIF support is unavailable because pillow-heif failed to load: {HEIF_IMPORT_ERROR}" | |
| raise gr.Error(f"Could not open image file.{heif_hint} Error: {exc}") from exc | |
| original_w, original_h = image_pil.size | |
| original_pixels = original_w * original_h | |
| max_pixels = sanitize_max_pixels(max_pixels) | |
| resized = False | |
| if original_pixels > max_pixels: | |
| ratio = (max_pixels / original_pixels) ** 0.5 | |
| new_w = max(1, int(original_w * ratio)) | |
| new_h = max(1, int(original_h * ratio)) | |
| image_pil = image_pil.resize((new_w, new_h), get_lanczos_resample_filter()) | |
| resized = True | |
| input_image = np.array(image_pil, dtype=np.uint8) | |
| current_h, current_w = input_image.shape[:2] | |
| status = ( | |
| f"Loaded {original_w}×{original_h} image. " | |
| f"Processing at {current_w}×{current_h} ({current_w * current_h:,} pixels)." | |
| ) | |
| if resized: | |
| status += f" Downsampled to stay under the {max_pixels:,}-pixel limit." | |
| if filepath_lower.endswith((".heic", ".heif")): | |
| status += " HEIC/HEIF decoded server-side." | |
| return input_image, status | |
| def estimate_depth_map(input_image: np.ndarray) -> np.ndarray: | |
| image = Image.fromarray(input_image).convert("RGB") | |
| with torch.no_grad(): | |
| inputs = processor(images=image, return_tensors="pt").to(device) | |
| outputs = model(**inputs) | |
| predicted_depth = outputs.predicted_depth | |
| depth = torch.nn.functional.interpolate( | |
| predicted_depth.unsqueeze(1), | |
| size=(image.height, image.width), | |
| mode="bilinear", | |
| align_corners=False, | |
| ).squeeze().cpu().numpy() | |
| return normalize_depth_for_emboss(depth) | |
| def apply_smoothing(values: np.ndarray, smoothing: int) -> np.ndarray: | |
| try: | |
| ksize = int(smoothing) | |
| except (TypeError, ValueError): | |
| ksize = 0 | |
| if ksize <= 1: | |
| return values.astype(np.float32, copy=True) | |
| if ksize % 2 == 0: | |
| ksize += 1 | |
| return cv2.GaussianBlur(values.astype(np.float32), (ksize, ksize), 0).astype(np.float32) | |
| def sanitize_curve_points(curve_points_json: str | None) -> list[dict[str, Any]]: | |
| try: | |
| raw_points = json.loads(curve_points_json or "") | |
| except (TypeError, json.JSONDecodeError): | |
| raw_points = DEFAULT_CURVE_POINTS | |
| points: list[dict[str, Any]] = [] | |
| if isinstance(raw_points, list): | |
| for point in raw_points: | |
| mode = "smooth" | |
| if isinstance(point, dict): | |
| x, y = point.get("x"), point.get("y") | |
| raw_mode = str(point.get("mode", "smooth")).lower() | |
| if raw_mode in {"corner", "broken", "linear"}: | |
| mode = "corner" | |
| elif isinstance(point, (list, tuple)) and len(point) >= 2: | |
| x, y = point[0], point[1] | |
| else: | |
| continue | |
| try: | |
| x_f = clamp(float(x), 0.0, 1.0) | |
| y_f = clamp(float(y), 0.0, 1.0) | |
| except (TypeError, ValueError): | |
| continue | |
| points.append({"x": x_f, "y": y_f, "mode": mode}) | |
| if len(points) < 2: | |
| points = [{**point, "mode": "smooth"} for point in DEFAULT_CURVE_POINTS] | |
| points.sort(key=lambda point: float(point["x"])) | |
| deduped: list[dict[str, Any]] = [] | |
| for point in points: | |
| if deduped and abs(float(point["x"]) - float(deduped[-1]["x"])) < 1e-5: | |
| deduped[-1] = point | |
| else: | |
| deduped.append(point) | |
| if len(deduped) < 2: | |
| deduped = [{**point, "mode": "smooth"} for point in DEFAULT_CURVE_POINTS] | |
| deduped[0]["x"] = 0.0 | |
| deduped[-1]["x"] = 1.0 | |
| deduped[0]["mode"] = "smooth" | |
| deduped[-1]["mode"] = "smooth" | |
| return deduped | |
| def curve_points_to_json(points: list[dict[str, Any]]) -> str: | |
| return json.dumps([ | |
| { | |
| "x": round(float(point["x"]), 5), | |
| "y": round(float(point["y"]), 5), | |
| "mode": "corner" if str(point.get("mode", "smooth")) == "corner" else "smooth", | |
| } | |
| for point in points | |
| ]) | |
| def pchip_endpoint_slope(h0: float, h1: float, delta0: float, delta1: float) -> float: | |
| slope = ((2.0 * h0 + h1) * delta0 - h0 * delta1) / max(h0 + h1, 1e-6) | |
| if slope == 0.0 or delta0 == 0.0 or np.sign(slope) != np.sign(delta0): | |
| return 0.0 | |
| if np.sign(delta0) != np.sign(delta1) and abs(slope) > abs(3.0 * delta0): | |
| return float(3.0 * delta0) | |
| return float(slope) | |
| def cubic_bezier_curve_data(points: list[dict[str, Any]]) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: | |
| x_values = np.array([point["x"] for point in points], dtype=np.float32) | |
| y_values = np.array([point["y"] for point in points], dtype=np.float32) | |
| interval_count = len(points) - 1 | |
| if interval_count <= 0: | |
| return x_values, y_values, np.zeros_like(x_values), np.zeros_like(x_values) | |
| h_values = np.maximum(np.diff(x_values), 1e-6) | |
| secant_slopes = np.diff(y_values) / h_values | |
| shared_slopes = np.zeros(len(points), dtype=np.float32) | |
| if len(points) == 2: | |
| shared_slopes[:] = secant_slopes[0] | |
| else: | |
| shared_slopes[0] = pchip_endpoint_slope(float(h_values[0]), float(h_values[1]), float(secant_slopes[0]), float(secant_slopes[1])) | |
| shared_slopes[-1] = pchip_endpoint_slope(float(h_values[-1]), float(h_values[-2]), float(secant_slopes[-1]), float(secant_slopes[-2])) | |
| for i in range(1, len(points) - 1): | |
| left_delta = float(secant_slopes[i - 1]) | |
| right_delta = float(secant_slopes[i]) | |
| if left_delta == 0.0 or right_delta == 0.0 or np.sign(left_delta) != np.sign(right_delta): | |
| shared_slopes[i] = 0.0 | |
| else: | |
| left_h = float(h_values[i - 1]) | |
| right_h = float(h_values[i]) | |
| w1 = 2.0 * right_h + left_h | |
| w2 = right_h + 2.0 * left_h | |
| shared_slopes[i] = float((w1 + w2) / ((w1 / left_delta) + (w2 / right_delta))) | |
| left_slopes = shared_slopes.copy() | |
| right_slopes = shared_slopes.copy() | |
| for i in range(1, len(points) - 1): | |
| if str(points[i].get("mode", "smooth")) == "corner": | |
| left_slopes[i] = secant_slopes[i - 1] | |
| right_slopes[i] = secant_slopes[i] | |
| return x_values, y_values, left_slopes.astype(np.float32), right_slopes.astype(np.float32) | |
| def evaluate_piecewise_cubic_bezier(values: np.ndarray, points: list[dict[str, Any]]) -> np.ndarray: | |
| x_values, y_values, left_slopes, right_slopes = cubic_bezier_curve_data(points) | |
| if len(points) <= 1: | |
| return np.full_like(values, y_values[0] if len(y_values) else 0.0, dtype=np.float32) | |
| flat_values = values.astype(np.float32, copy=False).reshape(-1) | |
| interval_indices = np.searchsorted(x_values, flat_values, side="right") - 1 | |
| interval_indices = np.clip(interval_indices, 0, len(x_values) - 2) | |
| interval_widths = np.maximum(x_values[interval_indices + 1] - x_values[interval_indices], 1e-6) | |
| t_values = np.clip((flat_values - x_values[interval_indices]) / interval_widths, 0.0, 1.0) | |
| omt_values = 1.0 - t_values | |
| p0 = y_values[interval_indices] | |
| p1 = y_values[interval_indices] + right_slopes[interval_indices] * (interval_widths / 3.0) | |
| p2 = y_values[interval_indices + 1] - left_slopes[interval_indices + 1] * (interval_widths / 3.0) | |
| p3 = y_values[interval_indices + 1] | |
| adjusted = ( | |
| (omt_values ** 3) * p0 | |
| + 3.0 * (omt_values ** 2) * t_values * p1 | |
| + 3.0 * omt_values * (t_values ** 2) * p2 | |
| + (t_values ** 3) * p3 | |
| ) | |
| return adjusted.reshape(values.shape).astype(np.float32) | |
| def apply_depth_curve(depth_map: np.ndarray, curve_points_json: str | None, invert_depth: bool) -> np.ndarray: | |
| working = depth_map.astype(np.float32, copy=True) | |
| if invert_depth: | |
| working = 1.0 - working | |
| points = sanitize_curve_points(curve_points_json) | |
| adjusted = evaluate_piecewise_cubic_bezier(working, points) | |
| return np.clip(adjusted, 0.0, 1.0) | |
| def depth_to_preview_image(depth_map: np.ndarray) -> Image.Image: | |
| preview = np.clip(depth_map, 0.0, 1.0) | |
| preview_uint8 = np.rint(np.clip(preview, 0.0, 1.0) * 255.0).astype(np.uint8) | |
| return Image.fromarray(preview_uint8) | |
| def histogram_json(depth_map: np.ndarray, bins: int = 96) -> str: | |
| counts, _ = np.histogram(np.clip(depth_map, 0.0, 1.0), bins=bins, range=(0.0, 1.0)) | |
| return json.dumps(counts.astype(int).tolist()) | |
| def prepare_depth(input_filepath: str): | |
| input_image, load_status = load_and_resize_image(input_filepath, CONFIGURED_MAX_RESOLUTION) | |
| depth_normalized = estimate_depth_map(input_image) | |
| curve_json = curve_points_to_json(DEFAULT_CURVE_POINTS) | |
| adjusted_depth = apply_depth_curve(depth_normalized, curve_json, invert_depth=False) | |
| height_map = combine_depth_and_texture( | |
| input_image, | |
| adjusted_depth, | |
| DEFAULT_TEXTURE_STRENGTH, | |
| DEFAULT_TEXTURE_SMOOTHING, | |
| ) | |
| state = { | |
| "image_rgb": input_image, | |
| "raw_depth": depth_normalized, | |
| "adjusted_depth": adjusted_depth, | |
| "height_map": height_map, | |
| "curve_points_json": curve_json, | |
| "depth_map_smoothing": 0, | |
| "invert_depth": False, | |
| "texture_strength": DEFAULT_TEXTURE_STRENGTH, | |
| "texture_smoothing": DEFAULT_TEXTURE_SMOOTHING, | |
| "load_status": load_status, | |
| } | |
| status = ( | |
| f"{load_status}\nDepth estimation complete. Use the curve and luminance texture controls to tune the 2D height map before generating the STL." | |
| ) | |
| return ( | |
| state, | |
| Image.fromarray(input_image), | |
| depth_to_preview_image(depth_normalized), | |
| depth_to_preview_image(height_map), | |
| histogram_json(depth_normalized), | |
| curve_json, | |
| status, | |
| "Preview is using the linear curve with the default luminance texture settings.", | |
| ) | |
| def clear_depth_outputs(): | |
| return ( | |
| None, | |
| None, | |
| None, | |
| None, | |
| "[]", | |
| curve_points_to_json(DEFAULT_CURVE_POINTS), | |
| "Upload an image to estimate its depth map.", | |
| "", | |
| ) | |
| def clear_preview_outputs(message: str = "Upload an image to estimate its depth map before tuning preview controls."): | |
| return ( | |
| None, | |
| None, | |
| "[]", | |
| curve_points_to_json(DEFAULT_CURVE_POINTS), | |
| message, | |
| ) | |
| def auto_prepare_depth(input_filepath: str): | |
| if input_filepath is None: | |
| return clear_depth_outputs() | |
| return prepare_depth(input_filepath) | |
| def get_smoothed_depth_from_state(depth_state: dict[str, Any] | None, depth_map_smoothing: int) -> np.ndarray: | |
| if not depth_state or "raw_depth" not in depth_state: | |
| raise gr.Error("Please upload an image and estimate its depth first.") | |
| return normalize_array(apply_smoothing(depth_state["raw_depth"], depth_map_smoothing)) | |
| def update_depth_preview( | |
| depth_state: dict[str, Any] | None, | |
| curve_points_json: str | None, | |
| depth_map_smoothing: int, | |
| invert_depth: bool, | |
| texture_strength: float, | |
| texture_smoothing: int, | |
| ): | |
| if not depth_state or "raw_depth" not in depth_state: | |
| return clear_preview_outputs() | |
| texture_strength = coerce_float(texture_strength, DEFAULT_TEXTURE_STRENGTH) | |
| texture_smoothing = int(coerce_float(texture_smoothing, DEFAULT_TEXTURE_SMOOTHING)) | |
| smoothed_depth = get_smoothed_depth_from_state(depth_state, depth_map_smoothing) | |
| adjusted_depth = apply_depth_curve(smoothed_depth, curve_points_json, invert_depth) | |
| points = sanitize_curve_points(curve_points_json) | |
| normalized_curve_json = curve_points_to_json(points) | |
| height_map = combine_depth_and_texture( | |
| depth_state["image_rgb"], | |
| adjusted_depth, | |
| texture_strength, | |
| texture_smoothing, | |
| ) | |
| if depth_state is not None: | |
| depth_state["adjusted_depth"] = adjusted_depth | |
| depth_state["height_map"] = height_map | |
| depth_state["curve_points_json"] = normalized_curve_json | |
| depth_state["depth_map_smoothing"] = int(depth_map_smoothing) | |
| depth_state["invert_depth"] = bool(invert_depth) | |
| depth_state["texture_strength"] = float(texture_strength) | |
| depth_state["texture_smoothing"] = int(texture_smoothing) | |
| status = ( | |
| f"Preview updated with {len(points)} curve points, texture strength {float(texture_strength):.3f}, " | |
| f"texture smoothing {int(texture_smoothing)}." | |
| ) | |
| return ( | |
| depth_state, | |
| depth_to_preview_image(height_map), | |
| histogram_json(smoothed_depth), | |
| normalized_curve_json, | |
| status, | |
| ) | |
| def get_adjusted_depth_for_settings( | |
| depth_state: dict[str, Any], | |
| curve_points_json: str | None, | |
| depth_map_smoothing: int, | |
| invert_depth: bool, | |
| ) -> np.ndarray: | |
| normalized_curve_json = curve_points_to_json(sanitize_curve_points(curve_points_json)) | |
| cached_depth = depth_state.get("adjusted_depth") | |
| cache_matches = ( | |
| cached_depth is not None | |
| and depth_state.get("curve_points_json") == normalized_curve_json | |
| and int(depth_state.get("depth_map_smoothing", -1)) == int(depth_map_smoothing) | |
| and bool(depth_state.get("invert_depth", False)) == bool(invert_depth) | |
| ) | |
| if cache_matches: | |
| return cached_depth | |
| smoothed_depth = get_smoothed_depth_from_state(depth_state, depth_map_smoothing) | |
| adjusted_depth = apply_depth_curve(smoothed_depth, normalized_curve_json, invert_depth) | |
| depth_state["adjusted_depth"] = adjusted_depth | |
| depth_state["curve_points_json"] = normalized_curve_json | |
| depth_state["depth_map_smoothing"] = int(depth_map_smoothing) | |
| depth_state["invert_depth"] = bool(invert_depth) | |
| return adjusted_depth | |
| def get_height_map_for_settings( | |
| depth_state: dict[str, Any], | |
| curve_points_json: str | None, | |
| depth_map_smoothing: int, | |
| invert_depth: bool, | |
| texture_strength: float, | |
| texture_smoothing: int, | |
| ) -> np.ndarray: | |
| texture_strength = coerce_float(texture_strength, DEFAULT_TEXTURE_STRENGTH) | |
| texture_smoothing = int(coerce_float(texture_smoothing, DEFAULT_TEXTURE_SMOOTHING)) | |
| normalized_curve_json = curve_points_to_json(sanitize_curve_points(curve_points_json)) | |
| cached_height_map = depth_state.get("height_map") | |
| cache_matches = ( | |
| cached_height_map is not None | |
| and depth_state.get("curve_points_json") == normalized_curve_json | |
| and int(depth_state.get("depth_map_smoothing", -1)) == int(depth_map_smoothing) | |
| and bool(depth_state.get("invert_depth", False)) == bool(invert_depth) | |
| and float(depth_state.get("texture_strength", -1.0)) == float(texture_strength) | |
| and int(depth_state.get("texture_smoothing", -1)) == int(texture_smoothing) | |
| ) | |
| if cache_matches: | |
| return cached_height_map | |
| adjusted_depth = get_adjusted_depth_for_settings( | |
| depth_state, | |
| normalized_curve_json, | |
| depth_map_smoothing, | |
| invert_depth, | |
| ) | |
| height_map = combine_depth_and_texture( | |
| depth_state["image_rgb"], | |
| adjusted_depth, | |
| texture_strength, | |
| texture_smoothing, | |
| ) | |
| depth_state["height_map"] = height_map | |
| depth_state["texture_strength"] = float(texture_strength) | |
| depth_state["texture_smoothing"] = int(texture_smoothing) | |
| return height_map | |
| def combine_depth_and_texture( | |
| image_rgb: np.ndarray, | |
| adjusted_depth: np.ndarray, | |
| texture_strength: float, | |
| texture_smoothing: int, | |
| ) -> np.ndarray: | |
| texture_strength = clamp(coerce_float(texture_strength, DEFAULT_TEXTURE_STRENGTH), 0.0, 0.5) | |
| if texture_strength <= 0.0: | |
| return np.clip(adjusted_depth, 0.0, 1.0).astype(np.float32) | |
| gray_image = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2GRAY).astype(np.float32) / 255.0 | |
| low_frequency = apply_smoothing(gray_image, 5) | |
| detail = gray_image - low_frequency | |
| texture_smoothing = int(coerce_float(texture_smoothing, DEFAULT_TEXTURE_SMOOTHING)) | |
| if texture_smoothing > 1: | |
| detail = apply_smoothing(detail, texture_smoothing) | |
| if detail.shape != adjusted_depth.shape: | |
| detail = cv2.resize( | |
| detail, | |
| (adjusted_depth.shape[1], adjusted_depth.shape[0]), | |
| interpolation=cv2.INTER_LINEAR, | |
| ) | |
| scale = float(np.percentile(np.abs(detail), 95.0)) | |
| if not np.isfinite(scale) or scale <= 1e-6: | |
| detail = np.zeros_like(adjusted_depth, dtype=np.float32) | |
| else: | |
| detail = np.clip(detail / scale, -1.0, 1.0).astype(np.float32) | |
| combined_map = adjusted_depth + (detail * texture_strength) | |
| return np.clip(combined_map, 0.0, 1.0).astype(np.float32) | |
| def apply_pca_correction_to_z(z_data: np.ndarray, x_length: float, min_z: float, max_z: float) -> np.ndarray: | |
| height, width = z_data.shape | |
| y_length = x_length * (height / width) | |
| x_coords_1d = np.linspace(0, x_length, width) | |
| y_coords_1d = np.linspace(y_length, 0, height) | |
| x_grid, y_grid = np.meshgrid(x_coords_1d, y_coords_1d) | |
| points = np.stack([x_grid.flatten(), y_grid.flatten(), z_data.flatten()], axis=1) | |
| n_points = points.shape[0] | |
| n_samples = min(n_points, 50_000) | |
| sample_indices = np.random.choice(n_points, n_samples, replace=False) | |
| pca = PCA(n_components=3) | |
| pca.fit(points[sample_indices]) | |
| normal = pca.components_[2] | |
| if normal[2] < 0: | |
| normal *= -1 | |
| p0 = pca.mean_ | |
| z_plane = p0[2] - (normal[0] * (x_grid - p0[0]) + normal[1] * (y_grid - p0[1])) / normal[2] | |
| corrected_z = z_data - z_plane | |
| corrected_normalized = normalize_array(corrected_z) | |
| return min_z + corrected_normalized * (max_z - min_z) | |
| def build_stl_faces(z_data: np.ndarray, x_length: float, close_body: bool) -> tuple[np.ndarray, float]: | |
| height, width = z_data.shape | |
| y_length = float(x_length * (height / width)) | |
| z_data = z_data.astype(np.float32, copy=False) | |
| x_coords = np.linspace(0, x_length, width, dtype=np.float32) | |
| y_coords = np.linspace(y_length, 0, height, dtype=np.float32) | |
| x_grid, y_grid = np.meshgrid(x_coords, y_coords) | |
| vertices = np.stack([x_grid, y_grid, z_data], axis=-1) | |
| tl = vertices[:-1, :-1] | |
| bl = vertices[1:, :-1] | |
| br = vertices[1:, 1:] | |
| tr = vertices[:-1, 1:] | |
| quad_count = (height - 1) * (width - 1) | |
| top_faces = np.empty((quad_count * 2, 3, 3), dtype=np.float32) | |
| top_faces[0::2] = np.stack([tl, bl, br], axis=2).reshape(-1, 3, 3) | |
| top_faces[1::2] = np.stack([tl, br, tr], axis=2).reshape(-1, 3, 3) | |
| if not close_body: | |
| return top_faces, y_length | |
| v_tl = vertices[0, 0] | |
| v_tr = vertices[0, width - 1] | |
| v_bl = vertices[height - 1, 0] | |
| v_br = vertices[height - 1, width - 1] | |
| b_tl = np.array([v_tl[0], v_tl[1], 0], dtype=np.float32) | |
| b_tr = np.array([v_tr[0], v_tr[1], 0], dtype=np.float32) | |
| b_bl = np.array([v_bl[0], v_bl[1], 0], dtype=np.float32) | |
| b_br = np.array([v_br[0], v_br[1], 0], dtype=np.float32) | |
| side_faces = np.array([ | |
| [v_tl, b_tl, b_tr], [v_tl, b_tr, v_tr], | |
| [v_br, b_br, b_bl], [v_br, b_bl, v_bl], | |
| [v_bl, b_bl, b_tl], [v_bl, b_tl, v_tl], | |
| [v_tr, b_tr, b_br], [v_tr, b_br, v_br], | |
| ], dtype=np.float32) | |
| base_faces = np.array([ | |
| [b_tl, b_br, b_bl], | |
| [b_tl, b_tr, b_br], | |
| ], dtype=np.float32) | |
| return np.concatenate([top_faces, side_faces, base_faces], axis=0), y_length | |
| def flip_faces_for_preview(faces: np.ndarray, y_length: float) -> np.ndarray: | |
| preview_faces = faces.copy() | |
| preview_faces[..., 1] = np.float32(y_length) - preview_faces[..., 1] | |
| preview_faces = preview_faces[:, [0, 2, 1], :] | |
| return preview_faces | |
| def save_stl_faces(faces: np.ndarray) -> str: | |
| surface = mesh.Mesh(np.zeros(faces.shape[0], dtype=mesh.Mesh.dtype)) | |
| surface.vectors = faces | |
| with tempfile.NamedTemporaryFile(delete=False, suffix=".stl") as tmpfile: | |
| surface.save(tmpfile.name) | |
| return tmpfile.name | |
| def build_stl_mesh(z_data: np.ndarray, x_length: float, close_body: bool, flip_for_preview: bool = False) -> str: | |
| faces, y_length = build_stl_faces(z_data, x_length, close_body) | |
| if flip_for_preview: | |
| faces = flip_faces_for_preview(faces, y_length) | |
| return save_stl_faces(faces) | |
| def generate_3d_model_from_adjusted_depth( | |
| depth_state: dict[str, Any] | None, | |
| curve_points_json: str | None, | |
| depth_map_smoothing: int, | |
| invert_depth: bool, | |
| texture_strength: float, | |
| texture_smoothing: int, | |
| max_z: float, | |
| min_z: float, | |
| x_length: float, | |
| do_pca_correction: bool, | |
| close_body: bool, | |
| ): | |
| min_z, max_z = coerce_z_bounds(min_z, max_z) | |
| x_length = coerce_float(x_length, 100.0) | |
| if max_z <= min_z: | |
| raise gr.Error("Max Z-height must be greater than Min Z-height.") | |
| if x_length <= 0: | |
| raise gr.Error("X Length must be positive.") | |
| if not depth_state or "image_rgb" not in depth_state: | |
| raise gr.Error("Please estimate a depth map before generating an STL.") | |
| height_map = get_height_map_for_settings( | |
| depth_state, | |
| curve_points_json, | |
| depth_map_smoothing, | |
| invert_depth, | |
| texture_strength, | |
| texture_smoothing, | |
| ) | |
| z_data = min_z + height_map * (max_z - min_z) | |
| if do_pca_correction: | |
| z_data = apply_pca_correction_to_z(z_data, x_length, min_z, max_z) | |
| export_faces, y_length = build_stl_faces(z_data, x_length, close_body) | |
| export_stl_path = save_stl_faces(export_faces) | |
| preview_stl_path = save_stl_faces(flip_faces_for_preview(export_faces, y_length)) | |
| point_count = len(sanitize_curve_points(curve_points_json)) | |
| status = ( | |
| f"Generated STL from previewed height map using {point_count} curve points and texture strength " | |
| f"{float(texture_strength):.3f}." | |
| ) | |
| return preview_stl_path, export_stl_path, status | |
| # CONTROLLED_EMBOSS_V1 | |
| DEFAULT_HEIGHT_GAMMA = 1.0 | |
| DEFAULT_BACKGROUND_FLATTEN = 0.08 | |
| def normalize_depth_for_emboss(values: np.ndarray) -> np.ndarray: | |
| values = values.astype(np.float32, copy=False) | |
| finite_values = values[np.isfinite(values)] | |
| if finite_values.size == 0: | |
| return np.zeros_like(values, dtype=np.float32) | |
| low, high = np.nanpercentile(finite_values, [2.0, 98.0]) | |
| if not np.isfinite(low) or not np.isfinite(high) or high <= low: | |
| return normalize_array(values) | |
| return np.clip((values - low) / (high - low), 0.0, 1.0).astype(np.float32) | |
| def apply_emboss_controls(height_map: np.ndarray, height_gamma: float, background_flatten: float) -> np.ndarray: | |
| gamma = clamp(coerce_float(height_gamma, DEFAULT_HEIGHT_GAMMA), 0.25, 4.0) | |
| flatten = clamp(coerce_float(background_flatten, DEFAULT_BACKGROUND_FLATTEN), 0.0, 0.8) | |
| controlled = np.power(np.clip(height_map, 0.0, 1.0), gamma).astype(np.float32) | |
| if flatten > 0.0: | |
| controlled = np.clip((controlled - flatten) / max(1.0 - flatten, 1e-6), 0.0, 1.0) | |
| return controlled.astype(np.float32) | |
| def height_to_hillshade_image(height_map: np.ndarray) -> Image.Image: | |
| heights = np.clip(height_map, 0.0, 1.0).astype(np.float32) | |
| gradient_y, gradient_x = np.gradient(heights) | |
| normals = np.dstack((-gradient_x * 4.0, -gradient_y * 4.0, np.ones_like(heights))) | |
| normals /= np.maximum(np.linalg.norm(normals, axis=2, keepdims=True), 1e-6) | |
| light = np.array([-0.45, -0.35, 0.82], dtype=np.float32) | |
| light /= np.linalg.norm(light) | |
| shade = 0.18 + 0.82 * np.tensordot(normals, light, axes=([2], [0])) | |
| return Image.fromarray(np.rint(np.clip(shade, 0.0, 1.0) * 255.0).astype(np.uint8)) | |
| def save_height_map_png(height_map: np.ndarray) -> str: | |
| height_uint16 = np.rint(np.clip(height_map, 0.0, 1.0) * 65535.0).astype(np.uint16) | |
| with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmpfile: | |
| Image.fromarray(height_uint16, mode="I;16").save(tmpfile.name, format="PNG") | |
| return tmpfile.name | |
| def build_controlled_height_map( | |
| image_rgb: np.ndarray, | |
| adjusted_depth: np.ndarray, | |
| texture_strength: float, | |
| texture_smoothing: int, | |
| height_gamma: float, | |
| background_flatten: float, | |
| ) -> np.ndarray: | |
| controlled_depth = apply_emboss_controls(adjusted_depth, height_gamma, background_flatten) | |
| return combine_depth_and_texture(image_rgb, controlled_depth, texture_strength, texture_smoothing) | |
| def estimate_depth_map(input_image: np.ndarray) -> np.ndarray: | |
| image = Image.fromarray(input_image).convert("RGB") | |
| with torch.no_grad(): | |
| inputs = processor(images=image, return_tensors="pt").to(device) | |
| outputs = model(**inputs) | |
| predicted_depth = outputs.predicted_depth | |
| depth = torch.nn.functional.interpolate( | |
| predicted_depth.unsqueeze(1), | |
| size=(image.height, image.width), | |
| mode="bilinear", | |
| align_corners=False, | |
| ).squeeze().cpu().numpy() | |
| return normalize_depth_for_emboss(depth) | |
| def prepare_depth(input_filepath: str): | |
| input_image, load_status = load_and_resize_image(input_filepath, CONFIGURED_MAX_RESOLUTION) | |
| depth_normalized = estimate_depth_map(input_image) | |
| curve_json = curve_points_to_json(DEFAULT_CURVE_POINTS) | |
| adjusted_depth = apply_depth_curve(depth_normalized, curve_json, invert_depth=False) | |
| height_map = build_controlled_height_map( | |
| input_image, | |
| adjusted_depth, | |
| DEFAULT_TEXTURE_STRENGTH, | |
| DEFAULT_TEXTURE_SMOOTHING, | |
| DEFAULT_HEIGHT_GAMMA, | |
| DEFAULT_BACKGROUND_FLATTEN, | |
| ) | |
| state = { | |
| "image_rgb": input_image, | |
| "raw_depth": depth_normalized, | |
| "adjusted_depth": adjusted_depth, | |
| "height_map": height_map, | |
| "curve_points_json": curve_json, | |
| "depth_map_smoothing": 0, | |
| "invert_depth": False, | |
| "texture_strength": DEFAULT_TEXTURE_STRENGTH, | |
| "texture_smoothing": DEFAULT_TEXTURE_SMOOTHING, | |
| "height_gamma": DEFAULT_HEIGHT_GAMMA, | |
| "background_flatten": DEFAULT_BACKGROUND_FLATTEN, | |
| "load_status": load_status, | |
| } | |
| status = ( | |
| f"{load_status}\nDepth estimation complete. Tune gamma and background flatten on the true grayscale height map before adding fine detail." | |
| ) | |
| return ( | |
| state, | |
| Image.fromarray(input_image), | |
| depth_to_preview_image(depth_normalized), | |
| depth_to_preview_image(height_map), | |
| height_to_hillshade_image(height_map), | |
| save_height_map_png(height_map), | |
| histogram_json(depth_normalized), | |
| curve_json, | |
| status, | |
| "Preview uses robust depth normalization, linear gamma, and the default background flatten threshold.", | |
| ) | |
| def clear_depth_outputs(): | |
| return ( | |
| None, | |
| None, | |
| None, | |
| None, | |
| None, | |
| None, | |
| "[]", | |
| curve_points_to_json(DEFAULT_CURVE_POINTS), | |
| "Upload an image to estimate its depth map.", | |
| "", | |
| ) | |
| def clear_preview_outputs(message: str = "Upload an image to estimate its depth map before tuning preview controls."): | |
| return ( | |
| None, | |
| None, | |
| None, | |
| None, | |
| "[]", | |
| curve_points_to_json(DEFAULT_CURVE_POINTS), | |
| message, | |
| ) | |
| def get_smoothed_depth_from_state(depth_state: dict[str, Any] | None, depth_map_smoothing: int) -> np.ndarray: | |
| if not depth_state or "raw_depth" not in depth_state: | |
| raise gr.Error("Please upload an image and estimate its depth first.") | |
| return normalize_depth_for_emboss(apply_smoothing(depth_state["raw_depth"], depth_map_smoothing)) | |
| def update_depth_preview( | |
| depth_state: dict[str, Any] | None, | |
| curve_points_json: str | None, | |
| depth_map_smoothing: int, | |
| invert_depth: bool, | |
| texture_strength: float, | |
| texture_smoothing: int, | |
| height_gamma: float, | |
| background_flatten: float, | |
| ): | |
| if not depth_state or "raw_depth" not in depth_state: | |
| return clear_preview_outputs() | |
| texture_strength = clamp(coerce_float(texture_strength, DEFAULT_TEXTURE_STRENGTH), 0.0, 0.5) | |
| texture_smoothing = int(coerce_float(texture_smoothing, DEFAULT_TEXTURE_SMOOTHING)) | |
| height_gamma = clamp(coerce_float(height_gamma, DEFAULT_HEIGHT_GAMMA), 0.25, 4.0) | |
| background_flatten = clamp(coerce_float(background_flatten, DEFAULT_BACKGROUND_FLATTEN), 0.0, 0.8) | |
| smoothed_depth = get_smoothed_depth_from_state(depth_state, depth_map_smoothing) | |
| adjusted_depth = apply_depth_curve(smoothed_depth, curve_points_json, invert_depth) | |
| points = sanitize_curve_points(curve_points_json) | |
| normalized_curve_json = curve_points_to_json(points) | |
| height_map = build_controlled_height_map( | |
| depth_state["image_rgb"], | |
| adjusted_depth, | |
| texture_strength, | |
| texture_smoothing, | |
| height_gamma, | |
| background_flatten, | |
| ) | |
| depth_state["adjusted_depth"] = adjusted_depth | |
| depth_state["height_map"] = height_map | |
| depth_state["curve_points_json"] = normalized_curve_json | |
| depth_state["depth_map_smoothing"] = int(depth_map_smoothing) | |
| depth_state["invert_depth"] = bool(invert_depth) | |
| depth_state["texture_strength"] = float(texture_strength) | |
| depth_state["texture_smoothing"] = int(texture_smoothing) | |
| depth_state["height_gamma"] = float(height_gamma) | |
| depth_state["background_flatten"] = float(background_flatten) | |
| status = ( | |
| f"Controlled preview updated with gamma {height_gamma:.2f}, background flatten {background_flatten:.2f}, " | |
| f"and fine detail strength {texture_strength:.3f}." | |
| ) | |
| return ( | |
| depth_state, | |
| depth_to_preview_image(height_map), | |
| height_to_hillshade_image(height_map), | |
| save_height_map_png(height_map), | |
| histogram_json(smoothed_depth), | |
| normalized_curve_json, | |
| status, | |
| ) | |
| def get_adjusted_depth_for_settings( | |
| depth_state: dict[str, Any], | |
| curve_points_json: str | None, | |
| depth_map_smoothing: int, | |
| invert_depth: bool, | |
| ) -> np.ndarray: | |
| normalized_curve_json = curve_points_to_json(sanitize_curve_points(curve_points_json)) | |
| cached_depth = depth_state.get("adjusted_depth") | |
| cache_matches = ( | |
| cached_depth is not None | |
| and depth_state.get("curve_points_json") == normalized_curve_json | |
| and int(depth_state.get("depth_map_smoothing", -1)) == int(depth_map_smoothing) | |
| and bool(depth_state.get("invert_depth", False)) == bool(invert_depth) | |
| ) | |
| if cache_matches: | |
| return cached_depth | |
| smoothed_depth = get_smoothed_depth_from_state(depth_state, depth_map_smoothing) | |
| adjusted_depth = apply_depth_curve(smoothed_depth, normalized_curve_json, invert_depth) | |
| depth_state["adjusted_depth"] = adjusted_depth | |
| depth_state["curve_points_json"] = normalized_curve_json | |
| depth_state["depth_map_smoothing"] = int(depth_map_smoothing) | |
| depth_state["invert_depth"] = bool(invert_depth) | |
| return adjusted_depth | |
| def get_height_map_for_settings( | |
| depth_state: dict[str, Any], | |
| curve_points_json: str | None, | |
| depth_map_smoothing: int, | |
| invert_depth: bool, | |
| texture_strength: float, | |
| texture_smoothing: int, | |
| height_gamma: float, | |
| background_flatten: float, | |
| ) -> np.ndarray: | |
| texture_strength = clamp(coerce_float(texture_strength, DEFAULT_TEXTURE_STRENGTH), 0.0, 0.5) | |
| texture_smoothing = int(coerce_float(texture_smoothing, DEFAULT_TEXTURE_SMOOTHING)) | |
| height_gamma = clamp(coerce_float(height_gamma, DEFAULT_HEIGHT_GAMMA), 0.25, 4.0) | |
| background_flatten = clamp(coerce_float(background_flatten, DEFAULT_BACKGROUND_FLATTEN), 0.0, 0.8) | |
| normalized_curve_json = curve_points_to_json(sanitize_curve_points(curve_points_json)) | |
| cached_height_map = depth_state.get("height_map") | |
| cache_matches = ( | |
| cached_height_map is not None | |
| and depth_state.get("curve_points_json") == normalized_curve_json | |
| and int(depth_state.get("depth_map_smoothing", -1)) == int(depth_map_smoothing) | |
| and bool(depth_state.get("invert_depth", False)) == bool(invert_depth) | |
| and float(depth_state.get("texture_strength", -1.0)) == float(texture_strength) | |
| and int(depth_state.get("texture_smoothing", -1)) == int(texture_smoothing) | |
| and float(depth_state.get("height_gamma", -1.0)) == float(height_gamma) | |
| and float(depth_state.get("background_flatten", -1.0)) == float(background_flatten) | |
| ) | |
| if cache_matches: | |
| return cached_height_map | |
| adjusted_depth = get_adjusted_depth_for_settings( | |
| depth_state, | |
| normalized_curve_json, | |
| depth_map_smoothing, | |
| invert_depth, | |
| ) | |
| height_map = build_controlled_height_map( | |
| depth_state["image_rgb"], | |
| adjusted_depth, | |
| texture_strength, | |
| texture_smoothing, | |
| height_gamma, | |
| background_flatten, | |
| ) | |
| depth_state["height_map"] = height_map | |
| depth_state["texture_strength"] = float(texture_strength) | |
| depth_state["texture_smoothing"] = int(texture_smoothing) | |
| depth_state["height_gamma"] = float(height_gamma) | |
| depth_state["background_flatten"] = float(background_flatten) | |
| return height_map | |
| def generate_3d_model_from_adjusted_depth( | |
| depth_state: dict[str, Any] | None, | |
| curve_points_json: str | None, | |
| depth_map_smoothing: int, | |
| invert_depth: bool, | |
| texture_strength: float, | |
| texture_smoothing: int, | |
| height_gamma: float, | |
| background_flatten: float, | |
| max_z: float, | |
| min_z: float, | |
| x_length: float, | |
| do_pca_correction: bool, | |
| close_body: bool, | |
| ): | |
| min_z, max_z = coerce_z_bounds(min_z, max_z) | |
| x_length = coerce_float(x_length, 100.0) | |
| if max_z <= min_z: | |
| raise gr.Error("Max Z-height must be greater than Min Z-height.") | |
| if x_length <= 0: | |
| raise gr.Error("X Length must be positive.") | |
| if not depth_state or "image_rgb" not in depth_state: | |
| raise gr.Error("Please estimate a depth map before generating an STL.") | |
| height_map = get_height_map_for_settings( | |
| depth_state, | |
| curve_points_json, | |
| depth_map_smoothing, | |
| invert_depth, | |
| texture_strength, | |
| texture_smoothing, | |
| height_gamma, | |
| background_flatten, | |
| ) | |
| z_data = min_z + height_map * (max_z - min_z) | |
| if do_pca_correction: | |
| z_data = apply_pca_correction_to_z(z_data, x_length, min_z, max_z) | |
| export_faces, y_length = build_stl_faces(z_data, x_length, close_body) | |
| export_stl_path = save_stl_faces(export_faces) | |
| preview_stl_path = save_stl_faces(flip_faces_for_preview(export_faces, y_length)) | |
| point_count = len(sanitize_curve_points(curve_points_json)) | |
| status = ( | |
| f"Generated STL from controlled height map using {point_count} curve points, gamma {float(height_gamma):.2f}, " | |
| f"and background flatten {float(background_flatten):.2f}." | |
| ) | |
| return preview_stl_path, export_stl_path, status | |
| # PUBLIC_TEST_EMBOSS_V1 | |
| TEST_LINEAR_CURVE = [ | |
| {"x": 0.0, "y": 0.0}, | |
| {"x": 1.0, "y": 1.0}, | |
| ] | |
| TEST_COMPRESS_HIGH_CURVE = [ | |
| {"x": 0.0, "y": 0.0}, | |
| {"x": 0.45, "y": 0.52}, | |
| {"x": 0.75, "y": 0.74}, | |
| {"x": 1.0, "y": 0.86}, | |
| ] | |
| def prepare_public_test_outputs( | |
| input_filepath: str | None, | |
| version: str, | |
| texture_strength: float, | |
| ) -> tuple[Image.Image | None, Image.Image | None, str]: | |
| if not input_filepath: | |
| return None, None, "Upload an image first." | |
| input_image, _ = load_and_resize_image(input_filepath, CONFIGURED_MAX_RESOLUTION) | |
| depth_map = estimate_depth_map(input_image) | |
| if version == "Version 2 — Compress high values": | |
| curve_points = TEST_COMPRESS_HIGH_CURVE | |
| else: | |
| curve_points = TEST_LINEAR_CURVE | |
| adjusted_depth = apply_depth_curve( | |
| depth_map, | |
| curve_points_to_json(curve_points), | |
| invert_depth=False, | |
| ) | |
| height_map = build_controlled_height_map( | |
| input_image, | |
| adjusted_depth, | |
| texture_strength, | |
| DEFAULT_TEXTURE_SMOOTHING, | |
| DEFAULT_HEIGHT_GAMMA, | |
| DEFAULT_BACKGROUND_FLATTEN, | |
| ) | |
| return ( | |
| depth_to_preview_image(height_map), | |
| height_to_hillshade_image(height_map), | |
| f"Completed {version} with Fine Detail Strength {float(texture_strength):.3f}.", | |
| ) | |
| with gr.Blocks(title="Test Embossed Effect") as demo: | |
| gr.Markdown("# Test Embossed Effect") | |
| gr.Markdown("Upload one image and compare the true grayscale height map with its embossed surface preview.") | |
| input_image = gr.Image( | |
| type="filepath", | |
| label="Upload one image", | |
| ) | |
| version = gr.Radio( | |
| choices=["Version 1 — Linear", "Version 2 — Compress high values"], | |
| value="Version 1 — Linear", | |
| label="Height interpretation version", | |
| ) | |
| texture_strength = gr.Slider( | |
| minimum=0.0, | |
| maximum=0.5, | |
| value=DEFAULT_TEXTURE_STRENGTH, | |
| step=0.001, | |
| label="Fine Detail Strength", | |
| info="Adds zero-mean fine source detail. Default uses strength 0.1 and smoothing 5.", | |
| ) | |
| rerun_button = gr.Button("Run / Rerun Preview", variant="primary") | |
| run_status = gr.Textbox( | |
| value="Upload an image, choose a version, then click Run / Rerun Preview.", | |
| label="Run status", | |
| interactive=False, | |
| ) | |
| with gr.Row(): | |
| height_map_output = gr.Image( | |
| format="png", | |
| image_mode="L", | |
| label="True grayscale height map", | |
| interactive=False, | |
| ) | |
| hillshade_output = gr.Image( | |
| format="png", | |
| image_mode="L", | |
| label="Embossed surface preview", | |
| interactive=False, | |
| ) | |
| test_inputs = [input_image, version, texture_strength] | |
| test_outputs = [height_map_output, hillshade_output, run_status] | |
| input_image.change( | |
| fn=prepare_public_test_outputs, | |
| inputs=test_inputs, | |
| outputs=test_outputs, | |
| show_progress="full", | |
| ) | |
| rerun_button.click( | |
| fn=prepare_public_test_outputs, | |
| inputs=test_inputs, | |
| outputs=test_outputs, | |
| show_progress="full", | |
| ) | |
| if __name__ == "__main__": | |
| demo.queue().launch( | |
| server_name=LAUNCH_ARGS.server_name, | |
| server_port=LAUNCH_ARGS.server_port, | |
| theme="base", | |
| ) | |