better name finding (use top part, full width)
This commit is contained in:
1 parent
9d0e40cc4f
commit
5a54b5ac2b
2 files changed
+155
-24
No files matched your search
@@ -188,3 +188,80 @@ class NameGridTests(unittest.TestCase):
|
||||
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)
|
||||
Reference in new issue
Block a user