sebastien
0a4f4a80d0
Bundle the trained heatmap/image orientation model and run it after provisional blank-page and split analysis. Confirm candidates with Gemini using low thinking; resolve mixed A3 sheet ordering with the complete statement and medium thinking. Keep A4 scans independent and export paired A3 scans in booklet reading order. Reanalyse uniformly inverted source sheets, validate ordering permutations, and preserve source provenance, corrected rotations and reading slots in manual review. Document the experiments, rejected approaches, calibration and runtime limits. Validation: focused splitting/orientation regression suites pass; bundled model checksum and wheel assets verified. Full Interro07 offline run: 44.26 seconds, eight candidates, mocked upright confirmations, source hashes and geometry intact.
291 lines
15 KiB
Python
291 lines
15 KiB
Python
import hashlib
|
|
import json
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import Mock, patch
|
|
|
|
import numpy as np
|
|
import pymupdf
|
|
from pydantic import ValidationError
|
|
|
|
from copienator.auto_page_split import run
|
|
from copienator.split_orientation import GeminiVerifier, PageOrder, sheets, validate_order
|
|
|
|
|
|
class SheetOrientationTests(unittest.TestCase):
|
|
def setUp(self):
|
|
self.temp = tempfile.TemporaryDirectory()
|
|
self.addCleanup(self.temp.cleanup)
|
|
self.root = Path(self.temp.name)
|
|
self.source = self.root/'copy.pdf'
|
|
self.statement = self.root/'enonce.pdf'
|
|
with pymupdf.open() as d:
|
|
d.new_page().insert_text((40,40),'Full exam statement')
|
|
d.save(self.statement)
|
|
self.analyses = []
|
|
|
|
def create(self, formats=('a3','a3'), rotations=None, mixed=False):
|
|
with pymupdf.open() as d:
|
|
for i,kind in enumerate(formats):
|
|
p = d.new_page(width=1190 if kind=='a3' else 595,height=842)
|
|
if kind=='a3':
|
|
p.insert_text((180,280) if mixed and i==0 else (50,70),
|
|
str(4 if i==0 else 2), fontsize=24,
|
|
rotate=180 if mixed and i==0 else 0)
|
|
p.insert_text((650,70),str(1 if i==0 else 3),fontsize=24)
|
|
else:
|
|
p.insert_text((50,70),f'A4-{i}',fontsize=24)
|
|
p.set_rotation((rotations or [0]*len(formats))[i])
|
|
d.save(self.source)
|
|
self.original = self.source.read_bytes()
|
|
|
|
def analyse(self, page, dpi, kind, *, blank=()):
|
|
self.analyses.append((page.number,page.rotation))
|
|
w = 200 if kind=='a3' else 100
|
|
outputs = []
|
|
for j,side in enumerate(['left','right'] if kind=='a3' else ['whole']):
|
|
keep = (page.number,j) not in blank
|
|
outputs.append(dict(side=side, keep=keep, state='content' if keep else 'blank',
|
|
rect_px=[j*100,0,(j+1)*100,100]))
|
|
return dict(format=kind,size_px=[w,100],outputs=outputs,crossings=[],
|
|
extensions_mm=[0,0],review_reasons=[]),np.full((100,w,3),255,np.uint8)
|
|
|
|
def execute(self, flags, verdicts=(), permutation=None, blank=()):
|
|
screen = Mock()
|
|
screen.predict.side_effect = [dict(flag=f,score=.8 if f else .1,threshold=.3)for f in flags]
|
|
verifier = Mock()
|
|
verifier.orientation.side_effect = list(verdicts)
|
|
verifier.order.return_value = permutation
|
|
with (patch('copienator.fast_page_split.analyse_pdf_page_fast',
|
|
side_effect=lambda p,d,k:self.analyse(p,d,k,blank=blank)),
|
|
patch('copienator.split_orientation.detector',return_value=screen),
|
|
patch('copienator.split_orientation.GeminiVerifier',return_value=verifier)):
|
|
summary = run(self.source,self.root/'result',dpi=100,orientation=True,statement=self.statement)
|
|
self.assertEqual(self.source.read_bytes(),self.original)
|
|
report = json.loads((self.root/'result/report.json').read_text())
|
|
return summary,report,screen,verifier
|
|
|
|
def texts(self):
|
|
with pymupdf.open(self.root/'result/pdfs/copy.pdf')as d:
|
|
return [p.get_text().strip()for p in d]
|
|
|
|
def test_upright_a3_booklet_order_and_no_requests(self):
|
|
self.create()
|
|
_,report,_,verifier = self.execute([False]*4)
|
|
self.assertEqual(self.texts(),['1','2','3','4'])
|
|
verifier.orientation.assert_not_called()
|
|
verifier.order.assert_not_called()
|
|
self.assertEqual([o['output_page']for p in report['pages']for o in p['outputs']],[4,1,2,3])
|
|
|
|
def test_a4_inversion_is_independent_and_keeps_positions(self):
|
|
self.create(('a4','a4'),rotations=[180,0])
|
|
_,report,_,verifier = self.execute([True,False],['inverted'])
|
|
self.assertEqual(self.texts(),['A4-0','A4-1'])
|
|
self.assertEqual(verifier.orientation.call_count,1)
|
|
verifier.order.assert_not_called()
|
|
self.assertEqual(report['pages'][0]['applied_rotation'],180)
|
|
self.assertNotIn('applied_rotation',report['pages'][1])
|
|
self.assertEqual(self.analyses,[(0,180),(1,0),(0,0)])
|
|
|
|
def test_entire_a3_sheet_inverted_reanalyses_both_source_faces(self):
|
|
self.create(rotations=[180,180])
|
|
_,report,_,verifier = self.execute([True,False,False,False],['inverted']*4)
|
|
self.assertEqual(self.texts(),['1','2','3','4'])
|
|
self.assertEqual(verifier.orientation.call_count,4)
|
|
verifier.order.assert_not_called()
|
|
self.assertEqual(self.analyses,[(0,180),(1,180),(0,0),(1,0)])
|
|
self.assertTrue(all(p['applied_rotation']==180 for p in report['pages']))
|
|
self.assertTrue(all('orientation_before'in p for p in report['pages']))
|
|
|
|
def test_blank_siblings_are_not_sent_to_gemini_and_are_retained_until_ordering(self):
|
|
self.create(rotations=[180,180])
|
|
summary,_,screen,verifier=self.execute([True],['inverted'],blank=((0,1),(1,0),(1,1)))
|
|
self.assertEqual(screen.predict.call_count,1)
|
|
self.assertEqual(verifier.orientation.call_count,1)
|
|
self.assertEqual(summary['blank_removed'],3)
|
|
self.assertEqual(len(self.texts()),1)
|
|
self.assertEqual(len(self.analyses),4)
|
|
|
|
def test_inconclusive_candidate_assumed_upright_without_sibling_requests(self):
|
|
self.create()
|
|
_,report,_,verifier=self.execute([True,False,False,False],['inconclusive'])
|
|
self.assertEqual(verifier.orientation.call_count,1)
|
|
self.assertEqual(len(self.analyses),2)
|
|
self.assertEqual(self.texts(),['1','2','3','4'])
|
|
self.assertEqual(report['pages'][0]['outputs'][0]['orientation']['verdict'],'upright')
|
|
|
|
def test_mixed_sheet_rotates_only_inverted_logical_page_and_uses_given_order(self):
|
|
self.create(mixed=True)
|
|
_,report,_,verifier=self.execute([True,False,False,False],
|
|
['inverted','upright','inconclusive','upright'],[2,3,4,1])
|
|
self.assertEqual(len(self.analyses),2)
|
|
self.assertEqual(self.texts(),['1','2','3','4'])
|
|
verifier.order.assert_called_once()
|
|
images,blank=verifier.order.call_args.args
|
|
self.assertEqual(len(images),4)
|
|
self.assertEqual(blank,[False]*4)
|
|
self.assertEqual(report['pages'][0]['outputs'][0]['rotation'],180)
|
|
with pymupdf.open(self.root/'result/pdfs/copy.pdf')as d:
|
|
self.assertGreater(d[3].get_text('blocks')[0][1],400)
|
|
|
|
def test_inconclusive_sibling_takes_mixed_branch(self):
|
|
self.create()
|
|
_,_,_,verifier=self.execute([True,False,False,False],
|
|
['inverted','inconclusive','inverted','inverted'],[2,3,4,1])
|
|
verifier.order.assert_called_once()
|
|
self.assertEqual(len(self.analyses),2)
|
|
|
|
def test_invalid_gemini_order_aborts_without_touching_source(self):
|
|
self.create()
|
|
with self.assertRaises(ValueError):
|
|
self.execute([True,False,False,False],['inverted','upright','upright','upright'],[1,1,3,4])
|
|
self.assertEqual(self.source.read_bytes(),self.original)
|
|
self.assertFalse((self.root/'result/pdfs/copy.pdf').exists())
|
|
|
|
def test_api_failure_aborts_without_guessing_upright(self):
|
|
self.create()
|
|
with self.assertRaises(RuntimeError):
|
|
self.execute([True,False,False,False],[RuntimeError('API unavailable')])
|
|
self.assertEqual(self.source.read_bytes(),self.original)
|
|
|
|
def test_orphan_a3_is_flagged_and_a4_breaks_pairing(self):
|
|
self.assertEqual(list(sheets([{'format':k}for k in ['a3','a4','a3','a3','a4']])),
|
|
[[0],[1],[2,3],[4]])
|
|
self.create(('a3',))
|
|
_,report,_,_=self.execute([False,False])
|
|
self.assertIn('unpaired-a3-sheet',report['pages'][0]['review_reasons'])
|
|
|
|
def review_saved(self, report, restore_blank=False):
|
|
import shutil
|
|
from copienator import EvaluationWorkspace
|
|
from copienator.commands.automatic_page_splitter import decision_path
|
|
from copienator.commands.page_splitter import PDFPreviewer
|
|
workspace=EvaluationWorkspace(self.root)
|
|
(self.root/'labels').write_text('Ex 1\n')
|
|
workspace.original_copies_dir.mkdir()
|
|
workspace.copies_dir.mkdir()
|
|
original=workspace.original_copies_dir/self.source.name
|
|
output=workspace.copies_dir/self.source.name
|
|
shutil.copy2(self.source,original)
|
|
shutil.copy2(self.root/'result/pdfs/copy.pdf',output)
|
|
path=decision_path(workspace,'copy');path.parent.mkdir(parents=True)
|
|
path.write_text(json.dumps(dict(source_sha256=hashlib.sha256(original.read_bytes()).hexdigest(),
|
|
output_sha256=hashlib.sha256(output.read_bytes()).hexdigest(),pages=report['pages'])))
|
|
expected=pymupdf.open(output)
|
|
self.addCleanup(lambda document=expected: None if document.is_closed else document.close())
|
|
preview=PDFPreviewer.__new__(PDFPreviewer)
|
|
preview.workspace=workspace;preview.inputs=[original]
|
|
preview.num=preview.global_rotation=0;preview.history=[];preview.output_dir=None
|
|
preview.failed=False;preview.master=Mock();preview.load_page=Mock()
|
|
self.assertTrue(preview.setup_next_file())
|
|
preview._initialize_current_page_settings()
|
|
for _ in range(2):preview.confirm_and_next_page(keep='both' if restore_blank else None)
|
|
self.assertFalse(preview.failed)
|
|
self.assertEqual(original.read_bytes(),self.original)
|
|
with pymupdf.open(output)as actual:
|
|
if restore_blank:
|
|
self.assertEqual([p.get_text().strip()for p in actual],['1','2','3','4'])
|
|
expected.close()
|
|
expected=pymupdf.open(output)
|
|
self.addCleanup(lambda document=expected: None if document.is_closed else document.close())
|
|
self.assertEqual(len(actual),len(expected))
|
|
for a,b in zip(actual,expected):
|
|
pa,pb=a.get_pixmap(dpi=40),b.get_pixmap(dpi=40)
|
|
self.assertEqual((pa.width,pa.height),(pb.width,pb.height))
|
|
diff=np.abs(np.frombuffer(pa.samples,np.uint8).astype(float)-np.frombuffer(pb.samples,np.uint8))
|
|
self.assertLess(diff.mean(),.1)
|
|
from copienator.orientation import load_provenance
|
|
self.assertEqual([p['origin']for p in load_provenance(workspace,output,4)],[1,2,2,1])
|
|
|
|
def test_manual_reinstatement_of_a_blank_page_uses_its_original_booklet_slot(self):
|
|
self.create()
|
|
_,report,_,_=self.execute([False]*3,blank=((0,0),))
|
|
self.assertEqual(report['pages'][0]['outputs'][0]['reading_order'],4)
|
|
self.assertEqual(self.texts(),['1','2','3'])
|
|
self.review_saved(report,restore_blank=True)
|
|
|
|
def test_manual_acceptance_preserves_uniform_source_rotation_and_booklet_order(self):
|
|
self.create(rotations=[180,180])
|
|
_,report,_,_=self.execute([True,False,False,False],['inverted']*4)
|
|
self.review_saved(report)
|
|
|
|
def test_manual_acceptance_preserves_mixed_logical_rotation_and_gemini_order(self):
|
|
self.create(mixed=True)
|
|
_,report,_,_=self.execute([True,False,False,False],
|
|
['inverted','upright','upright','upright'],[2,3,4,1])
|
|
self.review_saved(report)
|
|
|
|
def test_manual_review_keeps_a4_before_following_a3_pages(self):
|
|
from copienator.commands.page_splitter import PDFPreviewer
|
|
with pymupdf.open()as d:
|
|
d.new_page(width=595,height=842).insert_text((50,70),'A4')
|
|
p=d.new_page(width=1190,height=842)
|
|
p.insert_text((50,70),'B');p.insert_text((650,70),'C');d.save(self.source)
|
|
preview=PDFPreviewer.__new__(PDFPreviewer)
|
|
preview.doc=pymupdf.open(self.source)
|
|
preview.file_rotation=preview.global_rotation=0
|
|
preview.base_name='copy';preview.split_dir=self.root/'split';preview.reorder_dir=self.root/'reorder'
|
|
preview.final_file=self.root/'final.pdf';preview.output_dir=None;preview.automatic_review=True
|
|
preview.page_settings=[
|
|
dict(line_x=297.5,rotation=0,keep='as_is',automatic_outputs=[
|
|
dict(side='whole',crop_points=[0,0,595,842],reading_order=1)]),
|
|
dict(line_x=595,rotation=0,keep='both',automatic_outputs=[
|
|
dict(side='left',crop_points=[0,0,595,842],reading_order=2),
|
|
dict(side='right',crop_points=[595,0,1190,842],reading_order=3)])]
|
|
preview.split_pdf();preview.reorder_pdfs();preview.concate_files()
|
|
with pymupdf.open(preview.final_file)as d:
|
|
self.assertEqual([p.get_text().strip()for p in d],['A4','B','C'])
|
|
|
|
def test_secondary_cuts_follow_the_corrected_half_display(self):
|
|
from copienator.commands.page_splitter import PDFPreviewer
|
|
with pymupdf.open()as d:
|
|
p=d.new_page(width=1190,height=842)
|
|
for x,label in [(50,'A'),(400,'B'),(650,'C')]:p.insert_text((x,70),label,fontsize=24)
|
|
d.save(self.source)
|
|
preview=PDFPreviewer.__new__(PDFPreviewer)
|
|
preview.doc=pymupdf.open(self.source)
|
|
preview.file_rotation=preview.global_rotation=0
|
|
preview.base_name='copy';preview.split_dir=self.root/'split';preview.reorder_dir=self.root/'reorder'
|
|
preview.final_file=self.root/'final.pdf';preview.output_dir=None;preview.automatic_review=True
|
|
preview.page_settings=[dict(line_x=595,rotation=0,keep='both',secondary_splitters=[297.5],
|
|
automatic_outputs=[dict(side='left',crop_points=[0,0,595,842],rotation=180,reading_order=1),
|
|
dict(side='right',crop_points=[595,0,1190,842],reading_order=2)])]
|
|
preview.split_pdf();preview.reorder_pdfs();preview.concate_files()
|
|
with pymupdf.open(preview.final_file)as d:
|
|
self.assertEqual([p.get_text().strip()for p in d],['B','A','C'])
|
|
|
|
def test_order_schema_rejects_booleans_strings_duplicates_and_omissions(self):
|
|
for order in [[True,2,3,4],['1',2,3,4]]:
|
|
with self.assertRaises(ValidationError):PageOrder.model_validate({'order':order})
|
|
for order in [[1,2,2,4],[1,2,3],[1,2,3,5]]:
|
|
with self.assertRaises(ValueError):validate_order(order,4)
|
|
|
|
def test_gemini_order_includes_entire_statement_and_all_images(self):
|
|
verifier=GeminiVerifier(self.statement)
|
|
verifier.client=Mock()
|
|
verifier.client.models.generate_content.return_value.text='{"order":[2,3,4,1]}'
|
|
self.assertEqual(verifier.order([b'png']*4,[False,True,False,False]),[2,3,4,1])
|
|
config=verifier.client.models.generate_content.call_args.kwargs['config']
|
|
self.assertEqual(config.thinking_config.thinking_level.value,'MEDIUM')
|
|
parts=verifier.client.models.generate_content.call_args.kwargs['contents'][0].parts
|
|
blobs=[p.inline_data for p in parts if p.inline_data]
|
|
self.assertEqual([b.mime_type for b in blobs],['application/pdf']+['image/png']*4)
|
|
self.assertEqual(blobs[0].data,self.statement.read_bytes())
|
|
self.assertEqual(sum(b.data==b'png'for b in blobs),4)
|
|
|
|
def test_gemini_orientation_uses_low_thinking(self):
|
|
verifier=GeminiVerifier()
|
|
verifier.client=Mock()
|
|
verifier.client.models.generate_content.return_value.text='{"orientation":"upright"}'
|
|
self.assertEqual(verifier.orientation(b'png'),'upright')
|
|
config=verifier.client.models.generate_content.call_args.kwargs['config']
|
|
self.assertEqual(config.thinking_config.thinking_level.value,'LOW')
|
|
|
|
def test_missing_statement_is_an_error_for_mixed_ordering(self):
|
|
verifier=GeminiVerifier()
|
|
verifier.client=Mock()
|
|
with self.assertRaisesRegex(ValueError,'énoncé'):
|
|
verifier.order([b'png']*4,[False]*4)
|
|
verifier.client.models.generate_content.assert_not_called()
|