#!/usr/bin/env python3
"""leaf_relief.py — printable leaf relief from one top-down photo.
Segments the leaf, extracts veins (black-hat on luminance), builds a height
field: base + gentle dome + raised veins, and closes it into a solid STL.
Honest label: this is a 2.5D photometric relief, not a photogrammetric scan.
"""
import sys, numpy as np, cv2, trimesh
from PIL import Image, ImageOps
SRC = sys.argv[1] if len(sys.argv) > 1 else "/mnt/user-data/uploads/IMG_2187.jpeg"
LEAF_WIDTH_MM = float(sys.argv[2]) if len(sys.argv) > 2 else 100.0    # measure the real leaf!
BASE, VEIN, DOME, CELL = 1.6, 0.7, 0.5, 0.35

im = ImageOps.exif_transpose(Image.open(SRC)).convert("RGB"); img = np.array(im)
hsv = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)
mask = ((hsv[:,:,1] > 70) & (hsv[:,:,2] > 40)).astype(np.uint8)
mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, np.ones((5,5),np.uint8))
mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, np.ones((9,9),np.uint8))
n, lab, st, _ = cv2.connectedComponentsWithStats(mask); mask = (lab == 1 + np.argmax(st[1:, 4])).astype(np.uint8)
blade = cv2.morphologyEx(mask, cv2.MORPH_OPEN, np.ones((25,25),np.uint8))
n, lab, st, _ = cv2.connectedComponentsWithStats(blade); blade = (lab == 1 + np.argmax(st[1:, 4])).astype(np.uint8)
blade = (cv2.dilate(blade, np.ones((7,7),np.uint8)) & mask).astype(bool)
ys, xs = np.where(blade); y0,y1,x0,x1 = ys.min(), ys.max(), xs.min(), xs.max()

gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY).astype(np.float32)
bh = cv2.morphologyEx(gray, cv2.MORPH_BLACKHAT, cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(15,15)))
veins = cv2.GaussianBlur(bh, (0,0), 1.2); veins[~blade] = 0
veins = np.clip(veins / np.percentile(veins[blade], 99.5), 0, 1)
veins = np.clip((veins - 0.30) / 0.70, 0, 1)            # drop surface mottling, keep real veins
veins = cv2.GaussianBlur(veins, (0,0), 1.0)

crop = (slice(y0-4, y1+5), slice(x0-4, x1+5))
px = LEAF_WIDTH_MM / (x1 - x0); f = px / CELL; cell = px / f
m = cv2.resize(blade[crop].astype(np.float32), None, fx=f, fy=f, interpolation=cv2.INTER_AREA) > 0.5
v = cv2.resize(veins[crop], None, fx=f, fy=f, interpolation=cv2.INTER_AREA)
dist = cv2.distanceTransform(m.astype(np.uint8), cv2.DIST_L2, 5); dist /= max(dist.max(), 1e-6)
hmap = np.where(m, BASE + DOME*np.sqrt(dist) + VEIN*v, 0.0)
Hh, Ww = m.shape
print(f"leaf {x1-x0}x{y1-y0} px -> grid {Ww}x{Hh} at {cell:.3f} mm, height {hmap[m].min():.2f}..{hmap[m].max():.2f} mm")

pad = np.pad(hmap, 1); pm = np.pad(m, 1).astype(np.float32)
num = pad[:-1,:-1]+pad[:-1,1:]+pad[1:,:-1]+pad[1:,1:]; den = pm[:-1,:-1]+pm[:-1,1:]+pm[1:,:-1]+pm[1:,1:]
ch = np.where(den > 0, num/np.maximum(den,1), 0)
gy, gx = np.mgrid[0:Hh+1, 0:Ww+1]
top = np.stack([gx*cell, (Hh-gy)*cell, ch], -1).reshape(-1,3); bot = top.copy(); bot[:,2] = 0
V = np.vstack([top, bot]); NB = len(top); idx = lambda i,j: i*(Ww+1)+j
ci, cj = np.where(m); a,b,c,d = idx(ci,cj), idx(ci,cj+1), idx(ci+1,cj+1), idx(ci+1,cj)
F = [np.stack([a,c,b],1), np.stack([a,d,c],1), np.stack([a+NB,b+NB,c+NB],1), np.stack([a+NB,c+NB,d+NB],1)]
mp = np.pad(m, 1)
for (di,dj,p,q) in [(-1,0,(0,0),(0,1)), (0,1,(0,1),(1,1)), (1,0,(1,1),(1,0)), (0,-1,(1,0),(0,0))]:
    si, sj = np.where(m & ~mp[1+di:1+di+Hh, 1+dj:1+dj+Ww])
    if len(si):
        p0, q0 = idx(si+p[0], sj+p[1]), idx(si+q[0], sj+q[1])
        F += [np.stack([p0, q0+NB, q0],1), np.stack([p0, p0+NB, q0+NB],1)]
mesh = trimesh.Trimesh(V, np.vstack(F), process=True); mesh.remove_unreferenced_vertices(); mesh.merge_vertices()
trimesh.repair.fix_normals(mesh)
print("watertight:", mesh.is_watertight, "| extents mm:", np.round(mesh.extents,1), "| volume cm3: %.1f" % (mesh.volume/1000))
mesh.export("leaf_relief.stl")

import matplotlib; matplotlib.use("Agg"); import matplotlib.pyplot as plt
from matplotlib.colors import LightSource
fig,ax=plt.subplots(1,3,figsize=(16,7))
ax[0].imshow(img[crop]); ax[0].contour(blade[crop].astype(float), levels=[0.5], colors="cyan", linewidths=0.8); ax[0].set_title("segmentation"); ax[0].axis("off")
ax[1].imshow(veins[crop], cmap="magma"); ax[1].set_title("vein map"); ax[1].axis("off")
sh = LightSource(315, 35).hillshade(np.where(m, hmap, 0), vert_exag=6, dx=cell, dy=cell)
ax[2].imshow(np.where(m, sh, 1.0), cmap="gray", vmin=0, vmax=1); ax[2].set_title("relief, raking light"); ax[2].axis("off")
plt.tight_layout(); plt.savefig("leaf_relief_preview.png", dpi=110); print("saved")