419 lines
15 KiB
Python
419 lines
15 KiB
Python
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import csv
|
|
import json
|
|
import math
|
|
from dataclasses import asdict, dataclass
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
from PIL import Image, ImageDraw
|
|
from skimage import color, feature, measure, morphology, transform
|
|
|
|
|
|
IMAGE_EXTS = {".png", ".jpg", ".jpeg", ".webp"}
|
|
|
|
|
|
@dataclass
|
|
class ImageFeatures:
|
|
file: str
|
|
stem: str
|
|
width: int
|
|
height: int
|
|
foreground_fraction: float
|
|
component_count: int
|
|
largest_component_fraction: float
|
|
hole_count: int
|
|
weighted_compactness: float
|
|
edge_density: float
|
|
color_bins: int
|
|
circular_signal: float
|
|
rectilinear_signal: float
|
|
thin_profile_signal: float
|
|
symmetry_signal: float
|
|
gear_signal: float
|
|
standard_signal: float
|
|
complexity_penalty: float
|
|
restorability_score: float
|
|
recommended_backend: str
|
|
confidence: str
|
|
reasons: str
|
|
|
|
|
|
def clamp(value: float, lo: float = 0.0, hi: float = 1.0) -> float:
|
|
return max(lo, min(hi, value))
|
|
|
|
|
|
def percentile_score(value: float, p10: float, p90: float, invert: bool = False) -> float:
|
|
if p90 <= p10:
|
|
score = 0.5
|
|
else:
|
|
score = clamp((value - p10) / (p90 - p10))
|
|
return 1.0 - score if invert else score
|
|
|
|
|
|
def list_images(image_dir: Path) -> list[Path]:
|
|
return sorted(
|
|
p for p in image_dir.iterdir()
|
|
if p.is_file() and p.suffix.lower() in IMAGE_EXTS
|
|
)
|
|
|
|
|
|
def load_small_image(path: Path, max_width: int) -> tuple[np.ndarray, tuple[int, int]]:
|
|
image = Image.open(path).convert("RGB")
|
|
original_size = image.size
|
|
if image.width > max_width:
|
|
height = round(image.height * max_width / image.width)
|
|
image = image.resize((max_width, height), Image.Resampling.LANCZOS)
|
|
return np.asarray(image, dtype=np.uint8), original_size
|
|
|
|
|
|
def foreground_mask(rgb: np.ndarray) -> np.ndarray:
|
|
arr = rgb.astype(np.float32) / 255.0
|
|
hsv = color.rgb2hsv(arr)
|
|
value = hsv[..., 2]
|
|
saturation = hsv[..., 1]
|
|
|
|
# SOLIDWORKS screenshots have pale gradient backgrounds and many light-gray
|
|
# parts, so include medium-light gray entity pixels as well as colored parts.
|
|
mask = (value < 0.88) | ((saturation > 0.09) & (value < 0.99))
|
|
mask = morphology.closing(mask, morphology.disk(2))
|
|
mask = morphology.remove_small_objects(mask, max_size=max(32, mask.size // 9000))
|
|
mask = morphology.remove_small_holes(mask, max_size=max(32, mask.size // 14000))
|
|
return mask
|
|
|
|
|
|
def crop_component(mask: np.ndarray, label_image: np.ndarray, label_id: int) -> np.ndarray:
|
|
ys, xs = np.where(label_image == label_id)
|
|
if len(xs) == 0:
|
|
return np.zeros((1, 1), dtype=bool)
|
|
pad = 4
|
|
y0 = max(int(ys.min()) - pad, 0)
|
|
y1 = min(int(ys.max()) + pad + 1, mask.shape[0])
|
|
x0 = max(int(xs.min()) - pad, 0)
|
|
x1 = min(int(xs.max()) + pad + 1, mask.shape[1])
|
|
return label_image[y0:y1, x0:x1] == label_id
|
|
|
|
|
|
def symmetry_for(component: np.ndarray) -> float:
|
|
if component.size <= 1:
|
|
return 0.0
|
|
target = transform.resize(
|
|
component.astype(float),
|
|
(64, 64),
|
|
order=0,
|
|
preserve_range=True,
|
|
anti_aliasing=False,
|
|
) > 0.5
|
|
area = target.sum()
|
|
if area == 0:
|
|
return 0.0
|
|
horiz = np.logical_and(target, np.fliplr(target)).sum() / np.logical_or(target, np.fliplr(target)).sum()
|
|
vert = np.logical_and(target, np.flipud(target)).sum() / np.logical_or(target, np.flipud(target)).sum()
|
|
return float(max(horiz, vert))
|
|
|
|
|
|
def raw_features(path: Path, max_width: int) -> dict[str, float | int | str]:
|
|
rgb, (width, height) = load_small_image(path, max_width=max_width)
|
|
gray = color.rgb2gray(rgb)
|
|
mask = foreground_mask(rgb)
|
|
label_image = measure.label(mask)
|
|
props = [
|
|
prop for prop in measure.regionprops(label_image)
|
|
if prop.area >= max(48, mask.size // 4500)
|
|
]
|
|
props.sort(key=lambda item: item.area, reverse=True)
|
|
|
|
foreground_fraction = float(mask.mean())
|
|
component_count = len(props)
|
|
largest = props[0] if props else None
|
|
largest_component_fraction = float((largest.area / mask.size) if largest else 0.0)
|
|
total_area = sum(prop.area for prop in props) or 1
|
|
|
|
compactness_values = []
|
|
holes = 0
|
|
circular_signal = 0.0
|
|
rectilinear_signal = 0.0
|
|
thin_profile_signal = 0.0
|
|
gear_signal = 0.0
|
|
symmetry_signal = 0.0
|
|
|
|
for index, prop in enumerate(props[:16]):
|
|
area = max(float(prop.area), 1.0)
|
|
perimeter = max(float(prop.perimeter), 1.0)
|
|
compactness = perimeter * perimeter / (4.0 * math.pi * area)
|
|
compactness_values.append(compactness * (area / total_area))
|
|
component_holes = max(0, 1 - int(prop.euler_number))
|
|
holes += component_holes
|
|
|
|
minr, minc, maxr, maxc = prop.bbox
|
|
box_h = maxr - minr
|
|
box_w = maxc - minc
|
|
box_area = max(float(box_h * box_w), 1.0)
|
|
aspect = max(box_w / max(box_h, 1), box_h / max(box_w, 1))
|
|
extent = prop.area / box_area
|
|
solidity = float(getattr(prop, "solidity", 0.0))
|
|
eccentricity = float(getattr(prop, "eccentricity", 1.0))
|
|
|
|
circular = clamp((1.0 - eccentricity) * 1.7) * clamp(solidity * 1.25) * clamp(2.2 - compactness / 2.1)
|
|
rect = clamp((extent - 0.35) / 0.45) * clamp(solidity * 1.15)
|
|
thin = clamp((aspect - 3.0) / 8.0) * clamp(solidity * 1.2)
|
|
|
|
weight = area / total_area
|
|
circular_signal = max(circular_signal, circular)
|
|
rectilinear_signal = max(rectilinear_signal, rect)
|
|
thin_profile_signal = max(thin_profile_signal, thin)
|
|
|
|
if index < 4:
|
|
comp = crop_component(mask, label_image, prop.label)
|
|
symmetry_signal = max(symmetry_signal, symmetry_for(comp) * clamp(weight * 4.0))
|
|
|
|
teeth_like = clamp((compactness - 2.2) / 7.5)
|
|
gear_signal = max(
|
|
gear_signal,
|
|
clamp((1.0 - eccentricity) * 2.0) * teeth_like * clamp((component_holes + 1) / 3.0),
|
|
)
|
|
|
|
weighted_compactness = float(sum(compactness_values))
|
|
edges = feature.canny(gray, sigma=1.1)
|
|
edge_density = float(edges.mean())
|
|
|
|
object_pixels = rgb[mask]
|
|
if len(object_pixels) == 0:
|
|
color_bins = 0
|
|
else:
|
|
quantized = (object_pixels // 32).astype(np.uint8)
|
|
color_bins = int(len(np.unique(quantized.reshape(-1, 3), axis=0)))
|
|
|
|
return {
|
|
"file": str(path),
|
|
"stem": path.stem,
|
|
"width": width,
|
|
"height": height,
|
|
"foreground_fraction": foreground_fraction,
|
|
"component_count": component_count,
|
|
"largest_component_fraction": largest_component_fraction,
|
|
"hole_count": holes,
|
|
"weighted_compactness": weighted_compactness,
|
|
"edge_density": edge_density,
|
|
"color_bins": color_bins,
|
|
"circular_signal": float(circular_signal),
|
|
"rectilinear_signal": float(rectilinear_signal),
|
|
"thin_profile_signal": float(thin_profile_signal),
|
|
"symmetry_signal": float(symmetry_signal),
|
|
"gear_signal": float(gear_signal),
|
|
}
|
|
|
|
|
|
def score_rows(rows: list[dict[str, float | int | str]]) -> list[ImageFeatures]:
|
|
compactness = np.asarray([float(row["weighted_compactness"]) for row in rows])
|
|
edge_density = np.asarray([float(row["edge_density"]) for row in rows])
|
|
components = np.asarray([float(row["component_count"]) for row in rows])
|
|
color_bins = np.asarray([float(row["color_bins"]) for row in rows])
|
|
holes = np.asarray([float(row["hole_count"]) for row in rows])
|
|
|
|
p = {
|
|
"compactness": np.percentile(compactness, [10, 90]),
|
|
"edge": np.percentile(edge_density, [10, 90]),
|
|
"components": np.percentile(components, [10, 90]),
|
|
"colors": np.percentile(color_bins, [10, 90]),
|
|
"holes": np.percentile(holes, [50, 95]),
|
|
}
|
|
|
|
scored: list[ImageFeatures] = []
|
|
for row in rows:
|
|
comp_score = percentile_score(float(row["weighted_compactness"]), *p["compactness"], invert=True)
|
|
edge_score = percentile_score(float(row["edge_density"]), *p["edge"], invert=True)
|
|
component_score = percentile_score(float(row["component_count"]), *p["components"], invert=True)
|
|
color_score = percentile_score(float(row["color_bins"]), *p["colors"], invert=True)
|
|
many_hole_penalty = percentile_score(float(row["hole_count"]), *p["holes"])
|
|
|
|
circular = float(row["circular_signal"])
|
|
rect = float(row["rectilinear_signal"])
|
|
thin = float(row["thin_profile_signal"])
|
|
symmetry = float(row["symmetry_signal"])
|
|
gear = float(row["gear_signal"])
|
|
holes_good = 1.0 if 1 <= int(row["hole_count"]) <= 12 else 0.35 if int(row["hole_count"]) == 0 else 0.0
|
|
standard = max(circular, rect, thin, gear)
|
|
|
|
complexity_penalty = (
|
|
0.34 * (1.0 - comp_score)
|
|
+ 0.28 * (1.0 - edge_score)
|
|
+ 0.25 * (1.0 - component_score)
|
|
+ 0.13 * (1.0 - color_score)
|
|
)
|
|
|
|
special_family = (
|
|
gear >= 0.80
|
|
and circular >= 0.40
|
|
and float(row["edge_density"]) >= 0.048
|
|
)
|
|
|
|
score = (
|
|
25.0
|
|
+ 16.0 * comp_score
|
|
+ 12.0 * edge_score
|
|
+ 8.0 * component_score
|
|
+ 4.0 * color_score
|
|
+ 12.0 * standard
|
|
+ 5.0 * symmetry
|
|
+ 4.0 * holes_good
|
|
+ (7.0 * gear if special_family else 0.0)
|
|
- 10.0 * many_hole_penalty
|
|
)
|
|
score = max(0.0, min(100.0, score))
|
|
|
|
if special_family:
|
|
backend = "simplecadapi"
|
|
elif score >= 72.0:
|
|
backend = "build123d"
|
|
elif score >= 58.0:
|
|
backend = "build123d_or_manual_review"
|
|
else:
|
|
backend = "manual_review_before_modeling"
|
|
|
|
if score >= 82:
|
|
confidence = "high"
|
|
elif score >= 68:
|
|
confidence = "medium"
|
|
else:
|
|
confidence = "low"
|
|
|
|
reasons = []
|
|
if special_family:
|
|
reasons.append("gear/ring/rack-like standard mechanical family")
|
|
if circular >= 0.45:
|
|
reasons.append("circular or coaxial feature signal")
|
|
if rect >= 0.55:
|
|
reasons.append("simple plate/block extrusion signal")
|
|
if thin >= 0.45:
|
|
reasons.append("shaft/rail/thin extrusion signal")
|
|
if symmetry >= 0.55:
|
|
reasons.append("strong silhouette symmetry")
|
|
if holes_good >= 1.0:
|
|
reasons.append("moderate visible hole count")
|
|
if complexity_penalty >= 0.55:
|
|
reasons.append("penalized for dense edges/components/colors")
|
|
if not reasons:
|
|
reasons.append("generic low-complexity B-Rep signal")
|
|
|
|
scored.append(
|
|
ImageFeatures(
|
|
restorability_score=round(score, 3),
|
|
recommended_backend=backend,
|
|
confidence=confidence,
|
|
standard_signal=round(standard, 4),
|
|
complexity_penalty=round(complexity_penalty, 4),
|
|
reasons="; ".join(reasons),
|
|
**{
|
|
key: row[key] for key in (
|
|
"file",
|
|
"stem",
|
|
"width",
|
|
"height",
|
|
"foreground_fraction",
|
|
"component_count",
|
|
"largest_component_fraction",
|
|
"hole_count",
|
|
"weighted_compactness",
|
|
"edge_density",
|
|
"color_bins",
|
|
"circular_signal",
|
|
"rectilinear_signal",
|
|
"thin_profile_signal",
|
|
"symmetry_signal",
|
|
"gear_signal",
|
|
)
|
|
},
|
|
)
|
|
)
|
|
return sorted(scored, key=lambda item: item.restorability_score, reverse=True)
|
|
|
|
|
|
def write_csv(rows: list[ImageFeatures], path: Path) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
fields = list(asdict(rows[0]).keys()) if rows else []
|
|
with path.open("w", newline="", encoding="utf-8") as handle:
|
|
writer = csv.DictWriter(handle, fieldnames=fields)
|
|
writer.writeheader()
|
|
for row in rows:
|
|
writer.writerow(asdict(row))
|
|
|
|
|
|
def write_json(rows: list[ImageFeatures], path: Path) -> None:
|
|
path.write_text(
|
|
json.dumps([asdict(row) for row in rows], ensure_ascii=False, indent=2),
|
|
encoding="utf-8",
|
|
)
|
|
|
|
|
|
def make_contact_sheet(rows: list[ImageFeatures], output: Path, limit: int, title: str) -> None:
|
|
selected = rows[:limit]
|
|
if not selected:
|
|
return
|
|
thumb_w, thumb_h = 220, 316
|
|
label_h = 42
|
|
cols = 8
|
|
rows_count = math.ceil(len(selected) / cols)
|
|
sheet = Image.new("RGB", (cols * thumb_w, rows_count * (thumb_h + label_h) + 28), "white")
|
|
draw = ImageDraw.Draw(sheet)
|
|
draw.text((8, 8), title, fill=(0, 0, 0))
|
|
y_offset = 28
|
|
for index, item in enumerate(selected):
|
|
path = Path(item.file)
|
|
image = Image.open(path).convert("RGB")
|
|
image.thumbnail((thumb_w, thumb_h), Image.Resampling.LANCZOS)
|
|
col = index % cols
|
|
row = index // cols
|
|
x = col * thumb_w + (thumb_w - image.width) // 2
|
|
y = y_offset + row * (thumb_h + label_h)
|
|
sheet.paste(image, (x, y))
|
|
label = f"{index + 1:02d} {item.stem} {item.restorability_score:.1f}"
|
|
draw.text((col * thumb_w + 4, y + thumb_h + 3), label, fill=(0, 0, 0))
|
|
draw.text((col * thumb_w + 4, y + thumb_h + 20), item.recommended_backend[:28], fill=(70, 70, 70))
|
|
output.parent.mkdir(parents=True, exist_ok=True)
|
|
sheet.save(output, quality=90)
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description="Rank SOLIDWORKS screenshot restorability for cad-router backends.")
|
|
parser.add_argument("image_dir", type=Path)
|
|
parser.add_argument("--out-dir", type=Path, default=Path("models/sldprt_screen_rank"))
|
|
parser.add_argument("--max-width", type=int, default=420)
|
|
parser.add_argument("--contact-limit", type=int, default=80)
|
|
parser.add_argument("--skip-contact", action="store_true", help="Only write CSV/JSON rankings.")
|
|
args = parser.parse_args()
|
|
|
|
paths = list_images(args.image_dir)
|
|
if not paths:
|
|
raise SystemExit(f"No images found in {args.image_dir}")
|
|
|
|
raw = []
|
|
for index, path in enumerate(paths, start=1):
|
|
raw.append(raw_features(path, max_width=args.max_width))
|
|
if index % 100 == 0:
|
|
print(f"processed={index}/{len(paths)}", flush=True)
|
|
ranked = score_rows(raw)
|
|
|
|
args.out_dir.mkdir(parents=True, exist_ok=True)
|
|
write_csv(ranked, args.out_dir / "ranking.csv")
|
|
write_json(ranked, args.out_dir / "ranking.json")
|
|
if not args.skip_contact:
|
|
make_contact_sheet(ranked, args.out_dir / "top_80_contact_sheet.jpg", args.contact_limit, "Top restorability candidates")
|
|
make_contact_sheet(list(reversed(ranked)), args.out_dir / "bottom_40_contact_sheet.jpg", 40, "Lowest restorability candidates")
|
|
|
|
print(f"images={len(ranked)}")
|
|
print(f"csv={args.out_dir / 'ranking.csv'}")
|
|
print(f"json={args.out_dir / 'ranking.json'}")
|
|
print(f"top_sheet={args.out_dir / 'top_80_contact_sheet.jpg'}")
|
|
print("top10=")
|
|
for item in ranked[:10]:
|
|
print(f"{item.stem},{item.restorability_score:.1f},{item.recommended_backend},{item.reasons}")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|