156 lines
7.9 KiB
Python
156 lines
7.9 KiB
Python
"""First-page name crops, numbered contact sheets, and roster-checked recognition."""
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
|
|
from google.genai import types
|
|
from PIL import Image, ImageDraw
|
|
from pydantic import BaseModel, Field
|
|
|
|
from copienator import atomic_write_json, read_json
|
|
from copienator.annotation_layout import fingerprint
|
|
from copienator.filesystem import staged_files
|
|
from copienator.gemini_config import fresh_seed, thinking_config
|
|
from copienator.utils import natural_key
|
|
|
|
COLUMNS, ROWS = 6, 4
|
|
MARGIN_MM = 28
|
|
MISSING = "Missing"
|
|
|
|
|
|
def save_crop(page: Image.Image, pdf: Path, directory: Path, dpi: int) -> None:
|
|
"""Crop independently of the adjustable label strip, from the true top left."""
|
|
width = min(page.width, max(1, round(MARGIN_MM * dpi / 25.4)))
|
|
crop = page.crop((0, 0, width, max(1, page.height // 4)))
|
|
with staged_files(directory) as staging:
|
|
crop.save(staging / f"{pdf.stem}.jpg", "JPEG", quality=95)
|
|
atomic_write_json(staging / f"{pdf.stem}.json", {
|
|
"source_sha256": fingerprint(pdf), "width_mm": MARGIN_MM,
|
|
"height_fraction": 0.25,
|
|
})
|
|
|
|
|
|
def build_grids(directory: Path, copies_dir: Path) -> list[dict]:
|
|
"""Rebuild deterministic 6-by-4 sheets, excluding deleted or changed copies."""
|
|
crops = []
|
|
for path in sorted(directory.glob("*.jpg"), key=natural_key):
|
|
pdf = copies_dir / f"{path.stem}.pdf"
|
|
try:
|
|
if pdf.is_file() and read_json(path.with_suffix('.json'))['source_sha256'] == fingerprint(pdf):
|
|
crops.append(path)
|
|
except (OSError, ValueError, KeyError, TypeError):
|
|
continue
|
|
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]
|
|
images = []
|
|
for path in batch:
|
|
with Image.open(path) as image:
|
|
images.append(image.convert('RGB'))
|
|
width = max(image.width for image in images)
|
|
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')
|
|
draw = ImageDraw.Draw(sheet)
|
|
cells = []
|
|
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):
|
|
sheet.paste(images[index], (x, y + header))
|
|
draw.text((x + 5, y + 5), f"R{row + 1} C{column + 1} | {index + 1}", fill='black')
|
|
cells.append({'cell': index + 1, 'row': row + 1, 'column': column + 1,
|
|
'copy': batch[index].stem})
|
|
else:
|
|
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})
|
|
atomic_write_json(staging / 'manifest.json', grids)
|
|
return grids
|
|
|
|
|
|
class NameCell(BaseModel):
|
|
cell: int = Field(strict=True, ge=1, le=24)
|
|
name: str
|
|
|
|
|
|
class NameResponse(BaseModel):
|
|
cells: list[NameCell]
|
|
|
|
|
|
def identify_names(workspace, groups, *, client, model: str, names_text: str, overwrite: bool,
|
|
thinking_level: str | None = "low") -> dict[str, str]:
|
|
"""Identify pending copies before label requests; safely fall back on any failure."""
|
|
directory = workspace.cutleft_dir / 'Names'
|
|
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)
|
|
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}
|
|
pending = {}
|
|
for stem, cell in selected.items():
|
|
crop_hash = fingerprint(directory / f'{stem}.jpg')
|
|
cache_path = directory / f'{stem}_result.json'
|
|
try:
|
|
cached = read_json(cache_path)
|
|
if (not overwrite and cached['crop_sha256'] == crop_hash
|
|
and cached['names_text'] == names_text
|
|
and cached['name'] in roster | {MISSING}):
|
|
result[stem] = cached['name']
|
|
continue
|
|
except (OSError, ValueError, KeyError, TypeError):
|
|
pass
|
|
pending[stem] = (cell, crop_hash, cache_path)
|
|
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,
|
|
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. Ignore class names such as MPSI 3.
|
|
Occupied cells are exactly:
|
|
{positions}
|
|
All other cells are EMPTY padding: do not return them.
|
|
Return JSON {{"cells": [{{"cell": 1, "name": "exact authorized name"}}]}}.
|
|
Return one entry per occupied cell, using its cell number. Match only confidently
|
|
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)
|
|
try:
|
|
response = client.models.generate_content(
|
|
model=model,
|
|
contents=[types.Content(role='user', parts=[
|
|
types.Part.from_bytes(data=(directory / 'Grids' / grid['image']).read_bytes(), mime_type='image/jpeg'),
|
|
types.Part.from_text(text=prompt),
|
|
])],
|
|
config=types.GenerateContentConfig(
|
|
http_options=types.HttpOptions(timeout=80_000, retry_options=types.HttpRetryOptions(attempts=1)),
|
|
automatic_function_calling=types.AutomaticFunctionCallingConfig(disable=True),
|
|
thinking_config=thinking_config(thinking_level), seed=fresh_seed(),
|
|
response_mime_type='application/json', response_json_schema=NameResponse.model_json_schema(),
|
|
),
|
|
)
|
|
entries = NameResponse.model_validate_json(response.text).cells
|
|
by_cell = {}
|
|
for entry in entries:
|
|
by_cell.setdefault(entry.cell, []).append(entry.name)
|
|
for stem, (cell, crop_hash, cache_path) in pending.items():
|
|
candidates = by_cell.get(cell['cell'], [])
|
|
name = candidates[0] if len(candidates) == 1 and candidates[0] in roster else MISSING
|
|
result[stem] = name
|
|
atomic_write_json(cache_path, {'name': name, 'crop_sha256': crop_hash, 'names_text': names_text})
|
|
except Exception as exc: # Remote failures must preserve the existing fallback.
|
|
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
|