472 lines
19 KiB
Python
472 lines
19 KiB
Python
"""Render DWG/DXF/PDF → PNG image(s) for vision model consumption.
|
||
|
||
Strategy: multi-floor architectural drawings are split per detected legend.
|
||
Each floor renders to its own PNG at a resolution where individual symbols
|
||
remain distinguishable for the vision model.
|
||
"""
|
||
import logging
|
||
import subprocess
|
||
from pathlib import Path
|
||
|
||
import io
|
||
|
||
import ezdxf
|
||
from ezdxf.addons.drawing.config import (
|
||
BackgroundPolicy, ColorPolicy, Configuration,
|
||
)
|
||
from ezdxf.addons.drawing import Frontend, RenderContext, layout
|
||
from ezdxf.addons.drawing.svg import SVGBackend
|
||
from ezdxf.addons.drawing.matplotlib import MatplotlibBackend
|
||
import cairosvg
|
||
import matplotlib
|
||
matplotlib.use("Agg") # headless, no GUI thread
|
||
import matplotlib.pyplot as plt
|
||
from PIL import Image
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
RENDER_PX = 9000 # longer-edge target; matches counting MAX_DRAWING_PX so the
|
||
# matcher never has to downscale the (lossless) vector render
|
||
|
||
|
||
def dwg_to_dxf(dwg_path: Path, out_dir: Path) -> Path:
|
||
dxf_path = out_dir / f"{dwg_path.stem}.dxf"
|
||
r = subprocess.run(
|
||
["dwgread", "-O", "DXF", "-o", str(dxf_path), str(dwg_path)],
|
||
capture_output=True, text=True, timeout=180,
|
||
)
|
||
if not dxf_path.exists():
|
||
raise RuntimeError(f"DWG→DXF failed (exit {r.returncode}): {r.stderr or r.stdout}")
|
||
return dxf_path
|
||
|
||
|
||
SKIP_LAYER_PATTERNS = (
|
||
"dimens", "kota", "kóta", "koty", # dimensions
|
||
"sit_konst", "konstrukce", "ckkoty", # structural / dimensioning
|
||
"ckprofi", "profily", "csprofily", # steel/concrete profile rezy
|
||
"sanita", "vzt", "ov_kan", "tzb", # plumbing / HVAC (when not the target)
|
||
"viewport", "defpoints", "0_", "_b_", # CAD bookkeeping
|
||
"raster", "wipeout",
|
||
)
|
||
|
||
|
||
def _clean_doc(doc, drop_layers=False):
|
||
"""Strip elements that overwhelm vision rendering."""
|
||
msp = doc.modelspace()
|
||
|
||
# Force every layer visible. Electrical/symbol layers (power sockets,
|
||
# lights, etc.) are frequently turned OFF or FROZEN in the source DWG,
|
||
# so ezdxf skips them as "invisible" and they never reach the render —
|
||
# even though the geometry exists and the matcher can find it. Turning
|
||
# all layers on lets the user actually see and box those symbols.
|
||
for layer in doc.layers:
|
||
try:
|
||
layer.on()
|
||
layer.thaw()
|
||
except Exception:
|
||
pass
|
||
|
||
# Purge INSERTs that reference a missing block — in modelspace AND inside
|
||
# block definitions. dwgread-converted DXFs sometimes drop anonymous
|
||
# blocks ("*U", "*D") while a valid block still contains a nested INSERT
|
||
# of them; the renderer then crashes ("Required block definition for *U
|
||
# does not exist") when it explodes that reference.
|
||
valid_blocks = {b.name for b in doc.blocks}
|
||
for layout in (msp, *list(doc.blocks)):
|
||
for ins in list(layout.query("INSERT")):
|
||
if ins.dxf.name not in valid_blocks:
|
||
try:
|
||
layout.delete_entity(ins)
|
||
except Exception:
|
||
pass
|
||
# Clear per-entity invisibility flags (same reason as layers above).
|
||
for e in msp:
|
||
try:
|
||
if e.dxf.hasattr("invisible") and e.dxf.invisible:
|
||
e.dxf.invisible = 0
|
||
except Exception:
|
||
pass
|
||
# Remove annotation/fill noise by ENTITY TYPE (robust across drawings —
|
||
# unlike layer-name matching, which is per-drawing fragile). Dimensions,
|
||
# text and leaders are never the symbol being counted but dominate ink
|
||
# density on real architectural DWGs (~38% → a few %). Geometry that
|
||
# forms symbols (LINE/LWPOLYLINE/CIRCLE/ARC/INSERT) is kept.
|
||
for typ in ("HATCH", "SOLID", "MPOLYGON", "DIMENSION", "MTEXT", "TEXT",
|
||
"LEADER", "MULTILEADER", "ARC_DIMENSION"):
|
||
for e in list(msp.query(typ)):
|
||
msp.delete_entity(e)
|
||
if drop_layers:
|
||
for e in list(msp):
|
||
layer = str(getattr(e.dxf, "layer", "")).lower()
|
||
if any(p in layer for p in SKIP_LAYER_PATTERNS):
|
||
try:
|
||
msp.delete_entity(e)
|
||
except Exception:
|
||
pass
|
||
for e in msp:
|
||
try:
|
||
# 5 (0.05 mm) — known to render visibly at our raster scale.
|
||
# Hairline (0) went sub-pixel and produced a blank image.
|
||
e.dxf.lineweight = 5
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
def find_floors(dxf_path: Path) -> list[dict]:
|
||
"""Find legend 'LEGENDA' markers — each represents one floor.
|
||
|
||
Returns list of {legend_xy, floor_bbox} dicts ordered top to bottom.
|
||
"""
|
||
doc = ezdxf.readfile(str(dxf_path))
|
||
msp = doc.modelspace()
|
||
|
||
positions = []
|
||
for e in msp:
|
||
text = ""
|
||
if e.dxftype() == "MTEXT":
|
||
text = e.text
|
||
elif e.dxftype() == "TEXT":
|
||
text = e.dxf.text
|
||
if text and text.strip().upper() == "LEGENDA":
|
||
try:
|
||
positions.append((e.dxf.insert.x, e.dxf.insert.y))
|
||
except Exception:
|
||
pass
|
||
positions.sort(key=lambda p: -p[1])
|
||
|
||
if not positions:
|
||
return [{"legend_xy": None, "floor_bbox": None}]
|
||
|
||
# Use the same cleanup as render_region so extents reflect what'll be drawn
|
||
_clean_doc(doc)
|
||
from ezdxf.bbox import extents
|
||
try:
|
||
ext = extents(msp, fast=True)
|
||
mxmin, mymin = ext.extmin.x, ext.extmin.y
|
||
mxmax, mymax = ext.extmax.x, ext.extmax.y
|
||
except Exception:
|
||
mxmin, mymin = -1e9, -1e9
|
||
mxmax, mymax = 1e9, 1e9
|
||
logger.info("find_floors: model extents x=(%.0f,%.0f) y=(%.0f,%.0f)",
|
||
mxmin, mxmax, mymin, mymax)
|
||
|
||
floors = []
|
||
for i, (lx, ly) in enumerate(positions):
|
||
if i + 1 < len(positions):
|
||
y_height = ly - positions[i + 1][1]
|
||
elif i > 0:
|
||
y_height = positions[i - 1][1] - ly
|
||
else:
|
||
y_height = 60000
|
||
# Legend appears at the TOP of its floor view; plan extends downward
|
||
y_top = min(mymax, ly + 0.10 * y_height)
|
||
y_bot = max(mymin, ly - 0.95 * y_height)
|
||
# X span = whole model width (floor plans typically span the page width)
|
||
x_left = mxmin
|
||
x_right = mxmax
|
||
floors.append({
|
||
"legend_xy": (lx, ly),
|
||
"floor_bbox": (x_left, y_bot, x_right, y_top),
|
||
})
|
||
return floors
|
||
|
||
|
||
def render_region(dxf_path: Path, out_path: Path, bbox: tuple | None,
|
||
render_px: int | None = None,
|
||
backend: str = "cairosvg") -> Path:
|
||
"""Render a DXF (optionally clipped to bbox) to PNG.
|
||
|
||
Two backends:
|
||
- "cairosvg" (default): ezdxf → SVG → cairosvg → PNG. Faithful
|
||
vector linework. ~30–90 s on dense site plans at 4500 px.
|
||
- "matplotlib": ezdxf → Agg → PNG. Comparable speed, sometimes
|
||
faster on extremely entity-dense drawings.
|
||
Cropping is done in pixel space after rasterization.
|
||
"""
|
||
doc = ezdxf.readfile(str(dxf_path))
|
||
auditor = doc.audit()
|
||
if auditor.has_errors:
|
||
logger.info("DXF audit: %d errors", len(auditor.errors))
|
||
# Entity-type cleanup only (drop_layers=False). Layer-name matching is
|
||
# per-drawing fragile — it blanked this sample entirely. The unconditional
|
||
# DIMENSION/TEXT/HATCH removal in _clean_doc is the robust density fix.
|
||
_clean_doc(doc)
|
||
msp = doc.modelspace()
|
||
|
||
from ezdxf.bbox import extents
|
||
try:
|
||
ext = extents(msp, fast=True)
|
||
model_xmin = ext.extmin.x
|
||
model_ymin = ext.extmin.y
|
||
model_w = ext.size.x or 1
|
||
model_h = ext.size.y or 1
|
||
except Exception:
|
||
model_xmin = model_ymin = 0
|
||
model_w = model_h = 1
|
||
|
||
config = Configuration(
|
||
background_policy=BackgroundPolicy.WHITE,
|
||
color_policy=ColorPolicy.BLACK,
|
||
lineweight_scaling=0.5,
|
||
min_lineweight=0.05,
|
||
)
|
||
|
||
# If the caller asked for a sub-region (zoom into one area), copy
|
||
# the in-region entities into a fresh doc. ezdxf auto-fits the layout
|
||
# extents to the doc's geometry, so the zoom render's page covers
|
||
# ONLY the user's region at high effective DPI. (Without this, the
|
||
# SVG keeps the FULL drawing's coord range and our filtered geometry
|
||
# ends up at sub-pixel scale.)
|
||
if bbox is not None:
|
||
xmin, ymin, xmax, ymax = bbox
|
||
model_xmin = float(xmin)
|
||
model_ymin = float(ymin)
|
||
model_w = float(xmax - xmin) or 1.0
|
||
model_h = float(ymax - ymin) or 1.0
|
||
|
||
# Keep ANY entity whose bounding box overlaps the selected region —
|
||
# no entity-type whitelist, so symbols made of splines/solids/etc.
|
||
# and blocks whose insertion point sits outside the box are kept.
|
||
# A small margin avoids clipping symbols straddling the edge.
|
||
from ezdxf import bbox as _ezbbox
|
||
mx = model_w * 0.03
|
||
my = model_h * 0.03
|
||
|
||
def _in_bbox(e) -> bool:
|
||
try:
|
||
bb = _ezbbox.extents([e], fast=True)
|
||
if not bb.has_data:
|
||
return True # can't measure → keep rather than drop
|
||
return not (bb.extmax.x < xmin - mx or bb.extmin.x > xmax + mx or
|
||
bb.extmax.y < ymin - my or bb.extmin.y > ymax + my)
|
||
except Exception:
|
||
return True # on error keep the entity, don't lose symbols
|
||
|
||
keep = [e for e in msp if _in_bbox(e)]
|
||
sub_doc = ezdxf.new(dxfversion=doc.dxfversion)
|
||
# ezdxf's Importer copies entities WITH their dependencies (block
|
||
# definitions incl. anonymous "*U"/"*D" blocks, layers, linetypes).
|
||
# Hand-copying via blocks.new() can't recreate anonymous blocks and
|
||
# crashed rendering with "Required block definition for *U".
|
||
from ezdxf.addons import Importer
|
||
importer = Importer(doc, sub_doc)
|
||
for e in keep:
|
||
try:
|
||
importer.import_entity(e, sub_doc.modelspace())
|
||
except Exception:
|
||
pass
|
||
try:
|
||
importer.finalize()
|
||
except Exception:
|
||
logger.exception("zoom importer.finalize failed (continuing)")
|
||
doc = sub_doc
|
||
msp = sub_doc.modelspace()
|
||
|
||
aspect = model_w / max(model_h, 1)
|
||
longest_px = render_px or RENDER_PX
|
||
out_width = longest_px if aspect >= 1 else int(longest_px * aspect)
|
||
out_height = int(longest_px / aspect) if aspect >= 1 else longest_px
|
||
|
||
# Zoom (bbox) → matplotlib with a FIXED window pinned to the requested
|
||
# box. Large entities (sheet border, walls, long wires) that merely
|
||
# cross the box are clipped to it instead of stretching the view to the
|
||
# whole drawing. The transform then exactly equals the requested box.
|
||
if bbox is not None or backend == "matplotlib":
|
||
dpi = 100
|
||
fig = plt.figure(figsize=(max(out_width, 1) / dpi,
|
||
max(out_height, 1) / dpi), dpi=dpi)
|
||
ax = fig.add_axes([0, 0, 1, 1])
|
||
ax.set_axis_off()
|
||
ax.set_facecolor("white")
|
||
mpl_backend = MatplotlibBackend(ax)
|
||
Frontend(RenderContext(doc), mpl_backend, config=config).draw_layout(
|
||
msp, finalize=(bbox is None))
|
||
if bbox is not None:
|
||
ax.set_xlim(model_xmin, model_xmin + model_w)
|
||
ax.set_ylim(model_ymin, model_ymin + model_h)
|
||
ax.set_aspect("equal", adjustable="box")
|
||
fig.savefig(out_path, dpi=dpi, facecolor="white", pad_inches=0)
|
||
plt.close(fig)
|
||
img = Image.open(out_path).convert("RGB")
|
||
else:
|
||
if aspect >= 1:
|
||
page_w_mm, page_h_mm = 1000, 1000 / aspect
|
||
else:
|
||
page_w_mm, page_h_mm = 1000 * aspect, 1000
|
||
page = layout.Page(width=page_w_mm, height=page_h_mm,
|
||
units=layout.Units.mm,
|
||
margins=layout.Margins.all(0))
|
||
svg_backend = SVGBackend()
|
||
Frontend(RenderContext(doc), svg_backend, config=config).draw_layout(
|
||
msp, finalize=True)
|
||
svg_str = svg_backend.get_string(page)
|
||
png_bytes = cairosvg.svg2png(bytestring=svg_str.encode("utf-8"),
|
||
output_width=out_width)
|
||
img = Image.open(io.BytesIO(png_bytes))
|
||
if img.mode == "RGBA":
|
||
white = Image.new("RGB", img.size, (255, 255, 255))
|
||
white.paste(img, mask=img.split()[3])
|
||
img = white
|
||
|
||
cap_px = render_px or RENDER_PX
|
||
if max(img.size) > cap_px:
|
||
r = cap_px / max(img.size)
|
||
img = img.resize((int(img.size[0] * r), int(img.size[1] * r)), Image.LANCZOS)
|
||
|
||
img.save(out_path, "PNG", optimize=True)
|
||
logger.info("Rendered → %s (%dx%d)", out_path.name, img.size[0], img.size[1])
|
||
# Transform so callers can map model (x,y) → pixel on this image:
|
||
# px = (x - model_xmin) / model_w * img_w
|
||
# py = (model_ymin + model_h - y) / model_h * img_h (y is flipped)
|
||
# Only valid when bbox is None (no crop) — which is the block path.
|
||
return {
|
||
"path": out_path,
|
||
"transform": {
|
||
"model_xmin": float(model_xmin), "model_ymin": float(model_ymin),
|
||
"model_w": float(model_w), "model_h": float(model_h),
|
||
"img_w": int(img.size[0]), "img_h": int(img.size[1]),
|
||
"cropped": bbox is not None,
|
||
},
|
||
}
|
||
|
||
|
||
def render(input_path: Path, out_dir: Path) -> dict:
|
||
"""Convert input → list of floor images.
|
||
|
||
Returns: {"floors": [{"index":0, "png":"floor_0.png", "legend_xy":[x,y]}, ...]}.
|
||
"""
|
||
suffix = input_path.suffix.lower()
|
||
|
||
if suffix == ".pdf":
|
||
return _render_pdf(input_path, out_dir)
|
||
|
||
if suffix == ".dwg":
|
||
dxf_path = dwg_to_dxf(input_path, out_dir)
|
||
elif suffix == ".dxf":
|
||
dxf_path = input_path
|
||
else:
|
||
raise ValueError(f"Unsupported format: {suffix}")
|
||
|
||
# Render the WHOLE drawing (bbox=None → full modelspace extents). These
|
||
# DWGs are one big sheet with multiple floor rectangles inside; the old
|
||
# find_floors crop kept only the first rectangle and silently dropped the
|
||
# rest of the drawing. Counting must cover everything, so no crop.
|
||
png = out_dir / "floor_0.png"
|
||
rr = render_region(dxf_path, png, None)
|
||
return {"floors": [{"index": 0, "png": png.name, "legend_xy": None,
|
||
"transform": rr["transform"]}],
|
||
"dxf_path": str(dxf_path)}
|
||
|
||
|
||
def _render_pdf(pdf_path: Path, out_dir: Path) -> dict:
|
||
"""Render PDF → PNG, auto-rotate so LEGENDA reads horizontally.
|
||
|
||
Uses pdfplumber to find 'LEGENDA' text and its rotation, then renders
|
||
via pdf2image and applies image rotation so the legend is upright.
|
||
Returns the same shape as the DXF render path.
|
||
"""
|
||
from pdf2image import convert_from_path
|
||
import pdfplumber
|
||
# Allow large rasterizations (architectural PDFs can be 200M+ px at high DPI)
|
||
Image.MAX_IMAGE_PIXELS = None
|
||
|
||
# Pick DPI so longest page edge lands near RENDER_PX pixels
|
||
with pdfplumber.open(str(pdf_path)) as pdf:
|
||
first = pdf.pages[0]
|
||
pw_in = max(first.width, first.height) / 72 # PDF points → inches
|
||
target_dpi = max(150, min(600, int(RENDER_PX / max(pw_in, 1))))
|
||
logger.info("PDF page longest edge %.1f in → using dpi=%d", pw_in, target_dpi)
|
||
|
||
pages = convert_from_path(str(pdf_path), dpi=target_dpi)
|
||
out_paths = []
|
||
legend_info: list[dict] = []
|
||
|
||
with pdfplumber.open(str(pdf_path)) as pdf:
|
||
for i, plumb_page in enumerate(pdf.pages):
|
||
page_img = pages[i] if i < len(pages) else None
|
||
if page_img is None:
|
||
continue
|
||
pw, ph = plumb_page.width, plumb_page.height
|
||
iw, ih = page_img.size
|
||
# Find any text matching legend headings
|
||
words = plumb_page.extract_words(extra_attrs=["upright"]) or []
|
||
legend_word = None
|
||
for w in words:
|
||
text = w["text"].strip().upper()
|
||
if text in ("LEGENDA", "VYSVĚTLIVKY", "LEGENDA:", "POPIS"):
|
||
legend_word = w
|
||
break
|
||
rotation = 0
|
||
if legend_word is not None and not legend_word.get("upright", True):
|
||
# Sideways text → rotate the image so text is upright.
|
||
# PIL.rotate uses CCW for positive angles.
|
||
rotation = 90
|
||
page_img = page_img.rotate(90, expand=True)
|
||
iw, ih = page_img.size
|
||
logger.info("PDF page %d: rotated 90° CCW (LEGENDA was sideways)", i)
|
||
if legend_word is not None:
|
||
# Convert PDF coords to NORMALIZED image coords (after rotation)
|
||
x0, y0 = legend_word["x0"], legend_word["top"]
|
||
x1, y1 = legend_word["x1"], legend_word["bottom"]
|
||
if rotation == 90:
|
||
# CCW 90° rotation: original (x, y) → new (y, W-x).
|
||
# Rotated image has width=ph, height=pw.
|
||
nx0 = y0 / ph
|
||
nx1 = y1 / ph
|
||
ny0 = 1 - (x1 / pw)
|
||
ny1 = 1 - (x0 / pw)
|
||
else:
|
||
nx0, nx1 = x0 / pw, x1 / pw
|
||
ny0, ny1 = y0 / ph, y1 / ph
|
||
legend_info.append({"page": i, "norm_bbox": (nx0, ny0, nx1, ny1)})
|
||
logger.info("PDF page %d: LEGENDA at norm bbox %s",
|
||
i, (nx0, ny0, nx1, ny1))
|
||
|
||
if max(page_img.size) > RENDER_PX:
|
||
r = RENDER_PX / max(page_img.size)
|
||
page_img = page_img.resize(
|
||
(int(page_img.size[0] * r), int(page_img.size[1] * r)),
|
||
Image.LANCZOS,
|
||
)
|
||
p = out_dir / f"floor_{i}.png"
|
||
page_img.save(p, "PNG", optimize=True)
|
||
out_paths.append(p)
|
||
|
||
return {
|
||
"floors": [
|
||
{"index": i, "png": p.name, "legend_xy": None,
|
||
"legend_norm_bbox": next((li["norm_bbox"] for li in legend_info
|
||
if li["page"] == i), None)}
|
||
for i, p in enumerate(out_paths)
|
||
]
|
||
}
|
||
|
||
|
||
def crop_region(image_path: Path, bbox: dict, out_path: Path,
|
||
pad: float = 1.5, min_px: int = 120) -> Path:
|
||
"""Crop a region with generous padding so symbols are visible.
|
||
|
||
pad: multiplier of the bbox half-extent added on each side.
|
||
min_px: ensure the output is at least this many pixels wide/tall by
|
||
expanding the crop region if the requested area is smaller.
|
||
"""
|
||
img = Image.open(image_path)
|
||
W, H = img.size
|
||
bx, by, bw, bh = bbox["x"], bbox["y"], bbox["w"], bbox["h"]
|
||
cx, cy = bx + bw / 2, by + bh / 2
|
||
# Padded extents in normalized coords
|
||
half_w = bw / 2 + bw * pad
|
||
half_h = bh / 2 + bh * pad
|
||
# Enforce minimum pixel dimensions
|
||
if (2 * half_w) * W < min_px:
|
||
half_w = (min_px / 2) / W
|
||
if (2 * half_h) * H < min_px:
|
||
half_h = (min_px / 2) / H
|
||
x0 = max(0, int((cx - half_w) * W))
|
||
x1 = min(W, int((cx + half_w) * W))
|
||
y0 = max(0, int((cy - half_h) * H))
|
||
y1 = min(H, int((cy + half_h) * H))
|
||
crop = img.crop((x0, y0, x1, y1))
|
||
crop.save(out_path, "PNG")
|
||
return out_path
|