from __future__ import annotations import tempfile from collections.abc import Iterator from dataclasses import dataclass from pathlib import Path import fitz from PIL import Image, ImageDraw, ImageFont, ImageOps, ImageSequence from .config import SETTINGS SUPPORTED_SUFFIXES = {".pdf", ".png", ".jpg", ".jpeg", ".tif", ".tiff"} @dataclass class RenderedDocument: source_name: str pages: list[Image.Image] vector_text: list[str] @dataclass class Tile: page: int index: int image: Image.Image x: int y: int width: int height: int page_width: int page_height: int def _validate_file(path: Path) -> None: if not path.exists() or not path.is_file(): raise ValueError("The uploaded file is no longer available.") if path.suffix.lower() not in SUPPORTED_SUFFIXES: raise ValueError( "Unsupported file type. Upload PDF, PNG, JPG, JPEG, TIF, or TIFF." ) size_mb = path.stat().st_size / (1024 * 1024) if size_mb > SETTINGS.max_upload_mb: raise ValueError( f"File is {size_mb:.1f} MB; the configured limit is " f"{SETTINGS.max_upload_mb} MB." ) def load_document( file_path: str | Path, dpi: int = 300, max_pages: int | None = None, ) -> RenderedDocument: path = Path(file_path) _validate_file(path) page_limit = min(max_pages or SETTINGS.max_pages, SETTINGS.max_pages) if dpi not in {150, 200, 300, 450, 600}: raise ValueError("DPI must be one of 150, 200, 300, 450, or 600.") pages: list[Image.Image] = [] vector_text: list[str] = [] if path.suffix.lower() == ".pdf": document = fitz.open(path) try: if document.page_count > page_limit: raise ValueError( f"Drawing has {document.page_count} pages. This Space is " f"configured for at most {page_limit} pages per run." ) scale = dpi / 72.0 matrix = fitz.Matrix(scale, scale) for page in document: pixmap = page.get_pixmap(matrix=matrix, alpha=False) if pixmap.width * pixmap.height > SETTINGS.max_image_pixels: raise ValueError( f"Rendered page exceeds {SETTINGS.max_image_pixels:,} pixels. " "Reduce the DPI or split the drawing." ) image = Image.frombytes( "RGB", (pixmap.width, pixmap.height), pixmap.samples ) pages.append(image) vector_text.append(_extract_ordered_pdf_text(page)) finally: document.close() else: with Image.open(path) as source: for frame_index, frame in enumerate(ImageSequence.Iterator(source)): if frame_index >= page_limit: raise ValueError( f"Image contains more than the configured {page_limit} frames." ) image = ImageOps.exif_transpose(frame.copy()).convert("RGB") if image.width * image.height > SETTINGS.max_image_pixels: raise ValueError( f"Image exceeds {SETTINGS.max_image_pixels:,} pixels." ) pages.append(image) vector_text.append("") if not pages: raise ValueError("No readable drawing page was found in the uploaded file.") return RenderedDocument(path.name, pages, vector_text) def _extract_ordered_pdf_text(page: fitz.Page) -> str: words = page.get_text("words", sort=True) if not words: return "" lines: dict[tuple[int, int, int], list[tuple[int, str]]] = {} for item in words: block, line, word_index = int(item[5]), int(item[6]), int(item[7]) lines.setdefault((block, line, 0), []).append((word_index, str(item[4]))) ordered: list[str] = [] for key in sorted(lines): ordered.append(" ".join(text for _, text in sorted(lines[key]))) return "\n".join(ordered) def tile_image( image: Image.Image, page_number: int, tile_size: int = 1800, overlap: float = 0.18, ) -> Iterator[Tile]: """Yield overlapping square-ish tiles without losing page-edge content.""" tile_size = max(768, min(3072, int(tile_size))) overlap = max(0.0, min(0.4, float(overlap))) width, height = image.size if width <= tile_size and height <= tile_size: yield Tile(page_number, 0, image.copy(), 0, 0, width, height, width, height) return step = max(1, int(tile_size * (1 - overlap))) x_positions = _axis_positions(width, tile_size, step) y_positions = _axis_positions(height, tile_size, step) index = 0 for y in y_positions: for x in x_positions: right = min(width, x + tile_size) bottom = min(height, y + tile_size) crop = image.crop((x, y, right, bottom)) yield Tile( page_number, index, crop, x, y, crop.width, crop.height, width, height, ) index += 1 def _axis_positions(length: int, window: int, step: int) -> list[int]: if length <= window: return [0] positions = list(range(0, max(1, length - window + 1), step)) last = length - window if positions[-1] != last: positions.append(last) return positions def create_sample_drawing() -> str: """Create a deterministic test drawing; it is clearly labelled as a sample.""" output = Path(tempfile.gettempdir()) / "drawing_intelligence_sample.png" if output.exists(): return str(output) canvas = Image.new("RGB", (1800, 1180), "white") draw = ImageDraw.Draw(canvas) font = _font(26) small = _font(21) title = _font(34, bold=True) red = "#c9202e" charcoal = "#151b24" gray = "#75808f" draw.text((55, 35), "DRAWING INTELLIGENCE — DEMONSTRATION DRAWING", fill=charcoal, font=title) draw.text( (55, 82), "Synthetic training-free sample · not a production engineering drawing", fill=gray, font=small, ) draw.rectangle((45, 125, 1750, 1125), outline="#aeb5bf", width=3) # Simplified cylinder-head-like orthographic view. body = (250, 310, 1260, 805) draw.rounded_rectangle(body, radius=35, outline=charcoal, width=5) for cx in (450, 750, 1050): draw.ellipse((cx - 125, 410, cx + 125, 660), outline=charcoal, width=4) for dx, dy in ((-48, -48), (48, -48), (-48, 48), (48, 48)): draw.ellipse( (cx + dx - 28, 535 + dy - 28, cx + dx + 28, 535 + dy + 28), outline=charcoal, width=3, ) for x in range(310, 1230, 115): draw.ellipse((x - 16, 330, x + 16, 362), outline=charcoal, width=3) draw.ellipse((x - 16, 752, x + 16, 784), outline=charcoal, width=3) draw.line((250, 850, 1260, 850), fill=charcoal, width=3) draw.line((250, 835, 250, 865), fill=charcoal, width=3) draw.line((1260, 835, 1260, 865), fill=charcoal, width=3) draw.text((715, 860), "215", fill=charcoal, font=small) callouts = [ ("1", (385, 208), "4X ⌀12 H7 THRU", (440, 204)), ("5", (1030, 225), "M14×1.5-6H", (1090, 220)), ("7", (650, 185), "⌖ ⌀0.10 M | A | B | C", (710, 180)), ("14", (1340, 415), "Ra 1.6", (1410, 410)), ("27", (710, 705), "⌀12 H7", (760, 742)), ("32", (350, 980), "BREAK ALL SHARP EDGES 0.5×45°", (410, 975)), ] anchors = { "1": (420, 340), "5": (1100, 345), "7": (750, 310), "14": (1255, 520), "27": (735, 745), "32": (680, 800), } for balloon_id, center, text, text_at in callouts: _draw_balloon(draw, center, balloon_id, red, small) draw.line((center[0], center[1] + 22, *anchors[balloon_id]), fill=charcoal, width=3) draw.text(text_at, text, fill=charcoal, font=font) # Title block. draw.rectangle((1265, 805, 1750, 1125), outline=charcoal, width=3) for y in (875, 945, 1015): draw.line((1265, y, 1750, y), fill=charcoal, width=2) draw.text((1290, 825), "PART: SAMPLE CYLINDER HEAD", fill=charcoal, font=small) draw.text((1290, 895), "MATERIAL: FG260", fill=charcoal, font=small) draw.text((1290, 965), "UNLESS SPECIFIED: ±0.20", fill=charcoal, font=small) draw.text((1290, 1035), "SHEET 1 OF 1 · mm", fill=charcoal, font=small) canvas.save(output, optimize=True) return str(output) def _draw_balloon( draw: ImageDraw.ImageDraw, center: tuple[int, int], label: str, color: str, font: ImageFont.FreeTypeFont | ImageFont.ImageFont, ) -> None: x, y = center radius = 24 draw.ellipse((x - radius, y - radius, x + radius, y + radius), fill=color) box = draw.textbbox((0, 0), label, font=font) text_width, text_height = box[2] - box[0], box[3] - box[1] draw.text( (x - text_width / 2, y - text_height / 2 - 3), label, fill="white", font=font, ) def _font(size: int, bold: bool = False) -> ImageFont.FreeTypeFont | ImageFont.ImageFont: candidates = [ "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", "/usr/share/fonts/truetype/liberation2/LiberationSans-Bold.ttf" if bold else "/usr/share/fonts/truetype/liberation2/LiberationSans-Regular.ttf", ] for candidate in candidates: try: return ImageFont.truetype(candidate, size) except OSError: continue return ImageFont.load_default()