File size: 3,215 Bytes
ce847d4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
"""Debug detector output to understand word segmentation."""
import numpy as np
import onnxruntime as ort
from PIL import Image
from pathlib import Path
models_dir = Path("oneocr_extracted/onnx_models")
img = Image.open("image.png").convert("RGB")
w, h = img.size
# Detector setup
sess = ort.InferenceSession(str(next(models_dir.glob("model_00_*"))),
providers=["CPUExecutionProvider"])
scale = 800 / max(h, w)
dh = (int(h * scale) + 31) // 32 * 32
dw = (int(w * scale) + 31) // 32 * 32
img_d = np.array(img.resize((dw, dh), Image.LANCZOS), dtype=np.float32)
img_d = img_d[:, :, ::-1] - np.array([102.9801, 115.9465, 122.7717], dtype=np.float32)
data = img_d.transpose(2, 0, 1)[np.newaxis].astype(np.float32)
im_info = np.array([[dh, dw, scale]], dtype=np.float32)
outputs = sess.run(None, {"data": data, "im_info": im_info})
output_names = [o.name for o in sess.get_outputs()]
out_dict = dict(zip(output_names, outputs))
# Analyze FPN2 (highest resolution)
pixel_scores = out_dict["scores_hori_fpn2"][0, 0] # [56, 200]
link_scores = out_dict["link_scores_hori_fpn2"][0] # [8, 56, 200]
print(f"FPN2 shape: {pixel_scores.shape}")
print(f"Pixel scores: min={pixel_scores.min():.4f} max={pixel_scores.max():.4f}")
# Find text region
text_mask = pixel_scores > 0.6
print(f"Text pixels (>0.6): {text_mask.sum()}")
# Get the row/column range of text pixels
ys, xs = np.where(text_mask)
if len(ys) > 0:
print(f"Text region: rows [{ys.min()}-{ys.max()}], cols [{xs.min()}-{xs.max()}]")
# Check link scores within text region - do they separate words?
# Link 2 is East neighbor (right), Link 6 is West neighbor (left)
# If link between words is low, they should separate
row_mid = (ys.min() + ys.max()) // 2
print(f"\nHorizontal link scores at row {row_mid} (East neighbor):")
link_east = link_scores[2, row_mid, :] # E neighbor
for x in range(xs.min(), xs.max()+1):
ps = pixel_scores[row_mid, x]
le = link_east[x]
marker = "TEXT" if ps > 0.6 else " "
link_marker = "LINK" if le > 0.5 else "gap "
if ps > 0.3:
print(f" col={x:3d}: pixel={ps:.3f} [{marker}] east_link={le:.3f} [{link_marker}]")
# Also check if there are distinct "gap" regions in pixel scores
print(f"\nPixel scores along row {row_mid}:")
for x in range(max(0, xs.min()-2), min(pixel_scores.shape[1], xs.max()+3)):
ps = pixel_scores[row_mid, x]
bar = "█" * int(ps * 40)
print(f" col={x:3d}: {ps:.3f} {bar}")
# Try different thresholds
for thresh in [0.5, 0.6, 0.7, 0.8, 0.9]:
mask = pixel_scores > thresh
n = mask.sum()
# Connected components using simple scan
from scipy import ndimage
try:
labels, n_comps = ndimage.label(mask)
print(f"\nThreshold {thresh}: {n} pixels, {n_comps} components")
for c in range(1, min(n_comps+1, 10)):
comp_mask = labels == c
area = comp_mask.sum()
ys_c, xs_c = np.where(comp_mask)
print(f" Component {c}: area={area}, cols=[{xs_c.min()}-{xs_c.max()}]")
except ImportError:
# Fallback without scipy
print(f"Threshold {thresh}: {n} pixels")
|