977 lines
43 KiB
Python
977 lines
43 KiB
Python
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import base64
|
||
from collections import Counter
|
||
from concurrent.futures import ThreadPoolExecutor
|
||
import json
|
||
import math
|
||
from pathlib import Path
|
||
import random
|
||
import re
|
||
import struct
|
||
import threading
|
||
from typing import Any, Protocol
|
||
|
||
from .dataset import Sample, scan_dataset
|
||
from .featurescript_parser import parse_featurescript
|
||
from .lowering import _number, _point, plain
|
||
|
||
|
||
DESCRIPTION_SCHEMA_VERSION = "cadfs_to_cdsl.description.v1"
|
||
MANIFEST_NAME = "description_manifest.jsonl"
|
||
|
||
ENTITY_LABELS = {
|
||
"skLineSegment": "线段",
|
||
"skCircle": "圆",
|
||
"skArc": "圆弧",
|
||
"skEllipse": "椭圆",
|
||
"skFitSpline": "样条",
|
||
"skPoint": "点",
|
||
}
|
||
|
||
OPERATION_LABELS = {
|
||
"extrude": "拉伸",
|
||
"revolve": "旋转",
|
||
"fillet": "圆角",
|
||
"chamfer": "倒角",
|
||
"hole": "孔",
|
||
"linearPattern": "线性阵列",
|
||
"circularPattern": "圆周阵列",
|
||
"mirror": "镜像",
|
||
"shell": "抽壳",
|
||
"loft": "放样",
|
||
"sweep": "扫掠",
|
||
"booleanBodies": "布尔",
|
||
"cPlane": "参考平面",
|
||
"referenceAxis": "参考轴",
|
||
}
|
||
|
||
VISION_DESCRIPTION_TOOL = {
|
||
"type": "function",
|
||
"function": {
|
||
"name": "describe_cad_model",
|
||
"description": "Return a conservative semantic and geometric description for one CAD model.",
|
||
"parameters": {
|
||
"type": "object",
|
||
"additionalProperties": False,
|
||
"required": [
|
||
"category",
|
||
"category_confidence",
|
||
"candidate_names",
|
||
"summary_zh",
|
||
"possible_functions",
|
||
"applications",
|
||
"structural_features",
|
||
"geometric_features",
|
||
"keywords_zh",
|
||
"keywords_en",
|
||
"uncertainties",
|
||
],
|
||
"properties": {
|
||
"category": {"type": "string", "minLength": 1, "maxLength": 160},
|
||
"category_confidence": {"type": "number", "minimum": 0, "maximum": 1},
|
||
"candidate_names": {"type": "array", "items": {"type": "string", "minLength": 1, "maxLength": 120}, "maxItems": 8},
|
||
"summary_zh": {"type": "string", "minLength": 1, "maxLength": 1200},
|
||
"possible_functions": {"type": "array", "items": {"type": "string", "minLength": 1, "maxLength": 240}, "maxItems": 8},
|
||
"applications": {"type": "array", "items": {"type": "string", "minLength": 1, "maxLength": 240}, "maxItems": 8},
|
||
"structural_features": {"type": "array", "items": {"type": "string", "minLength": 1, "maxLength": 240}, "maxItems": 16},
|
||
"geometric_features": {"type": "array", "items": {"type": "string", "minLength": 1, "maxLength": 240}, "maxItems": 16},
|
||
"keywords_zh": {"type": "array", "items": {"type": "string", "minLength": 1, "maxLength": 80}, "maxItems": 32},
|
||
"keywords_en": {"type": "array", "items": {"type": "string", "minLength": 1, "maxLength": 80}, "maxItems": 32},
|
||
"uncertainties": {"type": "array", "items": {"type": "string", "minLength": 1, "maxLength": 240}, "maxItems": 12},
|
||
},
|
||
},
|
||
},
|
||
}
|
||
|
||
|
||
class VisionDescriptionClient(Protocol):
|
||
def describe(self, *, sample_id: str, image_path: Path, local_facts: dict[str, Any]) -> dict[str, Any]:
|
||
...
|
||
|
||
|
||
class ConfiguredVisionDescriptionClient:
|
||
def __init__(self) -> None:
|
||
from app.cad_agent.adapters.structured_llm import StructuredModelGateway
|
||
from app.settings import get_settings
|
||
|
||
self.settings = get_settings()
|
||
provider, model = self.settings.resolve_review_model()
|
||
if not model.vision:
|
||
raise ValueError("configured review model is not vision-capable")
|
||
self.provider_id = provider.id
|
||
self.model_id = model.id
|
||
self.gateway = StructuredModelGateway(self.settings)
|
||
|
||
def describe(self, *, sample_id: str, image_path: Path, local_facts: dict[str, Any]) -> dict[str, Any]:
|
||
try:
|
||
loop = asyncio.get_running_loop()
|
||
except RuntimeError:
|
||
loop = None
|
||
if loop and loop.is_running():
|
||
raise RuntimeError("vision description cannot run inside an active event loop")
|
||
return asyncio.run(self._describe(sample_id=sample_id, image_path=image_path, local_facts=local_facts))
|
||
|
||
async def _describe(self, *, sample_id: str, image_path: Path, local_facts: dict[str, Any]) -> dict[str, Any]:
|
||
payload = {
|
||
"sample_id": sample_id,
|
||
"local_category": local_facts.get("category"),
|
||
"local_candidate_names": local_facts.get("candidate_names"),
|
||
"local_geometric_features": local_facts.get("geometric_features"),
|
||
"local_operations": local_facts.get("operations"),
|
||
"local_dimensions": local_facts.get("dimensions"),
|
||
"annotation_excerpt": str(local_facts.get("annotation_excerpt") or "")[:4000],
|
||
"instruction": (
|
||
"Use the image and deterministic CAD facts to identify likely model class and retrieval terms. "
|
||
"Keep product identity and use cases as candidates when the evidence is not definitive. "
|
||
"Do not invent exact dimensions beyond the supplied facts."
|
||
),
|
||
}
|
||
content: list[dict[str, Any]] = [{"type": "text", "text": json.dumps(payload, ensure_ascii=False)}]
|
||
if image_path.is_file():
|
||
content.append(self._image_part(image_path))
|
||
response = await self.gateway.call_tool(
|
||
messages=[
|
||
{
|
||
"role": "system",
|
||
"content": (
|
||
"You describe CAD parts for vector search. Return only the required tool call. "
|
||
"Separate visible geometry from inferred semantics, and mark uncertain real-world identity conservatively."
|
||
),
|
||
},
|
||
{"role": "user", "content": content},
|
||
],
|
||
tool=VISION_DESCRIPTION_TOOL,
|
||
provider_id=self.provider_id,
|
||
model_id=self.model_id,
|
||
required_tool_name="describe_cad_model",
|
||
)
|
||
calls = response.get("tool_calls") or []
|
||
if len(calls) != 1:
|
||
raise ValueError("vision provider did not return exactly one tool call")
|
||
function = calls[0].get("function") if isinstance(calls[0], dict) else None
|
||
if not isinstance(function, dict) or function.get("name") != "describe_cad_model":
|
||
raise ValueError("vision provider returned an unexpected tool call")
|
||
arguments = json.loads(str(function.get("arguments") or "{}"))
|
||
if not isinstance(arguments, dict):
|
||
raise ValueError("vision provider arguments are not an object")
|
||
arguments["usage"] = response.get("usage") or {}
|
||
return arguments
|
||
|
||
@staticmethod
|
||
def _image_part(path: Path) -> dict[str, Any]:
|
||
media_type = "image/jpeg" if path.suffix.lower() in {".jpg", ".jpeg"} else "image/png"
|
||
data = base64.b64encode(path.read_bytes()).decode("ascii")
|
||
return {"type": "image_url", "image_url": {"url": f"data:{media_type};base64,{data}"}}
|
||
|
||
|
||
class VisionCallState:
|
||
def __init__(self, client: VisionDescriptionClient, *, failure_limit: int = 3) -> None:
|
||
self.client = client
|
||
self.failure_limit = failure_limit
|
||
self._lock = threading.Lock()
|
||
self._consecutive_failures = 0
|
||
self._disabled_reason = ""
|
||
|
||
def describe(self, *, sample_id: str, image_path: Path, local_facts: dict[str, Any]) -> tuple[dict[str, Any] | None, dict[str, Any] | None]:
|
||
with self._lock:
|
||
disabled = self._disabled_reason
|
||
if disabled:
|
||
return None, {"code": "vision_disabled", "message": disabled}
|
||
try:
|
||
result = self.client.describe(sample_id=sample_id, image_path=image_path, local_facts=local_facts)
|
||
except Exception as exc:
|
||
with self._lock:
|
||
self._consecutive_failures += 1
|
||
if self._consecutive_failures >= self.failure_limit:
|
||
self._disabled_reason = f"vision disabled after {self._consecutive_failures} consecutive failures"
|
||
return None, {"code": "vision_failed", "message": str(exc), "type": type(exc).__name__}
|
||
with self._lock:
|
||
self._consecutive_failures = 0
|
||
return result, None
|
||
|
||
|
||
def sample_shard(sample: Sample) -> str:
|
||
for modality in ("image", "featurescript", "annotation", "step", "stl"):
|
||
path = sample.files.get(modality)
|
||
if path:
|
||
return Path(path).parent.name
|
||
return sample.sample_id[:4]
|
||
|
||
|
||
def select_description_samples(
|
||
input_root: Path,
|
||
*,
|
||
shard: str | None = None,
|
||
sample_ids: list[str] | None = None,
|
||
offset: int = 0,
|
||
limit: int | None = None,
|
||
seed: int | None = None,
|
||
) -> list[Sample]:
|
||
samples = scan_dataset(input_root, include_hashes=False)
|
||
if shard:
|
||
samples = [sample for sample in samples if sample_shard(sample) == shard]
|
||
if sample_ids:
|
||
wanted = set(sample_ids)
|
||
samples = [sample for sample in samples if sample.sample_id in wanted]
|
||
missing = wanted - {sample.sample_id for sample in samples}
|
||
if missing:
|
||
raise ValueError("unknown sample ids: " + ", ".join(sorted(missing)))
|
||
if seed is not None and limit is not None:
|
||
samples = random.Random(seed).sample(samples, min(limit, len(samples)))
|
||
return sorted(samples, key=lambda sample: sample.sample_id)
|
||
return samples[offset:None if limit is None else offset + limit]
|
||
|
||
|
||
def describe_samples(
|
||
input_root: Path,
|
||
*,
|
||
shard: str | None = None,
|
||
sample_ids: list[str] | None = None,
|
||
mode: str = "hybrid",
|
||
offset: int = 0,
|
||
limit: int | None = None,
|
||
seed: int | None = None,
|
||
force: bool = False,
|
||
workers: int = 1,
|
||
vision_client: VisionDescriptionClient | None = None,
|
||
) -> tuple[list[dict[str, Any]], dict[str, Any]]:
|
||
if mode not in {"local", "hybrid", "vision"}:
|
||
raise ValueError("--mode must be local, hybrid, or vision")
|
||
if not shard and not sample_ids and limit is None:
|
||
raise ValueError("describe requires --shard, --sample-id, or --limit to avoid accidental full-corpus generation")
|
||
if not 1 <= workers <= 8:
|
||
raise ValueError("--workers must be between 1 and 8")
|
||
|
||
samples = select_description_samples(input_root, shard=shard, sample_ids=sample_ids, offset=offset, limit=limit, seed=seed)
|
||
existing_manifest = _read_manifest(input_root / "description_txt" / MANIFEST_NAME)
|
||
shared_diagnostics: list[dict[str, Any]] = []
|
||
vision_state = VisionCallState(vision_client) if vision_client is not None else None
|
||
if mode != "local" and vision_client is None:
|
||
try:
|
||
vision_client = ConfiguredVisionDescriptionClient()
|
||
vision_state = VisionCallState(vision_client)
|
||
except Exception as exc:
|
||
shared_diagnostics.append({"code": "vision_unavailable", "message": str(exc), "type": type(exc).__name__})
|
||
|
||
def process(sample: Sample) -> dict[str, Any]:
|
||
return describe_one(
|
||
sample,
|
||
input_root,
|
||
mode=mode,
|
||
force=force,
|
||
vision_state=vision_state,
|
||
existing_record=existing_manifest.get(sample.sample_id),
|
||
shared_diagnostics=shared_diagnostics,
|
||
)
|
||
|
||
if workers == 1:
|
||
records = [process(sample) for sample in samples]
|
||
else:
|
||
with ThreadPoolExecutor(max_workers=workers, thread_name_prefix="cadfs-describe") as executor:
|
||
records = list(executor.map(process, samples))
|
||
manifest_rows = _write_manifest(input_root, records)
|
||
return records, _summary(input_root, records, manifest_rows)
|
||
|
||
|
||
def describe_one(
|
||
sample: Sample,
|
||
input_root: Path,
|
||
*,
|
||
mode: str,
|
||
force: bool = False,
|
||
vision_state: VisionCallState | None = None,
|
||
existing_record: dict[str, Any] | None = None,
|
||
shared_diagnostics: list[dict[str, Any]] | None = None,
|
||
) -> dict[str, Any]:
|
||
shard = sample_shard(sample)
|
||
txt_path = input_root / "description_txt" / shard / f"{sample.sample_id}.txt"
|
||
if txt_path.exists() and not force:
|
||
return existing_record or _skipped_record(sample, input_root, txt_path, shard)
|
||
|
||
diagnostics = list(shared_diagnostics or [])
|
||
try:
|
||
local_facts = _local_facts(sample)
|
||
diagnostics.extend(local_facts.pop("diagnostics", []))
|
||
vision_result = None
|
||
if mode != "local":
|
||
if vision_state is None:
|
||
diagnostics.append({"code": "vision_unavailable", "message": "no configured vision client"})
|
||
else:
|
||
vision_result, vision_diagnostic = vision_state.describe(
|
||
sample_id=sample.sample_id,
|
||
image_path=Path(sample.files.get("image", "")),
|
||
local_facts=local_facts,
|
||
)
|
||
if vision_diagnostic is not None:
|
||
diagnostics.append(vision_diagnostic)
|
||
record = _compose_record(sample, input_root, shard, txt_path, local_facts, vision_result, diagnostics, mode)
|
||
_atomic_write_text(txt_path, render_description_text(record))
|
||
return record
|
||
except Exception as exc:
|
||
diagnostics.append({"code": "description_failed", "message": str(exc), "type": type(exc).__name__})
|
||
record = _failed_record(sample, input_root, txt_path, shard, diagnostics)
|
||
_atomic_write_text(txt_path, render_description_text(record))
|
||
return record
|
||
|
||
|
||
def render_description_text(record: dict[str, Any]) -> str:
|
||
confidence = _format_float(float(record.get("category_confidence") or 0.0))
|
||
keywords = _unique([*(record.get("keywords_zh") or []), *(record.get("keywords_en") or [])], limit=48)
|
||
evidence = [
|
||
"确定事实来自 FeatureScript、原始 CAD 操作说明、STL/STEP 文件存在性和可解析的几何范围。",
|
||
*_as_text_list(record.get("uncertainties")),
|
||
]
|
||
diagnostics = record.get("diagnostics") or []
|
||
if diagnostics:
|
||
codes = _unique([str(item.get("code") or "diagnostic") for item in diagnostics if isinstance(item, dict)], limit=8)
|
||
if codes:
|
||
evidence.append("诊断：" + "，".join(codes))
|
||
lines = [
|
||
f"样本ID：{record.get('sample_id', '')}",
|
||
f"分类：{record.get('category', '通用机械 CAD 零件（候选）')}",
|
||
"候选名称：" + _join_or_unknown(record.get("candidate_names")),
|
||
"模型概述：" + str(record.get("summary_zh") or "该模型为缺少语义上下文的 CAD 几何样本，描述以可见结构和建模特征为主。"),
|
||
"可能作用：" + _join_or_unknown(record.get("possible_functions")),
|
||
"典型应用：" + _join_or_unknown(record.get("applications")),
|
||
"结构特征：" + _join_or_unknown(record.get("structural_features")),
|
||
"建模与几何特征：" + _join_or_unknown(record.get("geometric_features")),
|
||
"检索关键词：" + _join_or_unknown(keywords),
|
||
"证据与不确定性：" + "；".join(evidence),
|
||
f"置信度：{confidence}（几何事实置信度较高；真实零件类别、用途和装配位置为候选判断）",
|
||
]
|
||
return "\n".join(lines) + "\n"
|
||
|
||
|
||
def _local_facts(sample: Sample) -> dict[str, Any]:
|
||
feature_path = Path(sample.files["featurescript"])
|
||
annotation_path = Path(sample.files.get("annotation", ""))
|
||
stl_path = Path(sample.files.get("stl", ""))
|
||
image_path = Path(sample.files.get("image", ""))
|
||
source = feature_path.read_text(encoding="utf-8")
|
||
annotation = annotation_path.read_text(encoding="utf-8") if annotation_path.is_file() else ""
|
||
model = parse_featurescript(source, sample.sample_id)
|
||
diagnostics: list[dict[str, Any]] = []
|
||
sketch_info, sketch_features, sketch_dimensions, entity_counts = _describe_sketches(model)
|
||
operations, operation_dimensions = _describe_operations(model)
|
||
stl_info = _stl_bbox(stl_path) if stl_path.is_file() else None
|
||
image_info = _image_metadata(image_path) if image_path.is_file() else None
|
||
shape_cues = _shape_cues(annotation, entity_counts, operations)
|
||
geometric_features = _unique(
|
||
[
|
||
f"包含 {len(model.sketches)} 个草图和 {len(model.features)} 个建模/修饰特征",
|
||
*sketch_features,
|
||
*shape_cues,
|
||
*_stl_features(stl_info),
|
||
],
|
||
limit=32,
|
||
)
|
||
dimensions = _unique([*sketch_dimensions, *operation_dimensions, *_stl_dimensions(stl_info)], limit=32)
|
||
classification = _classify_local(annotation, operations, geometric_features, stl_info)
|
||
keywords_zh, keywords_en = _keywords(classification, operations, geometric_features)
|
||
uncertainties = [
|
||
"CADFS 样本未提供真实装配上下文，类别、用途和安装位置只能作为候选语义。",
|
||
"确定描述优先依据几何、草图和特征操作；视觉判断仅作为补充证据。",
|
||
]
|
||
summary = _local_summary(classification, operations, geometric_features, dimensions)
|
||
if not source.strip():
|
||
diagnostics.append({"code": "empty_featurescript", "message": "FeatureScript file is empty"})
|
||
return {
|
||
"category": classification["category"],
|
||
"category_confidence": classification["confidence"],
|
||
"candidate_names": classification["candidate_names"],
|
||
"summary_zh": summary,
|
||
"possible_functions": classification["possible_functions"],
|
||
"applications": classification["applications"],
|
||
"structural_features": _unique([*classification["structural_features"], *sketch_info], limit=24),
|
||
"geometric_features": geometric_features,
|
||
"operations": operations,
|
||
"dimensions": dimensions,
|
||
"keywords_zh": keywords_zh,
|
||
"keywords_en": keywords_en,
|
||
"uncertainties": uncertainties,
|
||
"annotation_excerpt": annotation[:4000],
|
||
"image_metadata": image_info,
|
||
"diagnostics": diagnostics,
|
||
}
|
||
|
||
|
||
def _describe_sketches(model: Any) -> tuple[list[str], list[str], list[str], Counter[str]]:
|
||
sketch_info: list[str] = []
|
||
features: list[str] = []
|
||
dimensions: list[str] = []
|
||
entity_counts: Counter[str] = Counter()
|
||
for sketch in model.sketches:
|
||
counts = Counter(entity.operation for entity in sketch.entities)
|
||
entity_counts.update(counts)
|
||
labels = [f"{ENTITY_LABELS.get(name, name)} {count} 个" for name, count in sorted(counts.items())]
|
||
plane = _plane_name(sketch.workplane)
|
||
if labels:
|
||
sketch_info.append(f"{sketch.feature_id} 位于 {plane} 平面，包含" + "、".join(labels))
|
||
else:
|
||
sketch_info.append(f"{sketch.feature_id} 位于 {plane} 平面，没有可解析草图实体")
|
||
if counts:
|
||
features.append("草图包含" + "、".join(labels))
|
||
bbox = _sketch_bbox(sketch)
|
||
if bbox:
|
||
width = bbox[2] - bbox[0]
|
||
height = bbox[3] - bbox[1]
|
||
dimensions.append(f"{sketch.feature_id} 二维范围约 {_format_float(width)} × {_format_float(height)} mm")
|
||
return sketch_info, _unique(features, limit=12), dimensions, entity_counts
|
||
|
||
|
||
def _describe_operations(model: Any) -> tuple[list[str], list[str]]:
|
||
operations: list[str] = []
|
||
dimensions: list[str] = []
|
||
for feature in model.features:
|
||
params = feature.params
|
||
operation = feature.operation
|
||
if operation == "extrude":
|
||
operation_type = _enum_tail(params.get("operationType") or "NEW").upper()
|
||
action = "extrude_cut" if operation_type in {"REMOVE", "CUT"} else "extrude_add"
|
||
depth = _safe_number(params.get("depth"))
|
||
label = f"{action}"
|
||
if depth is not None:
|
||
label += f" depth_mm={_format_float(depth)}"
|
||
dimensions.append(f"{feature.feature_id} 拉伸深度 {_format_mm(depth)}")
|
||
if _truthy(params.get("hasSecondDirection")):
|
||
label += " two_sided=true"
|
||
operations.append(label)
|
||
elif operation == "revolve":
|
||
operation_type = _enum_tail(params.get("operationType") or params.get("surfaceOperationType") or "NEW").upper()
|
||
action = "revolve_cut" if operation_type in {"REMOVE", "CUT"} else "revolve_add"
|
||
angle = 360.0 if "FULL" in _enum_tail(params.get("revolveType") or "FULL").upper() else _safe_number(params.get("angle"))
|
||
operations.append(f"{action} angle_deg={_format_float(angle or 360.0)}")
|
||
elif operation == "hole":
|
||
diameter = _safe_number(params.get("holeDiameter"))
|
||
style = _enum_tail(params.get("style") or "simple").lower()
|
||
label = f"hole {style}"
|
||
if diameter is not None:
|
||
label += f" diameter_mm={_format_float(diameter)}"
|
||
dimensions.append(f"{feature.feature_id} 孔径 {_format_mm(diameter)}")
|
||
operations.append(label)
|
||
elif operation == "fillet":
|
||
radius = _safe_number(params.get("radius"))
|
||
operations.append("fillet" + (f" radius_mm={_format_float(radius)}" if radius is not None else ""))
|
||
if radius is not None:
|
||
dimensions.append(f"{feature.feature_id} 圆角半径 {_format_mm(radius)}")
|
||
elif operation == "chamfer":
|
||
width = _safe_number(params.get("width") or params.get("width1"))
|
||
operations.append("chamfer" + (f" distance_mm={_format_float(width)}" if width is not None else ""))
|
||
if width is not None:
|
||
dimensions.append(f"{feature.feature_id} 倒角距离 {_format_mm(width)}")
|
||
else:
|
||
operations.append(OPERATION_LABELS.get(operation, operation))
|
||
return _unique(operations, limit=32), _unique(dimensions, limit=32)
|
||
|
||
|
||
def _shape_cues(annotation: str, entity_counts: Counter[str], operations: list[str]) -> list[str]:
|
||
text = annotation.lower()
|
||
joined_ops = " ".join(operations).lower()
|
||
cues: list[str] = []
|
||
if "triangle" in text or "triangular" in text:
|
||
cues.append("具有三角形轮廓或三角截面")
|
||
if "rectangle" in text or entity_counts.get("skLineSegment", 0) >= 4:
|
||
cues.append("包含矩形/多边形直线轮廓")
|
||
if "circle" in text or entity_counts.get("skCircle", 0):
|
||
cues.append("包含圆形轮廓或圆孔候选特征")
|
||
if "arc" in text or entity_counts.get("skArc", 0):
|
||
cues.append("包含圆弧边界")
|
||
if "slot" in text or "notch" in text or "cutout" in text or "开口" in text:
|
||
cues.append("包含槽、缺口或开口候选结构")
|
||
if "hole" in text or "hole" in joined_ops:
|
||
cues.append("包含孔加工候选结构")
|
||
if "fillet" in joined_ops:
|
||
cues.append("包含圆角过渡")
|
||
if "chamfer" in joined_ops:
|
||
cues.append("包含倒角边")
|
||
if "pattern" in joined_ops:
|
||
cues.append("包含阵列重复特征")
|
||
if "mirror" in joined_ops:
|
||
cues.append("包含镜像对称特征")
|
||
if any(item.startswith("revolve") for item in operations):
|
||
cues.append("绕轴旋转形成回转体")
|
||
if any(item.startswith("extrude") for item in operations):
|
||
cues.append("由二维草图拉伸形成实体")
|
||
return _unique(cues, limit=16)
|
||
|
||
|
||
def _classify_local(annotation: str, operations: list[str], geometric_features: list[str], stl_info: dict[str, Any] | None) -> dict[str, Any]:
|
||
text = " ".join([annotation.lower(), " ".join(operations).lower(), " ".join(geometric_features)])
|
||
slender = _is_slender(stl_info)
|
||
has_extrude = any(item.startswith("extrude") for item in operations)
|
||
has_revolve = any(item.startswith("revolve") for item in operations)
|
||
has_hole = "hole" in text or "孔" in text
|
||
has_triangle = "triangle" in text or "triangular" in text or "三角" in text
|
||
has_rectangle = "rectangle" in text or "矩形" in text
|
||
has_cutout = any(word in text for word in ("slot", "notch", "cutout", "开口", "缺口"))
|
||
|
||
if has_revolve:
|
||
return {
|
||
"category": "回转轴套/法兰类零件（候选）" if has_hole else "回转体机械零件（候选）",
|
||
"confidence": 0.5,
|
||
"candidate_names": ["回转体", "轴套", "法兰盘", "轮毂状零件", "revolved part", "flange", "bushing"],
|
||
"possible_functions": ["可能用于同轴定位、连接、支承、隔套或旋转类结构的几何占位。"],
|
||
"applications": ["机械传动、夹具、管路连接、轴承座周边或需要轴线对称零件的装配场景。"],
|
||
"structural_features": ["回转外形", "轴向轮廓", "圆柱/圆盘候选结构"],
|
||
}
|
||
if has_triangle and has_extrude:
|
||
return {
|
||
"category": "三角棱柱/楔形梁类零件（候选）",
|
||
"confidence": 0.56,
|
||
"candidate_names": ["长条三角棱柱", "楔形梁", "三角截面导轨", "triangular prism", "wedge beam"],
|
||
"possible_functions": ["可能作为楔块、导向条、支撑肋、定位块或三角截面结构件使用。"],
|
||
"applications": ["夹具定位、机械支撑、导向结构、教育/仿真几何库或需要楔形截面的 CAD 检索场景。"],
|
||
"structural_features": ["三角截面", "长条拉伸体", "棱柱体", "斜面侧壁"],
|
||
}
|
||
if has_hole and has_rectangle:
|
||
return {
|
||
"category": "带孔板/安装板类零件（候选）",
|
||
"confidence": 0.52,
|
||
"candidate_names": ["安装板", "连接板", "带孔支架板", "mounting plate", "bracket plate"],
|
||
"possible_functions": ["可能用于螺钉安装、定位连接、固定支撑或作为装配转接板。"],
|
||
"applications": ["设备框架、夹具底板、连接支架、外壳内部固定件。"],
|
||
"structural_features": ["板状主体", "孔特征", "平面安装面"],
|
||
}
|
||
if has_cutout and has_extrude:
|
||
return {
|
||
"category": "开槽板/叉形支架类零件（候选）",
|
||
"confidence": 0.48,
|
||
"candidate_names": ["开槽板", "叉形支架", "U 形支架", "slotted plate", "fork bracket"],
|
||
"possible_functions": ["可能用于避让、卡接、导向、夹持或作为插槽式连接件。"],
|
||
"applications": ["支架、夹具、连接耳、导向槽结构或板件开口检索场景。"],
|
||
"structural_features": ["开口/缺口", "板状拉伸体", "平直侧壁"],
|
||
}
|
||
if has_extrude and slender:
|
||
return {
|
||
"category": "拉伸型梁/导轨类零件（候选）",
|
||
"confidence": 0.46,
|
||
"candidate_names": ["拉伸梁", "导轨条", "长条棱柱", "extruded beam", "rail"],
|
||
"possible_functions": ["可能用于导向、支撑、隔距、边框或长条结构件。"],
|
||
"applications": ["框架、滑轨、夹具、机械结构支撑或型材检索场景。"],
|
||
"structural_features": ["长条外形", "恒定截面候选", "拉伸成型"],
|
||
}
|
||
if has_extrude:
|
||
return {
|
||
"category": "拉伸棱柱/板块类零件（候选）",
|
||
"confidence": 0.42,
|
||
"candidate_names": ["拉伸实体", "板块", "棱柱体", "extruded solid", "prismatic part"],
|
||
"possible_functions": ["可能作为基础块、板件、支撑件或后续加工毛坯。"],
|
||
"applications": ["通用机械零件、夹具、支架、CAD 几何检索和相似形状匹配。"],
|
||
"structural_features": ["二维轮廓拉伸", "平面端面", "直壁结构"],
|
||
}
|
||
return {
|
||
"category": "通用机械 CAD 零件（候选）",
|
||
"confidence": 0.34,
|
||
"candidate_names": ["机械零件", "CAD 几何样本", "mechanical part", "CAD model"],
|
||
"possible_functions": ["可能作为机械结构、连接、支撑或几何检索样本，实际用途需结合装配上下文确认。"],
|
||
"applications": ["CAD 数据集检索、几何相似度匹配、零件分类训练和工程知识库索引。"],
|
||
"structural_features": ["可解析 CAD 特征组合"],
|
||
}
|
||
|
||
|
||
def _keywords(classification: dict[str, Any], operations: list[str], geometric_features: list[str]) -> tuple[list[str], list[str]]:
|
||
text = " ".join([classification["category"], " ".join(classification["candidate_names"]), " ".join(operations), " ".join(geometric_features)]).lower()
|
||
zh = ["CAD模型", "机械零件", "几何检索", "相似特征", *classification["candidate_names"][:4]]
|
||
en = ["cad model", "mechanical part", "geometry retrieval", "similar features"]
|
||
mappings = [
|
||
("三角", "三角棱柱", "triangular prism"),
|
||
("wedge", "楔形", "wedge"),
|
||
("extrude", "拉伸", "extruded"),
|
||
("revolve", "回转", "revolved"),
|
||
("hole", "孔", "hole"),
|
||
("fillet", "圆角", "fillet"),
|
||
("chamfer", "倒角", "chamfer"),
|
||
("slot", "槽", "slot"),
|
||
("flange", "法兰", "flange"),
|
||
("bushing", "轴套", "bushing"),
|
||
("plate", "板件", "plate"),
|
||
("bracket", "支架", "bracket"),
|
||
("rail", "导轨", "rail"),
|
||
]
|
||
for needle, zh_value, en_value in mappings:
|
||
if needle in text or zh_value in text:
|
||
zh.append(zh_value)
|
||
en.append(en_value)
|
||
return _unique(zh, limit=32), _unique(en, limit=32)
|
||
|
||
|
||
def _compose_record(
|
||
sample: Sample,
|
||
input_root: Path,
|
||
shard: str,
|
||
txt_path: Path,
|
||
local_facts: dict[str, Any],
|
||
vision_result: dict[str, Any] | None,
|
||
diagnostics: list[dict[str, Any]],
|
||
mode: str,
|
||
) -> dict[str, Any]:
|
||
vision = _normalize_vision(vision_result)
|
||
status = "described_local"
|
||
if mode != "local":
|
||
status = "described_hybrid" if vision else "local_fallback"
|
||
category = vision.get("category") or local_facts["category"]
|
||
category_confidence = _clamp_float(vision.get("category_confidence"), local_facts["category_confidence"])
|
||
candidate_names = _unique([*vision.get("candidate_names", []), *local_facts["candidate_names"]], limit=12)
|
||
structural_features = _unique([*local_facts["structural_features"], *vision.get("structural_features", [])], limit=32)
|
||
geometric_features = _unique([*local_facts["geometric_features"], *vision.get("geometric_features", [])], limit=40)
|
||
uncertainties = _unique([*local_facts["uncertainties"], *vision.get("uncertainties", [])], limit=20)
|
||
return {
|
||
"schema_version": DESCRIPTION_SCHEMA_VERSION,
|
||
"sample_id": sample.sample_id,
|
||
"shard": shard,
|
||
"status": status,
|
||
"txt_path": str(txt_path),
|
||
"category": category,
|
||
"category_confidence": category_confidence,
|
||
"candidate_names": candidate_names,
|
||
"summary_zh": vision.get("summary_zh") or local_facts["summary_zh"],
|
||
"possible_functions": _unique([*vision.get("possible_functions", []), *local_facts["possible_functions"]], limit=12),
|
||
"applications": _unique([*vision.get("applications", []), *local_facts["applications"]], limit=12),
|
||
"structural_features": structural_features,
|
||
"geometric_features": geometric_features,
|
||
"operations": local_facts["operations"],
|
||
"dimensions": local_facts["dimensions"],
|
||
"keywords_zh": _unique([*vision.get("keywords_zh", []), *local_facts["keywords_zh"]], limit=40),
|
||
"keywords_en": _unique([*vision.get("keywords_en", []), *local_facts["keywords_en"]], limit=40),
|
||
"uncertainties": uncertainties,
|
||
"source_files": _source_files(sample, input_root),
|
||
"diagnostics": diagnostics,
|
||
}
|
||
|
||
|
||
def _normalize_vision(value: dict[str, Any] | None) -> dict[str, Any]:
|
||
if not isinstance(value, dict):
|
||
return {}
|
||
result: dict[str, Any] = {}
|
||
for key in (
|
||
"candidate_names",
|
||
"possible_functions",
|
||
"applications",
|
||
"structural_features",
|
||
"geometric_features",
|
||
"keywords_zh",
|
||
"keywords_en",
|
||
"uncertainties",
|
||
):
|
||
result[key] = _as_text_list(value.get(key))
|
||
for key in ("category", "summary_zh"):
|
||
if isinstance(value.get(key), str) and value[key].strip():
|
||
result[key] = value[key].strip()
|
||
result["category_confidence"] = value.get("category_confidence")
|
||
return result
|
||
|
||
|
||
def _local_summary(classification: dict[str, Any], operations: list[str], geometric_features: list[str], dimensions: list[str]) -> str:
|
||
operation_text = "、".join(operations[:4]) if operations else "可解析 CAD 特征"
|
||
feature_text = "；".join(geometric_features[:4]) if geometric_features else "几何结构待进一步识别"
|
||
dimension_text = "；".join(dimensions[:3]) if dimensions else "未提取到稳定尺寸摘要"
|
||
return (
|
||
f"该模型可保守识别为{classification['category']}，主要由 {operation_text} 构成。"
|
||
f"确定结构包括：{feature_text}。尺寸线索：{dimension_text}。"
|
||
"真实产品身份和装配用途缺少上下文，因此以候选名称和相似几何特征用于检索。"
|
||
)
|
||
|
||
|
||
def _failed_record(sample: Sample, input_root: Path, txt_path: Path, shard: str, diagnostics: list[dict[str, Any]]) -> dict[str, Any]:
|
||
return {
|
||
"schema_version": DESCRIPTION_SCHEMA_VERSION,
|
||
"sample_id": sample.sample_id,
|
||
"shard": shard,
|
||
"status": "failed",
|
||
"txt_path": str(txt_path),
|
||
"category": "通用机械 CAD 零件（描述失败）",
|
||
"category_confidence": 0.0,
|
||
"candidate_names": ["CAD 模型"],
|
||
"summary_zh": "该样本描述生成失败，仅保留源文件索引和诊断信息。",
|
||
"possible_functions": ["无法可靠判断。"],
|
||
"applications": ["需重新运行描述生成或人工复核后再进入向量库。"],
|
||
"structural_features": [],
|
||
"geometric_features": [],
|
||
"operations": [],
|
||
"dimensions": [],
|
||
"keywords_zh": ["CAD模型", "描述失败"],
|
||
"keywords_en": ["cad model", "description failed"],
|
||
"uncertainties": ["描述生成过程失败，不能据此判断模型类别或用途。"],
|
||
"source_files": _source_files(sample, input_root),
|
||
"diagnostics": diagnostics,
|
||
}
|
||
|
||
|
||
def _skipped_record(sample: Sample, input_root: Path, txt_path: Path, shard: str) -> dict[str, Any]:
|
||
return {
|
||
"schema_version": DESCRIPTION_SCHEMA_VERSION,
|
||
"sample_id": sample.sample_id,
|
||
"shard": shard,
|
||
"status": "skipped_existing",
|
||
"txt_path": str(txt_path),
|
||
"category": "",
|
||
"category_confidence": 0.0,
|
||
"candidate_names": [],
|
||
"geometric_features": [],
|
||
"operations": [],
|
||
"dimensions": [],
|
||
"keywords_zh": [],
|
||
"keywords_en": [],
|
||
"uncertainties": ["已有描述文件，未使用 --force，因此本次未覆盖。"],
|
||
"source_files": _source_files(sample, input_root),
|
||
"diagnostics": [{"code": "existing_output_skipped"}],
|
||
}
|
||
|
||
|
||
def _source_files(sample: Sample, input_root: Path) -> dict[str, str]:
|
||
return {key: str(Path(value)) for key, value in sorted(sample.files.items()) if Path(value).is_file() or input_root}
|
||
|
||
|
||
def _read_manifest(path: Path) -> dict[str, dict[str, Any]]:
|
||
if not path.exists():
|
||
return {}
|
||
records: dict[str, dict[str, Any]] = {}
|
||
for line in path.read_text(encoding="utf-8").splitlines():
|
||
if not line.strip():
|
||
continue
|
||
try:
|
||
value = json.loads(line)
|
||
except json.JSONDecodeError:
|
||
continue
|
||
sample_id = str(value.get("sample_id") or "")
|
||
if sample_id:
|
||
records[sample_id] = value
|
||
return records
|
||
|
||
|
||
def _write_manifest(input_root: Path, records: list[dict[str, Any]]) -> int:
|
||
path = input_root / "description_txt" / MANIFEST_NAME
|
||
existing = _read_manifest(path)
|
||
for record in records:
|
||
existing[str(record["sample_id"])] = record
|
||
ordered = [existing[key] for key in sorted(existing)]
|
||
text = "".join(json.dumps(record, ensure_ascii=False, sort_keys=True) + "\n" for record in ordered)
|
||
_atomic_write_text(path, text)
|
||
return len(ordered)
|
||
|
||
|
||
def _summary(input_root: Path, records: list[dict[str, Any]], manifest_rows: int) -> dict[str, Any]:
|
||
counts = Counter(str(record.get("status") or "unknown") for record in records)
|
||
txt_count = sum(1 for record in records if Path(str(record.get("txt_path") or "")).is_file())
|
||
return {
|
||
"sample_count": len(records),
|
||
"statuses": dict(sorted(counts.items())),
|
||
"txt_count": txt_count,
|
||
"manifest": str(input_root / "description_txt" / MANIFEST_NAME),
|
||
"manifest_rows": manifest_rows,
|
||
}
|
||
|
||
|
||
def _atomic_write_text(path: Path, text: str) -> None:
|
||
path.parent.mkdir(parents=True, exist_ok=True)
|
||
tmp = path.with_name(path.name + ".tmp")
|
||
tmp.write_text(text, encoding="utf-8")
|
||
tmp.replace(path)
|
||
|
||
|
||
def _safe_number(value: Any) -> float | None:
|
||
try:
|
||
result = _number(value, True)
|
||
except Exception:
|
||
return None
|
||
return result if math.isfinite(result) else None
|
||
|
||
|
||
def _safe_point(value: Any) -> list[float] | None:
|
||
try:
|
||
point = _point(value)
|
||
except Exception:
|
||
return None
|
||
return point if len(point) >= 2 and all(math.isfinite(item) for item in point[:2]) else None
|
||
|
||
|
||
def _sketch_bbox(sketch: Any) -> list[float] | None:
|
||
xs: list[float] = []
|
||
ys: list[float] = []
|
||
for entity in sketch.entities:
|
||
params = entity.params
|
||
if entity.operation == "skLineSegment":
|
||
for key in ("start", "end"):
|
||
point = _safe_point(params.get(key))
|
||
if point:
|
||
xs.append(point[0])
|
||
ys.append(point[1])
|
||
elif entity.operation == "skCircle":
|
||
center = _safe_point(params.get("center"))
|
||
radius = _safe_number(params.get("radius"))
|
||
if center and radius is not None:
|
||
xs.extend([center[0] - radius, center[0] + radius])
|
||
ys.extend([center[1] - radius, center[1] + radius])
|
||
elif entity.operation == "skArc":
|
||
for key in ("start", "mid", "end"):
|
||
point = _safe_point(params.get(key))
|
||
if point:
|
||
xs.append(point[0])
|
||
ys.append(point[1])
|
||
elif entity.operation == "skPoint":
|
||
point = _safe_point(params.get("position"))
|
||
if point:
|
||
xs.append(point[0])
|
||
ys.append(point[1])
|
||
if not xs or not ys:
|
||
return None
|
||
return [min(xs), min(ys), max(xs), max(ys)]
|
||
|
||
|
||
def _stl_bbox(path: Path) -> dict[str, Any] | None:
|
||
try:
|
||
data = path.read_bytes()
|
||
except OSError:
|
||
return None
|
||
vertices = _binary_stl_vertices(data) or _ascii_stl_vertices(data)
|
||
if not vertices:
|
||
return None
|
||
mins = [min(vertex[i] for vertex in vertices) for i in range(3)]
|
||
maxs = [max(vertex[i] for vertex in vertices) for i in range(3)]
|
||
size = [maxs[i] - mins[i] for i in range(3)]
|
||
return {"min_mm": mins, "max_mm": maxs, "size_mm": size, "vertex_count": len(vertices)}
|
||
|
||
|
||
def _binary_stl_vertices(data: bytes) -> list[tuple[float, float, float]]:
|
||
if len(data) < 84:
|
||
return []
|
||
triangle_count = struct.unpack("<I", data[80:84])[0]
|
||
expected = 84 + triangle_count * 50
|
||
if triangle_count <= 0 or expected > len(data):
|
||
return []
|
||
vertices: list[tuple[float, float, float]] = []
|
||
offset = 84
|
||
for _ in range(triangle_count):
|
||
offset += 12
|
||
for _ in range(3):
|
||
vertices.append(struct.unpack("<fff", data[offset:offset + 12]))
|
||
offset += 12
|
||
offset += 2
|
||
return vertices
|
||
|
||
|
||
def _ascii_stl_vertices(data: bytes) -> list[tuple[float, float, float]]:
|
||
try:
|
||
text = data[:5_000_000].decode("utf-8", errors="ignore")
|
||
except Exception:
|
||
return []
|
||
pattern = re.compile(r"\bvertex\s+([-+0-9.eE]+)\s+([-+0-9.eE]+)\s+([-+0-9.eE]+)")
|
||
vertices = []
|
||
for match in pattern.finditer(text):
|
||
try:
|
||
vertices.append((float(match.group(1)), float(match.group(2)), float(match.group(3))))
|
||
except ValueError:
|
||
continue
|
||
return vertices
|
||
|
||
|
||
def _image_metadata(path: Path) -> dict[str, Any]:
|
||
try:
|
||
from PIL import Image
|
||
|
||
with Image.open(path) as image:
|
||
return {"width": image.width, "height": image.height, "format": str(image.format or "").lower()}
|
||
except Exception as exc:
|
||
return {"error": f"image metadata unavailable: {type(exc).__name__}"}
|
||
|
||
|
||
def _stl_features(stl_info: dict[str, Any] | None) -> list[str]:
|
||
if not stl_info:
|
||
return []
|
||
size = stl_info.get("size_mm") or []
|
||
if len(size) != 3:
|
||
return []
|
||
features = ["STL 网格提供三维包围盒"]
|
||
if _is_slender(stl_info):
|
||
features.append("整体呈长条比例")
|
||
if _is_plate_like(stl_info):
|
||
features.append("整体呈薄板比例")
|
||
return features
|
||
|
||
|
||
def _stl_dimensions(stl_info: dict[str, Any] | None) -> list[str]:
|
||
if not stl_info:
|
||
return []
|
||
size = stl_info.get("size_mm") or []
|
||
if len(size) != 3:
|
||
return []
|
||
return [f"STL 三维包围盒约 {_format_float(size[0])} × {_format_float(size[1])} × {_format_float(size[2])} mm"]
|
||
|
||
|
||
def _is_slender(stl_info: dict[str, Any] | None) -> bool:
|
||
if not stl_info:
|
||
return False
|
||
values = [abs(float(item)) for item in stl_info.get("size_mm") or [] if abs(float(item)) > 1e-9]
|
||
return len(values) >= 2 and max(values) / max(min(values), 1e-9) >= 3.0
|
||
|
||
|
||
def _is_plate_like(stl_info: dict[str, Any] | None) -> bool:
|
||
if not stl_info:
|
||
return False
|
||
values = sorted(abs(float(item)) for item in stl_info.get("size_mm") or [] if abs(float(item)) > 1e-9)
|
||
return len(values) == 3 and values[0] * 4 <= values[1]
|
||
|
||
|
||
def _plane_name(value: Any) -> str:
|
||
text = json.dumps(plain(value), ensure_ascii=False)
|
||
for name in ("Top", "Front", "Right"):
|
||
if f"{name}.planeOp" in text:
|
||
return name
|
||
return "未知"
|
||
|
||
|
||
def _enum_tail(value: Any) -> str:
|
||
return str(value or "").split(".")[-1]
|
||
|
||
|
||
def _truthy(value: Any) -> bool:
|
||
return value is True or (isinstance(value, str) and value.lower() == "true")
|
||
|
||
|
||
def _format_float(value: float) -> str:
|
||
rounded = round(float(value), 4)
|
||
if rounded == int(rounded):
|
||
return str(int(rounded))
|
||
return f"{rounded:.4f}".rstrip("0").rstrip(".")
|
||
|
||
|
||
def _format_mm(value: float) -> str:
|
||
return f"{_format_float(value)} mm"
|
||
|
||
|
||
def _clamp_float(value: Any, fallback: float) -> float:
|
||
try:
|
||
number = float(value)
|
||
except (TypeError, ValueError):
|
||
number = float(fallback)
|
||
if not math.isfinite(number):
|
||
number = float(fallback)
|
||
return max(0.0, min(1.0, number))
|
||
|
||
|
||
def _as_text_list(value: Any) -> list[str]:
|
||
if not isinstance(value, list):
|
||
return []
|
||
return [str(item).strip() for item in value if str(item).strip()]
|
||
|
||
|
||
def _unique(values: list[Any], *, limit: int | None = None) -> list[str]:
|
||
result: list[str] = []
|
||
seen: set[str] = set()
|
||
for value in values:
|
||
text = str(value).strip()
|
||
if not text or text in seen:
|
||
continue
|
||
seen.add(text)
|
||
result.append(text)
|
||
if limit is not None and len(result) >= limit:
|
||
break
|
||
return result
|
||
|
||
|
||
def _join_or_unknown(values: Any) -> str:
|
||
items = _as_text_list(values) if isinstance(values, list) else _unique(list(values or [])) if isinstance(values, tuple) else []
|
||
return "、".join(items) if items else "无法可靠判断"
|