diff --git a/copienator/name_grids.py b/copienator/name_grids.py index 510c980..f948154 100644 --- a/copienator/name_grids.py +++ b/copienator/name_grids.py @@ -31,10 +31,13 @@ def save_crop(page: Image.Image, pdf: Path, directory: Path, dpi: int) -> None: }) -def build_grids(directory: Path, copies_dir: Path) -> list[dict]: - """Rebuild deterministic 6-by-4 sheets, excluding deleted or changed copies.""" +def build_grids(directory: Path, copies_dir: Path, *, columns: int = COLUMNS, + rows: int = ROWS, stems: set[str] | None = None) -> list[dict]: + """Rebuild numbered sheets, excluding deleted or changed copies.""" crops = [] for path in sorted(directory.glob("*.jpg"), key=natural_key): + if stems is not None and path.stem not in stems: + continue pdf = copies_dir / f"{path.stem}.pdf" try: if pdf.is_file() and read_json(path.with_suffix('.json'))['source_sha256'] == fingerprint(pdf): @@ -44,8 +47,8 @@ def build_grids(directory: Path, copies_dir: Path) -> list[dict]: grids_dir = directory / "Grids" grids = [] with staged_files(grids_dir, remove=[p.name for p in grids_dir.glob('*') if p.is_file()]) as staging: - for offset in range(0, len(crops), COLUMNS * ROWS): - batch = crops[offset:offset + COLUMNS * ROWS] + for offset in range(0, len(crops), columns * rows): + batch = crops[offset:offset + columns * rows] images = [] for path in batch: with Image.open(path) as image: @@ -54,11 +57,11 @@ def build_grids(directory: Path, copies_dir: Path) -> list[dict]: height = max(image.height for image in images) border, header = 6, 32 cell_width, cell_height = width + border, height + header + border - sheet = Image.new('RGB', (COLUMNS * cell_width + border, ROWS * cell_height + border), 'black') + sheet = Image.new('RGB', (columns * cell_width + border, rows * cell_height + border), 'black') draw = ImageDraw.Draw(sheet) cells = [] - for index in range(COLUMNS * ROWS): - row, column = divmod(index, COLUMNS) + for index in range(columns * rows): + row, column = divmod(index, columns) x, y = border + column * cell_width, border + row * cell_height draw.rectangle((x, y, x + width - 1, y + height + header - 1), fill='white') if index < len(images): @@ -70,7 +73,7 @@ def build_grids(directory: Path, copies_dir: Path) -> list[dict]: draw.text((x + 5, y + 5), 'EMPTY', fill='black') filename = f"grid_{len(grids) + 1:03d}.jpg" sheet.save(staging / filename, 'JPEG', quality=95) - grids.append({'image': filename, 'columns': COLUMNS, 'rows': ROWS, 'cells': cells}) + grids.append({'image': filename, 'columns': columns, 'rows': rows, 'cells': cells}) atomic_write_json(staging / 'manifest.json', grids) return grids @@ -103,27 +106,31 @@ def name_bottom_height_fraction(workspace, stem: str) -> float | None: This is intended as a future lower boundary for margin analysis. """ directory = workspace.cutleft_dir / 'Names' - try: - cached = read_json(directory / f'{stem}_result.json') - value = cached['name_bottom_height_fraction'] - if (cached['source_sha256'] == fingerprint(workspace.copies_dir / f'{stem}.pdf') - and cached['crop_sha256'] == fingerprint(directory / f'{stem}.jpg') - and cached['name'] != MISSING - and type(value) in (int, float) and 0 < value <= 1): - return float(value) - except (OSError, ValueError, KeyError, TypeError): - pass + for candidate in (directory, directory / 'Wide'): + try: + cached = read_json(candidate / f'{stem}_result.json') + value = cached['name_bottom_height_fraction'] + if (cached['source_sha256'] == fingerprint(workspace.copies_dir / f'{stem}.pdf') + and cached['crop_sha256'] == fingerprint(candidate / f'{stem}.jpg') + and cached['name'] != MISSING + and type(value) in (int, float) and 0 < value <= 1): + return float(value) + except (OSError, ValueError, KeyError, TypeError): + pass return None -def identify_names(workspace, groups, *, client, model: str, names_text: str, overwrite: bool, - thinking_level: str | None = "low") -> dict[str, str]: +def _identify_names(workspace, groups, *, client, model: str, names_text: str, overwrite: bool, + thinking_level: str | None = "low", wide: bool = False) -> dict[str, str]: """Identify pending copies before label requests; safely fall back on any failure.""" directory = workspace.cutleft_dir / 'Names' + if wide: + directory /= 'Wide' result = {stem: MISSING for stem in groups} if not groups or not directory.is_dir(): return result # Older workspaces retain their existing label/name workflow. - grids = build_grids(directory, workspace.copies_dir) + grids = build_grids(directory, workspace.copies_dir, columns=1 if wide else COLUMNS, + rows=4 if wide else ROWS, stems=set(groups) if wide else None) roster = {line.strip() for line in names_text.splitlines() if line.strip()} - {MISSING, 'Unknown', 'Continued'} for grid in grids: selected = {cell['copy']: cell for cell in grid['cells'] if cell['copy'] in groups} @@ -145,8 +152,9 @@ def identify_names(workspace, groups, *, client, model: str, names_text: str, ov if not pending: continue positions = '\n'.join(f"Cell {c['cell']}: row {c['row']}, column {c['column']}" for c in grid['cells']) - prompt = f"""Identify student names from the top-left corners of exam first pages. -The image is a grid with exactly 6 columns and 4 rows. Read rows left to right, + region = 'full-width top quarters' if wide else 'top-left corners' + prompt = f"""Identify student names from the {region} of exam first pages. +The image is a grid with exactly {grid['columns']} columns and {grid['rows']} rows. Read rows left to right, then top to bottom. Thick black borders separate cells; headers show row, column, and cell number. Each cell belongs to a DIFFERENT student. Never carry a name from one cell into another. Do not include class names such as MPSI 3 in the name. @@ -165,7 +173,7 @@ readable handwriting to an exact authorized name below. Do not guess or assign names by elimination. Return "Missing" for an unreadable, absent, or ambiguous name. Authorized names: {names_text}""" - print(f"Identifying student names: {grid['image']} ({len(grid['cells'])} occupied cells, model={model}, thinking={thinking_level or 'default'}).", flush=True) + print(f"Identifying student names{' (full-width retry)' if wide else ''}: {grid['image']} ({len(grid['cells'])} occupied cells, model={model}, thinking={thinking_level or 'default'}).", flush=True) try: response = client.models.generate_content( model=model, @@ -202,3 +210,49 @@ Authorized names: print(f"Name recognition failed for {grid['image']}: {type(exc).__name__}: {exc}; using label/name fallback.") # Do not cache transport/schema failures: a resumed run should retry. return result + + +def identify_names(workspace, groups, *, client, model: str, names_text: str, overwrite: bool, + thinking_level: str | None = "low") -> dict[str, str]: + """Try narrow grids, then collated full-width headers before label fallback.""" + options = dict(client=client, model=model, names_text=names_text, + overwrite=overwrite, thinking_level=thinking_level) + result = _identify_names(workspace, groups, **options) + directory = workspace.cutleft_dir / 'Names' + if not directory.is_dir(): + return result + wide_directory = directory / 'Wide' + pending = {} + for stem, name in result.items(): + if name != MISSING: + continue + pdf = workspace.copies_dir / f'{stem}.pdf' + try: + source_hash = fingerprint(pdf) + crop_path = wide_directory / f'{stem}.jpg' + try: + current = (crop_path.is_file() and + read_json(crop_path.with_suffix('.json'))['source_sha256'] == source_hash) + except (OSError, ValueError, KeyError, TypeError): + current = False + if not current: + from pdf2image import convert_from_path + page = convert_from_path(pdf, dpi=200, first_page=1, last_page=1, + use_cropbox=False)[0] + try: + height = max(1, page.height // 4) + with staged_files(wide_directory) as staging: + page.crop((0, 0, page.width, height)).save( + staging / crop_path.name, 'JPEG', quality=95) + atomic_write_json(staging / f'{stem}.json', { + 'source_sha256': source_hash, + 'height_fraction': height / page.height, + }) + finally: + page.close() + pending[stem] = groups[stem] + except Exception as exc: + print(f'Full-width name crop failed for {stem}: {type(exc).__name__}: {exc}; using label/name fallback.') + if pending: + result.update(_identify_names(workspace, pending, wide=True, **options)) + return result diff --git a/tests/test_name_grids.py b/tests/test_name_grids.py index def8a58..c941c6a 100644 --- a/tests/test_name_grids.py +++ b/tests/test_name_grids.py @@ -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)