Spaces:
Runtime error
Runtime error
| 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"} | |
| class RenderedDocument: | |
| source_name: str | |
| pages: list[Image.Image] | |
| vector_text: list[str] | |
| 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() | |