268 lines
15 KiB
Python
268 lines
15 KiB
Python
import json
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import Mock, patch
|
|
|
|
from PIL import Image
|
|
|
|
from copienator import EvaluationWorkspace, ExitCode, read_json
|
|
from copienator import name_grids as names
|
|
from copienator.commands import cutleft, gemini_for_labels as labels
|
|
|
|
|
|
class NameGridTests(unittest.TestCase):
|
|
def setUp(self):
|
|
temporary = tempfile.TemporaryDirectory()
|
|
self.addCleanup(temporary.cleanup)
|
|
self.workspace = EvaluationWorkspace(Path(temporary.name))
|
|
self.workspace.copies_dir.mkdir()
|
|
self.workspace.cutleft_dir.mkdir()
|
|
self.directory = self.workspace.cutleft_dir / 'Names'
|
|
(self.workspace.root / 'names').write_text('Alice\nBob\n')
|
|
(self.workspace.root / 'labels').write_text('Ex 1\n')
|
|
self.client = Mock()
|
|
|
|
def crop(self, number):
|
|
pdf = self.workspace.copies_dir / f'Copie{number:02d}.pdf'
|
|
pdf.write_bytes(b'pdf')
|
|
page = Image.new('RGB', (1654, 2338), 'white')
|
|
page.putpixel((0, 0), (0, 0, 0))
|
|
names.save_crop(page, pdf, self.directory, 200)
|
|
(self.workspace.cutleft_dir / f'{pdf.stem}_01.jpg').write_bytes(b'image')
|
|
return pdf
|
|
|
|
def identify(self, groups, overwrite=False):
|
|
return names.identify_names(self.workspace, groups, client=self.client, model='test',
|
|
names_text='Alice\nBob\n', overwrite=overwrite)
|
|
|
|
def test_crop_geometry_and_partial_grid_mapping(self):
|
|
for number in range(1, 26):
|
|
self.crop(number)
|
|
with Image.open(self.directory / 'Copie01.jpg') as crop:
|
|
self.assertEqual(crop.size, (220, 584))
|
|
self.assertLess(crop.getpixel((0, 0))[0], 80)
|
|
grids = names.build_grids(self.directory, self.workspace.copies_dir)
|
|
self.assertEqual([len(grid['cells']) for grid in grids], [24, 1])
|
|
self.assertEqual(grids[0]['cells'][6], {'cell': 7, 'row': 2, 'column': 1, 'copy': 'Copie07'})
|
|
self.assertEqual(grids[1]['cells'][0]['copy'], 'Copie25')
|
|
with Image.open(self.directory / 'Grids' / grids[1]['image']) as grid:
|
|
self.assertEqual(grid.size, (1362, 2494))
|
|
self.assertLess(grid.getpixel((228, 500))[0], 40)
|
|
# Deleted/stale PDFs cannot be assigned names from old crops.
|
|
(self.workspace.copies_dir / 'Copie25.pdf').unlink()
|
|
(self.workspace.copies_dir / 'Copie24.pdf').write_bytes(b'changed')
|
|
grids = names.build_grids(self.directory, self.workspace.copies_dir)
|
|
self.assertEqual(len(grids), 1)
|
|
self.assertEqual(len(grids[0]['cells']), 23)
|
|
self.assertFalse((self.directory / 'Grids' / 'grid_002.jpg').exists())
|
|
|
|
def test_validation_cache_and_roster_changes(self):
|
|
for number in range(1, 5):
|
|
self.crop(number)
|
|
self.client.models.generate_content.return_value = Mock(text=json.dumps({'cells': [
|
|
{'cell': 1, 'name': 'Alice'}, {'cell': 2, 'name': 'Outsider'},
|
|
{'cell': 3, 'name': 'Bob'}, {'cell': 3, 'name': 'Alice'},
|
|
{'cell': 24, 'name': 'Bob'},
|
|
]}))
|
|
groups = {f'Copie{i:02d}': [] for i in range(1, 5)}
|
|
expected = {'Copie01': 'Alice', **{f'Copie{i:02d}': 'Missing' for i in range(2, 5)}}
|
|
self.assertEqual(self.identify(groups), expected)
|
|
request = self.client.models.generate_content.call_args.kwargs
|
|
self.assertEqual(request['config'].thinking_config.thinking_level.value, 'LOW')
|
|
text = request['contents'][0].parts[1].text
|
|
self.assertIn('6 columns and 4 rows', text)
|
|
self.assertIn('Cell 4: row 1, column 4', text)
|
|
self.assertNotIn('Cell 5:', text)
|
|
self.assertEqual(self.identify(groups), expected)
|
|
self.assertEqual(self.client.models.generate_content.call_count, 1)
|
|
self.identify(groups, overwrite=True)
|
|
self.assertEqual(self.client.models.generate_content.call_count, 2)
|
|
names.identify_names(self.workspace, groups, client=self.client, model='test',
|
|
names_text='Bob\n', overwrite=False)
|
|
self.assertEqual(read_json(self.directory / 'Copie01_result.json')['name'], 'Missing')
|
|
|
|
def test_api_failure_is_retried_on_resume(self):
|
|
self.crop(1)
|
|
self.client.models.generate_content.side_effect = TimeoutError()
|
|
self.assertEqual(self.identify({'Copie01': []}), {'Copie01': 'Missing'})
|
|
self.assertFalse((self.directory / 'Copie01_result.json').exists())
|
|
self.identify({'Copie01': []})
|
|
self.assertEqual(self.client.models.generate_content.call_count, 2)
|
|
|
|
def test_name_box_is_saved_in_full_page_height_units(self):
|
|
pdf = self.crop(1)
|
|
self.client.models.generate_content.return_value = Mock(text=json.dumps({'cells': [
|
|
{'cell': 1, 'name': 'Alice', 'box_2d': [100, 20, 800, 950]},
|
|
]}))
|
|
self.assertEqual(self.identify({'Copie01': []}), {'Copie01': 'Alice'})
|
|
cached = read_json(self.directory / 'Copie01_result.json')
|
|
self.assertEqual(cached['box_2d'], [100, 20, 800, 950])
|
|
self.assertAlmostEqual(cached['name_bottom_height_fraction'], .8 * 584 / 2338)
|
|
self.assertAlmostEqual(names.name_bottom_height_fraction(self.workspace, 'Copie01'),
|
|
.8 * 584 / 2338)
|
|
self.identify({'Copie01': []})
|
|
self.assertEqual(self.client.models.generate_content.call_count, 1)
|
|
pdf.write_bytes(b'changed')
|
|
self.assertIsNone(names.name_bottom_height_fraction(self.workspace, 'Copie01'))
|
|
|
|
def test_invalid_or_unmatched_boxes_do_not_produce_a_height(self):
|
|
self.crop(1)
|
|
for name, box in [('Alice', [800, 20, 100, 950]),
|
|
('Alice', [0, 0, 1001, 900]), ('Alice', [1, 2]),
|
|
('Alice', ['bad', 0, 800, 900]), ('Alice', None),
|
|
('Outsider', [100, 20, 800, 950]),
|
|
('Missing', [100, 20, 800, 950])]:
|
|
with self.subTest(name=name, box=box):
|
|
self.client.models.generate_content.return_value = Mock(text=json.dumps({'cells': [
|
|
{'cell': 1, 'name': name, 'box_2d': box},
|
|
]}))
|
|
self.assertEqual(self.identify({'Copie01': []}, overwrite=True),
|
|
{'Copie01': 'Alice' if name == 'Alice' else 'Missing'})
|
|
self.assertIsNone(names.name_bottom_height_fraction(self.workspace, 'Copie01'))
|
|
|
|
def test_legacy_name_cache_is_refreshed_to_request_position(self):
|
|
self.crop(1)
|
|
self.client.models.generate_content.return_value = Mock(text=json.dumps({'cells': [
|
|
{'cell': 1, 'name': 'Alice', 'box_2d': [100, 20, 800, 950]},
|
|
]}))
|
|
self.identify({'Copie01': []})
|
|
path = self.directory / 'Copie01_result.json'
|
|
cached = read_json(path)
|
|
del cached['name_bottom_height_fraction']
|
|
path.write_text(json.dumps(cached))
|
|
self.identify({'Copie01': []})
|
|
self.assertEqual(self.client.models.generate_content.call_count, 2)
|
|
self.assertIsNotNone(names.name_bottom_height_fraction(self.workspace, 'Copie01'))
|
|
|
|
def test_known_and_missing_names_choose_different_label_requests(self):
|
|
for number in (1, 2):
|
|
self.crop(number)
|
|
(self.workspace.cutleft_dir / 'Copie01_02.jpg').write_bytes(b'image')
|
|
self.client.models.generate_content.side_effect = [
|
|
Mock(text='{"cells":[{"cell":1,"name":"Alice"},{"cell":2,"name":"Missing"}]}'),
|
|
Mock(text='{"list":[{"box_2d":[1,2,3,4],"label":"Ex 1"}]}'), Mock(text='{"list":[]}'),
|
|
Mock(text='{"name":"Bob","list":[]}'),
|
|
]
|
|
with patch.object(labels.config, 'MODEL_FOR_NAME_ID', 'custom-name-model'), \
|
|
patch.object(labels.config, 'THINKING_LEVEL_FOR_NAME', 'medium'):
|
|
status = labels.run(self.workspace, [self.workspace.root], client=self.client,
|
|
sleep=lambda _: None, max_workers=1)
|
|
self.assertEqual(status, ExitCode.SUCCESS)
|
|
requests = self.client.models.generate_content.call_args_list
|
|
self.assertEqual(len(requests), 4)
|
|
self.assertEqual(requests[0].kwargs['model'], 'custom-name-model')
|
|
self.assertEqual(requests[0].kwargs['config'].thinking_config.thinking_level.value, 'MEDIUM')
|
|
self.assertEqual(requests[1].kwargs['model'], labels.MODEL_ID)
|
|
for request in requests[1:3]:
|
|
text = request.kwargs['contents'][0].parts[1].text
|
|
self.assertNotIn('Alice', text)
|
|
self.assertNotIn('Bob', text)
|
|
self.assertNotIn('student name', text)
|
|
self.assertNotIn('name', request.kwargs['config'].response_json_schema['properties'])
|
|
self.assertIn('Alice', requests[3].kwargs['contents'][0].parts[1].text)
|
|
for stem, expected in [('Copie01_01', 'Alice'), ('Copie01_02', 'Continued'), ('Copie02_01', 'Bob')]:
|
|
self.assertEqual(read_json(self.workspace.copies_dir / f'{stem}.json')['name'], expected)
|
|
self.client.models.generate_content.reset_mock()
|
|
labels.run(self.workspace, [self.workspace.root], client=self.client, sleep=lambda _: None)
|
|
self.client.models.generate_content.assert_not_called()
|
|
|
|
def test_save_label_crops_also_saves_name_grid(self):
|
|
pdf = self.workspace.copies_dir / 'Copie01.pdf'
|
|
pdf.write_bytes(b'pdf')
|
|
page = Image.new('RGB', (1654, 2338), 'white')
|
|
with patch.object(cutleft, 'get_pdf_pages', return_value=[page]):
|
|
result = cutleft.process_single_pdf(pdf, automatic_margin=False)
|
|
cutleft.save_results(result, pdf, self.workspace.cutleft_dir)
|
|
self.assertTrue((self.directory / 'Copie01.jpg').is_file())
|
|
self.assertEqual(read_json(self.directory / 'Grids' / 'manifest.json')[0]['cells'][0]['copy'], 'Copie01')
|
|
|
|
def test_known_name_survives_label_retry(self):
|
|
self.crop(1)
|
|
image = self.workspace.cutleft_dir / 'Copie01_01.jpg'
|
|
self.client.models.generate_content.side_effect = [TimeoutError(), Mock(text='{"list":[]}')]
|
|
labels.process_copy_group(self.workspace, 'Copie01', [image], client=self.client,
|
|
labels_text='Ex 1', names_text='Alice\nBob', valid_labels={'Ex 1'},
|
|
valid_names={'Alice', 'Bob'}, overwrite=False, known_name='Alice',
|
|
sleep=lambda _: None, target_interval=0)
|
|
self.assertEqual(read_json(self.workspace.copies_dir / 'Copie01_01.json')['name'], 'Alice')
|
|
for call in self.client.models.generate_content.call_args_list:
|
|
self.assertNotIn('Alice', call.kwargs['contents'][0].parts[1].text)
|
|
|
|
|
|
class WideNameGridTests(unittest.TestCase):
|
|
setUp = NameGridTests.setUp
|
|
crop = NameGridTests.crop
|
|
identify = NameGridTests.identify
|
|
|
|
def test_recovered_name_is_used_by_label_detection(self):
|
|
self.crop(1)
|
|
self.client.models.generate_content.side_effect = [
|
|
Mock(text='{"cells":[{"cell":1,"name":"Missing"}]}'),
|
|
Mock(text='{"cells":[{"cell":1,"name":"Bob"}]}'),
|
|
Mock(text='{"list":[]}'),
|
|
]
|
|
with patch('pdf2image.convert_from_path', return_value=[Image.new('RGB', (800, 1200))]):
|
|
status = labels.run(self.workspace, [self.workspace.root], client=self.client,
|
|
sleep=lambda _: None, max_workers=1)
|
|
self.assertEqual(status, ExitCode.SUCCESS)
|
|
self.assertEqual(read_json(self.workspace.copies_dir / 'Copie01_01.json')['name'], 'Bob')
|
|
self.assertNotIn('name', self.client.models.generate_content.call_args.kwargs[
|
|
'config'].response_json_schema['properties'])
|
|
|
|
def test_missing_names_are_collated_and_cached_before_label_fallback(self):
|
|
for number in range(1, 7):
|
|
self.crop(number)
|
|
groups = {f'Copie{i:02d}': [] for i in range(1, 7)}
|
|
self.client.models.generate_content.side_effect = [
|
|
Mock(text=json.dumps({'cells': [{'cell': 1, 'name': 'Alice'}]})),
|
|
Mock(text=json.dumps({'cells': [
|
|
{'cell': 1, 'name': 'Bob', 'box_2d': [100, 700, 300, 950]},
|
|
{'cell': 2, 'name': 'Missing'},
|
|
{'cell': 3, 'name': 'Outsider'},
|
|
{'cell': 4, 'name': 'Missing'}]})),
|
|
Mock(text=json.dumps({'cells': [{'cell': 1, 'name': 'Alice'}]})),
|
|
]
|
|
def render(*args, **kwargs):
|
|
self.assertEqual(kwargs['first_page'], 1)
|
|
self.assertEqual(kwargs['last_page'], 1)
|
|
page = Image.new('RGB', (1654, 2338), 'white')
|
|
page.paste('red', (1400, 50, 1500, 100))
|
|
return [page]
|
|
with patch('pdf2image.convert_from_path', side_effect=render) as rendering:
|
|
result = self.identify(groups)
|
|
self.assertEqual(rendering.call_count, 5)
|
|
self.assertEqual(result, {'Copie01': 'Alice', 'Copie02': 'Bob',
|
|
'Copie03': 'Missing', 'Copie04': 'Missing',
|
|
'Copie05': 'Missing', 'Copie06': 'Alice'})
|
|
self.assertEqual(self.identify(groups), result)
|
|
self.assertEqual(rendering.call_count, 5)
|
|
self.assertEqual(self.client.models.generate_content.call_count, 3)
|
|
wide = self.directory / 'Wide'
|
|
manifest = read_json(wide / 'Grids' / 'manifest.json')
|
|
self.assertEqual([[c['copy'] for c in g['cells']] for g in manifest],
|
|
[['Copie02', 'Copie03', 'Copie04', 'Copie05'], ['Copie06']])
|
|
with Image.open(wide / 'Copie02.jpg') as crop:
|
|
self.assertEqual(crop.size, (1654, 584))
|
|
self.assertGreater(crop.getpixel((1450, 75))[0], 200)
|
|
self.assertLess(crop.getpixel((1450, 75))[1], 40)
|
|
self.assertAlmostEqual(names.name_bottom_height_fraction(self.workspace, 'Copie02'),
|
|
.3 * 584 / 2338)
|
|
for request in self.client.models.generate_content.call_args_list[1:]:
|
|
prompt = request.kwargs['contents'][0].parts[1].text
|
|
self.assertIn('full-width top quarters', prompt)
|
|
self.assertIn('1 columns and 4 rows', prompt)
|
|
|
|
def test_wide_failure_is_retried_on_resume(self):
|
|
self.crop(1)
|
|
self.client.models.generate_content.side_effect = [
|
|
Mock(text='{"cells":[{"cell":1,"name":"Missing"}]}'), TimeoutError(),
|
|
Mock(text='{"cells":[{"cell":1,"name":"Bob"}]}'),
|
|
]
|
|
with patch('pdf2image.convert_from_path', return_value=[Image.new('RGB', (800, 1200))]) as render:
|
|
self.assertEqual(self.identify({'Copie01': []}), {'Copie01': 'Missing'})
|
|
self.assertFalse((self.directory / 'Wide' / 'Copie01_result.json').exists())
|
|
self.assertEqual(self.identify({'Copie01': []}), {'Copie01': 'Bob'})
|
|
render.assert_called_once()
|
|
self.assertEqual(self.client.models.generate_content.call_count, 3)
|