259 lines
13 KiB
Python
259 lines
13 KiB
Python
"""First-page name crops, numbered contact sheets, and roster-checked recognition."""
|
||
from __future__ import annotations
|
||
|
||
from pathlib import Path
|
||
import math
|
||
|
||
from google.genai import types
|
||
from PIL import Image, ImageDraw
|
||
from pydantic import BaseModel, Field, field_validator
|
||
|
||
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": crop.height / page.height,
|
||
})
|
||
|
||
|
||
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):
|
||
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
|
||
box_2d: list[float] | None = Field(
|
||
default=None, description="[ymin, xmin, ymax, xmax], 0–1000 relative to the individual crop, excluding its header")
|
||
|
||
@field_validator('box_2d', mode='before')
|
||
@classmethod
|
||
def valid_box(cls, value):
|
||
# Bad localization must not discard an otherwise readable name.
|
||
if (not isinstance(value, list) or len(value) != 4
|
||
or any(type(v) not in (int, float) or not math.isfinite(v)
|
||
or not 0 <= v <= 1000 for v in value)):
|
||
return None
|
||
return value if value[0] < value[2] and value[1] < value[3] else None
|
||
|
||
|
||
class NameResponse(BaseModel):
|
||
cells: list[NameCell]
|
||
|
||
|
||
def name_bottom_height_fraction(workspace, stem: str) -> float | None:
|
||
"""Read the name/class bottom on page 1, in cutleft's 0–1 height units.
|
||
|
||
Returns None when localization is unavailable or the source/crop changed.
|
||
This is intended as a future lower boundary for margin analysis.
|
||
"""
|
||
directory = workspace.cutleft_dir / 'Names'
|
||
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", 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, 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}
|
||
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 'name_bottom_height_fraction' in cached
|
||
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'])
|
||
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.
|
||
Occupied cells are exactly:
|
||
{positions}
|
||
All other cells are EMPTY padding: do not return them.
|
||
Return JSON {{"cells": [{{"cell": 1, "name": "exact authorized name", "box_2d": [100, 50, 400, 900]}}]}}.
|
||
For each identified student, box_2d must tightly enclose the handwritten name AND
|
||
any adjacent class designation, including a class written just below the name.
|
||
Use [ymin, xmin, ymax, xmax] coordinates normalized to 0–1000 relative to that
|
||
individual page crop, NOT the whole grid. Exclude the numbered header, black
|
||
borders and any padding outside the crop. Return box_2d=null for Missing names
|
||
or if the position cannot be determined reliably.
|
||
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{' (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,
|
||
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)
|
||
for stem, (cell, crop_hash, cache_path) in pending.items():
|
||
candidates = by_cell.get(cell['cell'], [])
|
||
entry = candidates[0] if len(candidates) == 1 else None
|
||
name = entry.name if entry is not None and entry.name in roster else MISSING
|
||
box = entry.box_2d if name != MISSING else None
|
||
metadata = read_json(directory / f'{stem}.json')
|
||
bottom = box[2] / 1000 * metadata['height_fraction'] if box else None
|
||
result[stem] = name
|
||
atomic_write_json(cache_path, {
|
||
'name': name, 'crop_sha256': crop_hash, 'names_text': names_text,
|
||
'source_sha256': metadata['source_sha256'],
|
||
'box_2d': box, # 0–1000 coordinates relative to the name crop.
|
||
'name_bottom_height_fraction': bottom, # Fraction of full page 1.
|
||
})
|
||
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
|
||
|
||
|
||
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
|