229 lines
9.9 KiB
Python
229 lines
9.9 KiB
Python
import os
|
||
os.environ["DISABLE_MODEL_SOURCE_CHECK"] = "True"
|
||
|
||
import cv2
|
||
import numpy as np
|
||
import math
|
||
import fitz # PyMuPDF
|
||
import paddlex as px # Using Paddlex for Layout
|
||
from paddleocr import SealTextDetection
|
||
|
||
# Tool: Standard Polar Unwarp
|
||
def polar_unwarp(img, center, radius, start_theta, angular_extent):
|
||
if angular_extent <= 0: return None
|
||
strip_w = int(angular_extent * radius)
|
||
strip_h = int(radius * 0.6)
|
||
if strip_w <= 0 or strip_h <= 0: return None
|
||
strip = np.zeros((strip_h, strip_w, 3), dtype=np.uint8)
|
||
for y in range(strip_h):
|
||
r = radius - y
|
||
for x in range(strip_w):
|
||
theta = start_theta + angular_extent * (x / strip_w)
|
||
src_x = center[0] + r * math.cos(theta)
|
||
src_y = center[1] + r * math.sin(theta)
|
||
if 0 <= src_x < img.shape[1]-1 and 0 <= src_y < img.shape[0]-1:
|
||
patch = cv2.getRectSubPix(img, (1, 1), (float(src_x), float(src_y)))
|
||
strip[y, x] = patch[0, 0]
|
||
else:
|
||
strip[y, x] = [255, 255, 255]
|
||
return strip
|
||
|
||
def calculate_precise_arc(polygons, center):
|
||
initial_clusters = []
|
||
gap_thresh = math.radians(15)
|
||
for poly in polygons:
|
||
thetas = sorted([math.atan2(p[1] - center[1], p[0] - center[0]) for i, p in enumerate(poly)])
|
||
if not thetas: continue
|
||
max_gap = 0
|
||
gap_idx = -1
|
||
for i in range(len(thetas)):
|
||
gap = (thetas[0] + 2*math.pi - thetas[i]) if i == len(thetas)-1 else (thetas[i+1]-thetas[i])
|
||
if gap > max_gap: max_gap = gap; gap_idx = i
|
||
if gap_idx == len(thetas) - 1:
|
||
t_arc = thetas
|
||
else:
|
||
t_arc = thetas[gap_idx+1:] + [t + 2*math.pi for t in thetas[:gap_idx+1]]
|
||
if not t_arc: continue
|
||
curr = [t_arc[0]]
|
||
for i in range(1, len(t_arc)):
|
||
if t_arc[i] - t_arc[i-1] > gap_thresh:
|
||
initial_clusters.append({'start': curr[0], 'end': curr[-1]})
|
||
curr = [t_arc[i]]
|
||
else:
|
||
curr.append(t_arc[i])
|
||
initial_clusters.append({'start': curr[0], 'end': curr[-1]})
|
||
if not initial_clusters: return 0.0, 0.0
|
||
initial_clusters.sort(key=lambda x: x['start'])
|
||
merged = []
|
||
merge_thresh = math.radians(45)
|
||
if initial_clusters:
|
||
curr = initial_clusters[0]
|
||
for i in range(1, len(initial_clusters)):
|
||
nxt = initial_clusters[i]
|
||
if nxt['start'] - curr['end'] < merge_thresh:
|
||
curr['end'] = max(curr['end'], nxt['end'])
|
||
else:
|
||
merged.append(curr)
|
||
curr = nxt
|
||
merged.append(curr)
|
||
candidates = []
|
||
for m in merged:
|
||
st, en = m['start'], m['end']
|
||
ex = en - st
|
||
mid = (st + en) / 2
|
||
dist_to_top = abs(((mid + math.pi/2 + math.pi) % (2*math.pi)) - math.pi)
|
||
weight = math.exp(-0.5 * (dist_to_top / (math.pi/2))**2)
|
||
candidates.append({'start': st, 'end': en, 'extent': ex, 'score': ex * weight})
|
||
candidates.sort(key=lambda x: x['score'], reverse=True)
|
||
best = candidates[0]
|
||
return best['start'], best['end'] - best['start']
|
||
|
||
def extract_pdf_page(pdf_path, page_num=0):
|
||
doc = fitz.open(pdf_path)
|
||
page = doc.load_page(page_num)
|
||
pix = page.get_pixmap(matrix=fitz.Matrix(2, 2))
|
||
img = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.h, pix.w, pix.n)
|
||
if pix.n == 4: img = cv2.cvtColor(img, cv2.COLOR_RGBA2RGB)
|
||
return cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
|
||
|
||
def run_layout_detection(image_path):
|
||
print("Initializing Paddlex PP-DocLayout-L...")
|
||
model = px.create_model("PP-DocLayout-L")
|
||
output = model.predict(image_path, batch_size=1)
|
||
|
||
all_regions = []
|
||
for res in output:
|
||
# Paddlex 3.0 result structure: res is a dict with 'boxes' key
|
||
boxes = res.get('boxes', [])
|
||
for box in boxes:
|
||
# box structure: label_name, label, score, coordinate
|
||
label_name = box.get('label_name', box.get('label', 'unknown'))
|
||
score = box.get('score', 0.0)
|
||
coords = box.get('coordinate')
|
||
print(f"Detected: {label_name} (Score: {score:.2f}) at {coords}")
|
||
all_regions.append({
|
||
'label': label_name,
|
||
'score': score,
|
||
'box': coords
|
||
})
|
||
return all_regions
|
||
|
||
def process_full_workflow(pdf_path, output_dir="report_viz"):
|
||
if not os.path.exists(output_dir): os.makedirs(output_dir)
|
||
print(f"Rendering PDF {pdf_path} Page 1...")
|
||
page_img = extract_pdf_page(pdf_path)
|
||
doc_path = os.path.join(output_dir, "doc_page.png")
|
||
cv2.imwrite(doc_path, page_img)
|
||
|
||
print("Running Layout Detection via Paddlex...")
|
||
all_regions = run_layout_detection(doc_path)
|
||
|
||
page_viz = page_img.copy()
|
||
seal_boxes = []
|
||
for reg in all_regions:
|
||
box = reg.get('box')
|
||
label = reg.get('label')
|
||
score = reg.get('score', 0.0)
|
||
|
||
# In Paddlex 3.0 DocLayout, 'seal' is index 16 or name 'seal'
|
||
# Let's match by name.
|
||
is_seal = (label == 'seal')
|
||
|
||
if score > 0.2: # Low threshold for debugging
|
||
x1, y1, x2, y2 = [int(v) for v in box]
|
||
color = (0, 0, 255) if is_seal else (0, 255, 0)
|
||
cv2.rectangle(page_viz, (x1, y1), (x2, y2), color, 2)
|
||
cv2.putText(page_viz, f"{label} {score:.2f}", (x1, y1-5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 1)
|
||
|
||
if is_seal:
|
||
seal_boxes.append(box)
|
||
|
||
cv2.imwrite(os.path.join(output_dir, "doc_layout_viz.png"), page_viz)
|
||
|
||
seal_results = []
|
||
print(f"Processing {len(seal_boxes)} detected seals...")
|
||
det_model = SealTextDetection(model_name="PP-OCRv4_server_seal_det")
|
||
|
||
for i, box in enumerate(seal_boxes):
|
||
x1, y1, x2, y2 = [int(v) for v in box]
|
||
pad = 40
|
||
y1_p, y2_p = max(0, y1-pad), min(page_img.shape[0], y2+pad)
|
||
x1_p, x2_p = max(0, x1-pad), min(page_img.shape[1], x2+pad)
|
||
seal_crop = page_img[y1_p:y2_p, x1_p:x2_p]
|
||
crop_path = os.path.join(output_dir, f"seal_crop_{i}.png")
|
||
cv2.imwrite(crop_path, seal_crop)
|
||
|
||
print(f"Refining Seal #{i}...")
|
||
output = det_model.predict(crop_path, batch_size=1)
|
||
all_polygons = []
|
||
for res in output:
|
||
# SealTextDetection returns dt_polys directly in the result
|
||
polys = res.get('dt_polys') if isinstance(res, dict) else None
|
||
if polys:
|
||
all_polygons.extend(polys)
|
||
print(f" Found {len(polys)} text polygons in seal #{i}")
|
||
|
||
ch, cw = seal_crop.shape[:2]
|
||
center = [cw // 2, ch // 2]
|
||
radius = min(cw, ch) // 2 - 10
|
||
|
||
start_theta, extent = calculate_precise_arc(all_polygons, center)
|
||
marked = seal_crop.copy()
|
||
for p in all_polygons:
|
||
cv2.polylines(marked, [np.array(p, dtype=np.int32)], True, (0, 255, 0), 2)
|
||
|
||
unwarp_name = f"seal_unwarp_{i}.png"
|
||
unwarp_path = os.path.join(output_dir, unwarp_name)
|
||
unwarp = None
|
||
if extent > 0:
|
||
unwarp = polar_unwarp(seal_crop, center, radius, start_theta, extent)
|
||
if unwarp is not None:
|
||
cv2.imwrite(unwarp_path, unwarp)
|
||
def draw_line(m, theta, color):
|
||
x = center[0] + radius * math.cos(theta)
|
||
y = center[1] + radius * math.sin(theta)
|
||
cv2.line(m, (int(center[0]), int(center[1])), (int(x), int(y)), color, 2)
|
||
draw_line(marked, start_theta, (255, 0, 0))
|
||
draw_line(marked, start_theta + extent, (0, 0, 255))
|
||
|
||
marked_name = f"seal_marked_{i}.png"
|
||
cv2.imwrite(os.path.join(output_dir, marked_name), marked)
|
||
seal_results.append({'index': i, 'crop': f"seal_crop_{i}.png", 'marked': marked_name, 'unwarp': unwarp_name if unwarp is not None else None})
|
||
|
||
# Integrated HTML Template
|
||
html = f"""
|
||
<html><body style="font-family: sans-serif; padding: 20px; background: #fdfdfd;">
|
||
<h1>Integrated Workflow: Paddlex Layout Analysis + OCR</h1>
|
||
<div style="background: white; padding: 20px; border-radius: 8px; box-shadow: 0 2px 10px rgba(0,0,0,0.05); margin-bottom: 40px;">
|
||
<h3>1. Document Layout Detection (Paddlex PP-DocLayout-L)</h3>
|
||
<p>File: WTS2025-21283.pdf | Detected Regions: {len(all_regions)}</p>
|
||
<img src="doc_layout_viz.png" style="max-width: 100%; border: 1px solid #999;">
|
||
</div>
|
||
<div>
|
||
<h2>2. Refined Seal Extraction & Unwarping</h2>
|
||
{"".join([f'''
|
||
<div style="margin-bottom: 40px; border-bottom: 2px solid #eee; padding-bottom: 20px;">
|
||
<h3>Seal Area #{s['index']}</h3>
|
||
<div style="display: flex; gap: 20px;">
|
||
<div style="background:white; padding:10px; border-radius:4px; box-shadow: 0 1px 3px rgba(0,0,0,0.1);">
|
||
<p style="margin-top:0;">Detection Overlay</p>
|
||
<img src="{s['marked']}" style="max-height: 350px;">
|
||
</div>
|
||
<div style="flex-grow:1; background:white; padding:10px; border-radius:4px; box-shadow: 0 1px 3px rgba(0,0,0,0.1);">
|
||
<p style="margin-top:0;">Unwarped Organization Name</p>
|
||
{f'<img src="{s["unwarp"]}" style="max-width: 100%; border: 1px solid #ddd;">' if s['unwarp'] else '<p style="color:red;">No text arc found in this crop.</p>'}
|
||
</div>
|
||
</div>
|
||
</div>
|
||
''' for s in seal_results]) if seal_results else '<p>No seals detected for unwarping.</p>'}
|
||
</div>
|
||
</body></html>
|
||
"""
|
||
with open(os.path.join(output_dir, "index.html"), "w", encoding="utf-8") as f:
|
||
f.write(html)
|
||
print(f"Workflow Complete: {output_dir}/index.html")
|
||
|
||
if __name__ == "__main__":
|
||
pdf_path = r"src/test/resources/data/pdfs/关于中检测试技术（广东）集团有限公司检验检测资质的调查取证函（局长件）_pages11-14.pdf"
|
||
process_full_workflow(pdf_path)
|