Ukawka
131 lines
7.5 KiB
Python
131 lines
7.5 KiB
Python
"""Measured mask QA and previews made only from the original character pixels."""
|
|
import colorsys
|
|
import html
|
|
import json
|
|
import math
|
|
import re
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import cv2
|
|
from PIL import Image, ImageDraw, ImageFont
|
|
|
|
|
|
def assess_candidate(mask, alpha, bbox, points):
|
|
mask = np.asarray(mask, dtype=bool)
|
|
if mask.shape != alpha.shape:
|
|
raise ValueError('Mask dimensions do not match the source')
|
|
clipped = mask & (alpha > 0)
|
|
_, _, stats, _ = cv2.connectedComponentsWithStats(clipped.astype(np.uint8), 8)
|
|
area, foreground_area = int(mask.sum()), int(clipped.sum())
|
|
x0, y0, x1, y1 = bbox
|
|
hit = sum(bool(clipped[y, x]) for x, y in points)
|
|
return {
|
|
'raw_area_px': area, 'area_px': foreground_area,
|
|
'outside_alpha_ratio': (area - foreground_area) / max(1, area),
|
|
'inside_roi_fraction': int(clipped[y0:y1, x0:x1].sum()) / max(1, foreground_area),
|
|
'seed_hit_fraction': hit / max(1, len(points)),
|
|
'components_at_least_4px': int((stats[1:, cv2.CC_STAT_AREA] >= 4).sum()),
|
|
}
|
|
|
|
|
|
def render_outputs(image, parts, out_dir):
|
|
out = Path(out_dir)
|
|
(out / 'parts').mkdir(parents=True, exist_ok=True)
|
|
(out / 'masks').mkdir(parents=True, exist_ok=True)
|
|
rgba = np.array(image.convert('RGBA'))
|
|
h, w = rgba.shape[:2]
|
|
coverage = np.zeros((h, w), np.uint16)
|
|
white = Image.new('RGBA', (w, h), 'white')
|
|
overlay = Image.alpha_composite(white, image.convert('RGBA'))
|
|
records, crops, clipped_masks = [], [], []
|
|
seen = set()
|
|
for index, part in enumerate(parts):
|
|
name = part['id']
|
|
if not re.fullmatch('[a-z][a-z0-9_-]{0,63}', name) or name in seen:
|
|
raise ValueError(f'Invalid or duplicate part ID: {name}')
|
|
seen.add(name)
|
|
mask = np.asarray(part['mask'], dtype=bool)
|
|
if mask.shape != (h, w):
|
|
raise ValueError(f'{name}: mask dimensions mismatch')
|
|
mask = mask & (rgba[:, :, 3] > 0)
|
|
clipped_masks.append(mask)
|
|
ys, xs = np.nonzero(mask)
|
|
if not len(xs):
|
|
raise ValueError(f'{name}: empty foreground mask')
|
|
coverage += mask
|
|
bbox = [int(xs.min()), int(ys.min()), int(xs.max()) + 1, int(ys.max()) + 1]
|
|
color = [round(v * 255) for v in colorsys.hsv_to_rgb((index * .61803398875) % 1, .75, .95)]
|
|
cut = rgba.copy()
|
|
cut[~mask] = 0
|
|
crop = Image.fromarray(cut).crop(bbox)
|
|
crop.save(out / 'parts' / f'{name}.png')
|
|
Image.fromarray(mask.astype(np.uint8) * 255).save(out / 'masks' / f'{name}.png')
|
|
tint = np.zeros_like(rgba)
|
|
tint[mask, :3] = color
|
|
tint[mask, 3] = 100
|
|
overlay = Image.alpha_composite(overlay, Image.fromarray(tint))
|
|
records.append({k: v for k, v in part.items() if k != 'mask'} | {
|
|
'color_rgb': color, 'crop_xyxy': bbox, 'area_px': int(mask.sum()),
|
|
'image': f'parts/{name}.png', 'mask': f'masks/{name}.png',
|
|
})
|
|
crops.append(crop)
|
|
near_duplicates = []
|
|
for i, a in enumerate(clipped_masks):
|
|
for j in range(i + 1, len(clipped_masks)):
|
|
b = clipped_masks[j]
|
|
iou = int((a & b).sum()) / max(1, int((a | b).sum()))
|
|
if iou >= .8:
|
|
near_duplicates.append({'part_a': records[i]['id'], 'part_b': records[j]['id'], 'iou': iou})
|
|
for index in [i, j]:
|
|
warnings = records[index].setdefault('warnings', []).copy()
|
|
if 'NEAR_DUPLICATE_MASK' not in warnings:
|
|
warnings.append('NEAR_DUPLICATE_MASK')
|
|
records[index]['warnings'] = warnings
|
|
# Keep native scale, separated by padding and a label band. No rotation or
|
|
# resampling of source pixels is used in the exploded sheet.
|
|
pad, label = 24, 26
|
|
cell_widths = [max(c.width, len(r['id']) * 8 + 8) for c, r in zip(crops, records)]
|
|
sheet_width = max(640, max(cell_widths, default=1) + pad * 2,
|
|
int(math.sqrt(sum((cw + pad) * (c.height + label + pad) for cw, c in zip(cell_widths, crops)) * 1.5)))
|
|
x, y, row_height = pad, pad, 0
|
|
for record, crop, cw in zip(records, crops, cell_widths):
|
|
if x + cw + pad > sheet_width and x != pad:
|
|
x, y, row_height = pad, y + row_height + pad, 0
|
|
record['sheet_xyxy'] = [x, y + label, x + crop.width, y + label + crop.height]
|
|
x += cw + pad
|
|
row_height = max(row_height, label + crop.height)
|
|
sheet = Image.new('RGB', (sheet_width, max(100, y + row_height + pad)), 'white')
|
|
draw = ImageDraw.Draw(sheet)
|
|
font = ImageFont.load_default(size=14)
|
|
for record, crop in zip(records, crops):
|
|
x0, y0, _, _ = record['sheet_xyxy']
|
|
draw.rectangle((x0, y0 - label + 4, x0 + 10, y0 - label + 14), fill=tuple(record['color_rgb']))
|
|
draw.text((x0 + 16, y0 - label + 1), record['id'], fill='black', font=font)
|
|
sheet.paste(crop, (x0, y0), crop)
|
|
overlay.convert('RGB').save(out / 'annotated.png')
|
|
sheet.save(out / 'parts-sheet.png')
|
|
Image.fromarray(((rgba[:, :, 3] > 0) & (coverage == 0)).astype(np.uint8) * 255).save(out / 'unassigned-mask.png')
|
|
Image.fromarray((coverage > 1).astype(np.uint8) * 255).save(out / 'overlap-mask.png')
|
|
total = int((rgba[:, :, 3] > 0).sum())
|
|
return {'parts': records, 'near_duplicate_masks': near_duplicates, 'coverage': {
|
|
'foreground_pixels': total, 'covered_pixels': int((coverage > 0).sum()),
|
|
'coverage_ratio': int((coverage > 0).sum()) / max(1, total),
|
|
'overlap_pixels': int((coverage > 1).sum()),
|
|
'unassigned_pixels': int(((rgba[:, :, 3] > 0) & (coverage == 0)).sum()),
|
|
}, 'outputs': {'annotated': 'annotated.png', 'parts_sheet': 'parts-sheet.png'}}
|
|
|
|
|
|
def write_report(out_dir, result):
|
|
out = Path(out_dir)
|
|
rows = ''.join(f'<tr><td><span style="background:rgb({",".join(map(str, p["color_rgb"]))})"> </span> <a href="{p["image"]}">{html.escape(p["id"])}</a></td><td>{html.escape(p["semantic"])}</td><td>{p["model_score"]:.3f}</td><td>{html.escape(", ".join(p.get("warnings", [])))}</td></tr>' for p in result['parts'])
|
|
body = f'''<!doctype html><html lang="zh-CN"><meta charset="utf-8"><title>SAM 部件分割</title>
|
|
<style>body{{font:16px system-ui;max-width:1400px;margin:24px auto;padding:16px}}img{{max-width:100%;height:auto}}td,th{{border:1px solid #ccc;padding:8px}}table{{border-collapse:collapse}}pre{{white-space:pre-wrap}}</style>
|
|
<h1>SAM 本地分割结果</h1><p>颜色表示候选部件 mask。白底拆解图保留原始可见像素，尚未补全遮挡部分。重叠 mask 保留并计入诊断，不代表已经获得互斥、语义正确的最终拆件。</p>
|
|
<h2>原图半透明标注</h2><img src="annotated.png"><h2>白底拆解图</h2><img src="parts-sheet.png">
|
|
<h2>部件</h2><table><tr><th>ID / 颜色</th><th>语义</th><th>模型分数</th><th>需复核</th></tr>{rows}</table>
|
|
<h2>未生成的部件</h2><pre>{html.escape(json.dumps(result['skipped'], ensure_ascii=False, indent=2))}</pre>
|
|
<h2>覆盖与重叠</h2><p>覆盖率只表示有 mask 的像素比例，不是分割准确率。模型分数也不代表部件语义已验证。</p><p><a href="unassigned-mask.png">漏分区域（白色）</a> · <a href="overlap-mask.png">重叠区域（白色）</a></p><pre>{html.escape(json.dumps(result['coverage'], indent=2))}</pre>
|
|
<h2>疑似重复分割</h2><p>以下不同部件的 mask IoU ≥ 0.8，需要检查是否分到了同一物体。</p><pre>{html.escape(json.dumps(result['near_duplicate_masks'], indent=2))}</pre></html>'''
|
|
(out / 'report.html').write_text(body, encoding='utf-8')
|