Standardisation 7
This commit is contained in:
+262
-156
@@ -1,18 +1,30 @@
|
||||
from google import genai
|
||||
from google.genai import types
|
||||
import base64
|
||||
from pathlib import Path
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import List, Dict
|
||||
import sys
|
||||
import os
|
||||
import time
|
||||
import json
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import re
|
||||
import time
|
||||
import typing
|
||||
from collections import defaultdict
|
||||
from collections.abc import Callable, Sequence
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from pathlib import Path
|
||||
|
||||
from google import genai
|
||||
from google.genai import types
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
import config
|
||||
from copienator import (
|
||||
CliError,
|
||||
EvaluationWorkspace,
|
||||
ExitCode,
|
||||
atomic_write_json,
|
||||
execute,
|
||||
read_json,
|
||||
target_parser,
|
||||
workspace_from_target,
|
||||
)
|
||||
from utils import natural_key, read_all_labels
|
||||
|
||||
MODEL_ID = config.MODEL_FOR_LABEL_ID
|
||||
api_key = config.API_KEY
|
||||
@@ -125,12 +137,14 @@ Since this copy isn't the first part of a sequence, simply set the
|
||||
name to `\"Continued\"`."""
|
||||
|
||||
class BoxItem(BaseModel):
|
||||
box_2d: List[int] = Field(description="Bounding box coordinates (e.g., [ymin, xmin, ymax, xmax])")
|
||||
box_2d: list[int] = Field(description="Bounding box coordinates (e.g., [ymin, xmin, ymax, xmax])")
|
||||
label: str = Field(description="The label associated with the specific box")
|
||||
|
||||
class AnnotationData(BaseModel):
|
||||
name: str = Field(description="The name identifier")
|
||||
list: List[BoxItem] = Field(description="List of bounding box items")
|
||||
list: typing.List[BoxItem] = Field( # noqa: UP006 - field name shadows list
|
||||
description="List of bounding box items"
|
||||
)
|
||||
|
||||
|
||||
def generate_request(file, labels, names, context_labels, wrong_labels):
|
||||
@@ -176,169 +190,261 @@ def generate_request(file, labels, names, context_labels, wrong_labels):
|
||||
)
|
||||
return (contents, generate_content_config)
|
||||
|
||||
# Argument Parsing
|
||||
parser = argparse.ArgumentParser(description="Process a directory or specific files using Gemini.")
|
||||
parser.add_argument("input_paths", nargs='+', help="The input directory or specific files")
|
||||
parser.add_argument("--overwrite", action="store_true", help="Regenerate output even if it exists")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
# input_arg = Path(args.input_path)
|
||||
image_files = []
|
||||
|
||||
from utils import natural_key, read_all_labels
|
||||
|
||||
for path_str in args.input_paths:
|
||||
input_arg = Path(path_str)
|
||||
target_files = []
|
||||
|
||||
# 1. Determine which files to process
|
||||
if input_arg.is_file():
|
||||
INPUT_DIR = input_arg.parent.parent
|
||||
target_files = [input_arg]
|
||||
elif input_arg.is_dir():
|
||||
INPUT_DIR = input_arg
|
||||
COPIES_DIR = INPUT_DIR / "Copies"
|
||||
target_files = list(COPIES_DIR.glob("Copie*.pdf"))
|
||||
if not target_files:
|
||||
print(f"Warning: No Copie*.pdf files found in {input_arg}")
|
||||
else:
|
||||
print(f"Error: {input_arg} is not a valid file or directory.")
|
||||
continue
|
||||
|
||||
# 2. Run the logic for all collected files
|
||||
for target_file in target_files:
|
||||
# INPUT_DIR = target_file.parent
|
||||
CUTLEFT_DIR = INPUT_DIR / 'Cutleft'
|
||||
|
||||
# Matches stem_01.jpg, stem_02.jpg, etc.
|
||||
found_files = sorted(
|
||||
CUTLEFT_DIR.glob(f"{target_file.stem}_*.jpg"),
|
||||
key=natural_key
|
||||
)
|
||||
|
||||
if found_files:
|
||||
image_files.extend(found_files)
|
||||
else:
|
||||
print(f"Warning: No variants found for {target_file.stem} in {CUTLEFT_DIR}")
|
||||
|
||||
all_labels = read_all_labels(INPUT_DIR)
|
||||
labels_txt = "\n".join(all_labels) + "\n"
|
||||
valid_labels_set = set(all_labels)
|
||||
names_path = (INPUT_DIR / "names")
|
||||
if not os.path.exists(names_path):
|
||||
names_path = Path("names")
|
||||
names_txt = names_path.read_text()
|
||||
|
||||
valid_names_set = set(line.strip() for line in names_txt.splitlines() if line.strip())
|
||||
valid_names_set.add("Unknown")
|
||||
valid_names_set.add("Continued")
|
||||
|
||||
client = genai.Client(api_key=api_key)
|
||||
|
||||
# Group files by Copy ID (e.g. Copie01_01.jpg -> Copie01)
|
||||
# regex: match everything before the last underscore if it ends in digits
|
||||
file_groups = defaultdict(list)
|
||||
for img in image_files:
|
||||
stem = img.stem
|
||||
# match CopieXX_YY -> Group CopieXX
|
||||
match = re.match(r"(.+)_(\d+)$", stem)
|
||||
if match:
|
||||
group_key = match.group(1)
|
||||
file_groups[group_key].append(img)
|
||||
else:
|
||||
# Fallback for files without underscore numbering
|
||||
file_groups[stem].append(img)
|
||||
|
||||
# Sort files within each group to ensure sequential processing
|
||||
for key in file_groups:
|
||||
file_groups[key].sort(key=lambda x: x.name)
|
||||
|
||||
TARGET_INTERVAL = 3.5
|
||||
Sleep = Callable[[float], None]
|
||||
|
||||
def process_copy_group(group_key, files):
|
||||
"""Processes a list of files belonging to one copy sequentially to maintain context."""
|
||||
|
||||
# Context accumulator for this specific copy
|
||||
accumulated_labels = []
|
||||
def selected_images(
|
||||
workspace: EvaluationWorkspace,
|
||||
targets: list[Path],
|
||||
) -> tuple[list[Path], list[str]]:
|
||||
"""Resolve evaluation, copy-PDF, or Cutleft-image targets."""
|
||||
workspace.require_directories("Copies", "Cutleft")
|
||||
images: list[Path] = []
|
||||
warnings: list[str] = []
|
||||
for target in targets:
|
||||
if target.is_dir():
|
||||
copy_pdfs = sorted(
|
||||
workspace.copies_dir.glob("Copie*.pdf"), key=natural_key
|
||||
)
|
||||
if not copy_pdfs:
|
||||
warnings.append(f"No Copie*.pdf files found in {workspace.copies_dir}")
|
||||
stems = [path.stem for path in copy_pdfs]
|
||||
elif target.suffix.casefold() in {".jpg", ".jpeg"}:
|
||||
if target.parent != workspace.cutleft_dir:
|
||||
raise CliError(
|
||||
f"Image target is not in {workspace.cutleft_dir}: {target}",
|
||||
ExitCode.INVALID_ARGUMENTS,
|
||||
)
|
||||
images.append(target)
|
||||
continue
|
||||
elif target.suffix.casefold() == ".pdf":
|
||||
stems = [target.stem]
|
||||
else:
|
||||
raise CliError(
|
||||
f"Unsupported target for label detection: {target}",
|
||||
ExitCode.INVALID_ARGUMENTS,
|
||||
)
|
||||
for stem in stems:
|
||||
found = sorted(
|
||||
workspace.cutleft_dir.glob(f"{stem}_*.jpg"), key=natural_key
|
||||
)
|
||||
if found:
|
||||
images.extend(found)
|
||||
else:
|
||||
warnings.append(
|
||||
f"No Cutleft image variants found for {stem} in "
|
||||
f"{workspace.cutleft_dir}"
|
||||
)
|
||||
return list(dict.fromkeys(images)), warnings
|
||||
|
||||
|
||||
def group_images(image_files: list[Path]) -> dict[str, list[Path]]:
|
||||
groups: defaultdict[str, list[Path]] = defaultdict(list)
|
||||
for image in image_files:
|
||||
match = re.match(r"(.+)_(\d+)$", image.stem)
|
||||
groups[match.group(1) if match else image.stem].append(image)
|
||||
for files in groups.values():
|
||||
files.sort(key=natural_key)
|
||||
return dict(groups)
|
||||
|
||||
|
||||
def _existing_context(output_json: Path) -> list[str]:
|
||||
try:
|
||||
loaded = read_json(output_json)
|
||||
if not isinstance(loaded, dict):
|
||||
return []
|
||||
return [
|
||||
str(item["label"])
|
||||
for item in loaded.get("list", [])
|
||||
if isinstance(item, dict) and "label" in item
|
||||
]
|
||||
except (OSError, TypeError, ValueError):
|
||||
return []
|
||||
|
||||
|
||||
def process_copy_group(
|
||||
workspace: EvaluationWorkspace,
|
||||
group_key: str,
|
||||
files: list[Path],
|
||||
*,
|
||||
client,
|
||||
labels_text: str,
|
||||
names_text: str,
|
||||
valid_labels: set[str],
|
||||
valid_names: set[str],
|
||||
overwrite: bool,
|
||||
sleep: Sleep = time.sleep,
|
||||
target_interval: float = TARGET_INTERVAL,
|
||||
) -> int:
|
||||
"""Process one student's image parts sequentially to preserve context."""
|
||||
accumulated_labels: list[str] = []
|
||||
generated = 0
|
||||
for image_file in files:
|
||||
start_time = time.time()
|
||||
base_name = image_file.stem
|
||||
output_json = INPUT_DIR / "Copies" / f"{base_name}.json"
|
||||
|
||||
# Check existing
|
||||
if output_json.exists() and not args.overwrite:
|
||||
started = time.monotonic()
|
||||
output_json = workspace.copies_dir / f"{image_file.stem}.json"
|
||||
if output_json.exists() and not overwrite:
|
||||
print(f"[{group_key}] Skipping {image_file.name}, output exists.")
|
||||
# If skipping, we should try to load existing labels to keep context for next parts
|
||||
try:
|
||||
with open(output_json, 'r') as f:
|
||||
data = json.load(f)
|
||||
for item in data.get('list', []):
|
||||
accumulated_labels.append(item['label'])
|
||||
except:
|
||||
pass # If read fails, next part has no context
|
||||
accumulated_labels.extend(_existing_context(output_json))
|
||||
continue
|
||||
|
||||
print(f"[{group_key}] Processing {image_file.name} with {len(accumulated_labels)} accumulated labels...")
|
||||
|
||||
attempt = -1
|
||||
wrong_labels = []
|
||||
print(
|
||||
f"[{group_key}] Processing {image_file.name} with "
|
||||
f"{len(accumulated_labels)} accumulated labels..."
|
||||
)
|
||||
attempt = 0
|
||||
wrong_labels: list[str] = []
|
||||
while True:
|
||||
attempt += 1
|
||||
if attempt > 0:
|
||||
time.sleep(10 * attempt)
|
||||
sleep(10 * attempt)
|
||||
try:
|
||||
contents, config = generate_request(image_file, labels_txt, names_txt, accumulated_labels,
|
||||
wrong_labels)
|
||||
|
||||
contents, request_config = generate_request(
|
||||
image_file,
|
||||
labels_text,
|
||||
names_text,
|
||||
accumulated_labels,
|
||||
wrong_labels,
|
||||
)
|
||||
response = client.models.generate_content(
|
||||
model=MODEL_ID,
|
||||
contents=contents,
|
||||
config=config
|
||||
config=request_config,
|
||||
)
|
||||
|
||||
annota = AnnotationData.model_validate_json(response.text)
|
||||
unknown = [item.label for item in annota.list if item.label not in valid_labels_set]
|
||||
name = annota.name
|
||||
annotation = AnnotationData.model_validate_json(response.text)
|
||||
unknown = [
|
||||
item.label
|
||||
for item in annotation.list
|
||||
if item.label not in valid_labels
|
||||
]
|
||||
if unknown:
|
||||
print(f"Error: {image_file.name} contained unknown labels: {unknown}")
|
||||
print(
|
||||
f"Error: {image_file.name} contained unknown labels: "
|
||||
f"{unknown}"
|
||||
)
|
||||
wrong_labels.extend(unknown)
|
||||
print("Retrying request...")
|
||||
continue # Retry immediately
|
||||
|
||||
if name not in valid_names_set:
|
||||
print(f"Error: {image_file.name} returned unknown name : {name}")
|
||||
attempt += 1
|
||||
continue
|
||||
if annotation.name not in valid_names:
|
||||
print(
|
||||
f"Error: {image_file.name} returned unknown name: "
|
||||
f"{annotation.name}"
|
||||
)
|
||||
if attempt == 0:
|
||||
print("Retrying request...")
|
||||
continue # Retry immediately
|
||||
else:
|
||||
name = "Unknown"
|
||||
annota.name = name
|
||||
# Save result
|
||||
with open(output_json, "w", encoding="utf-8") as f:
|
||||
json.dump(annota.model_dump(), f, indent=2)
|
||||
attempt += 1
|
||||
continue
|
||||
annotation.name = "Unknown"
|
||||
|
||||
# Update context for the next part in this group
|
||||
for box in annota.list:
|
||||
accumulated_labels.append(box.label)
|
||||
break # exit retry loop
|
||||
except Exception as e:
|
||||
print(f"Error processing {image_file.name}: {e}\n\tIt will be retried.")
|
||||
atomic_write_json(output_json, annotation.model_dump())
|
||||
accumulated_labels.extend(box.label for box in annotation.list)
|
||||
generated += 1
|
||||
break
|
||||
except KeyboardInterrupt:
|
||||
raise
|
||||
except Exception as exc: # noqa: BLE001 - remote API retry boundary
|
||||
print(
|
||||
f"Error processing {image_file.name}: {exc}\n"
|
||||
"\tIt will be retried."
|
||||
)
|
||||
attempt += 1
|
||||
sleep(max(0.0, target_interval - (time.monotonic() - started)))
|
||||
return generated
|
||||
|
||||
# Rate Limiting
|
||||
elapsed = time.time() - start_time
|
||||
time.sleep(max(0, TARGET_INTERVAL - elapsed))
|
||||
|
||||
# Run ThreadPool on GROUPS (Copies), not individual files
|
||||
# Each thread handles one student's full exam copy sequentially
|
||||
with ThreadPoolExecutor(max_workers=12) as executor:
|
||||
# Convert dict items to arguments for map
|
||||
# executor.map expects a function and an iterable.
|
||||
# We use a lambda or separate function to unpack the tuple if needed,
|
||||
# but here we'll just submit futures.
|
||||
futures = [executor.submit(process_copy_group, k, v) for k, v in file_groups.items()]
|
||||
def run(
|
||||
workspace: EvaluationWorkspace,
|
||||
targets: list[Path],
|
||||
*,
|
||||
overwrite: bool = False,
|
||||
client=None,
|
||||
sleep: Sleep = time.sleep,
|
||||
max_workers: int = 12,
|
||||
) -> ExitCode:
|
||||
workspace.require_files("labels")
|
||||
images, warnings = selected_images(workspace, targets)
|
||||
for warning in warnings:
|
||||
print(f"Warning: {warning}")
|
||||
if not images:
|
||||
return ExitCode.PARTIAL
|
||||
|
||||
# Wait for all to complete
|
||||
for future in futures:
|
||||
future.result()
|
||||
all_labels = read_all_labels(workspace.root)
|
||||
labels_text = "\n".join(all_labels) + "\n"
|
||||
names_path = workspace.names_file()
|
||||
if not names_path.is_file():
|
||||
raise CliError(f"Names file not found: {names_path}", ExitCode.INVALID_WORKSPACE)
|
||||
names_text = names_path.read_text(encoding="utf-8")
|
||||
valid_names = {
|
||||
line.strip() for line in names_text.splitlines() if line.strip()
|
||||
} | {"Unknown", "Continued"}
|
||||
if client is None:
|
||||
client = genai.Client(api_key=api_key)
|
||||
|
||||
groups = group_images(images)
|
||||
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
||||
futures = [
|
||||
executor.submit(
|
||||
process_copy_group,
|
||||
workspace,
|
||||
group_key,
|
||||
files,
|
||||
client=client,
|
||||
labels_text=labels_text,
|
||||
names_text=names_text,
|
||||
valid_labels=set(all_labels),
|
||||
valid_names=valid_names,
|
||||
overwrite=overwrite,
|
||||
sleep=sleep,
|
||||
)
|
||||
for group_key, files in groups.items()
|
||||
]
|
||||
for future in futures:
|
||||
future.result()
|
||||
return ExitCode.PARTIAL if warnings else ExitCode.SUCCESS
|
||||
|
||||
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
parser = target_parser("Detect handwritten question labels with Gemini")
|
||||
parser.add_argument(
|
||||
"additional_targets",
|
||||
nargs="*",
|
||||
type=Path,
|
||||
help="Additional copy PDFs or Cutleft images from the same evaluation",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--overwrite",
|
||||
action="store_true",
|
||||
help="Regenerate JSON outputs that already exist",
|
||||
)
|
||||
return parser
|
||||
|
||||
|
||||
def main(argv: Sequence[str] | None = None) -> int:
|
||||
parser = build_parser()
|
||||
|
||||
def handle(args: argparse.Namespace) -> ExitCode:
|
||||
workspace, target = workspace_from_target(
|
||||
args, repository=Path(__file__).resolve().parent
|
||||
)
|
||||
targets = [target]
|
||||
for additional in args.additional_targets:
|
||||
resolved = additional.expanduser().resolve()
|
||||
if not resolved.exists():
|
||||
raise CliError(
|
||||
f"Target does not exist: {resolved}",
|
||||
ExitCode.INVALID_WORKSPACE,
|
||||
)
|
||||
additional_workspace = EvaluationWorkspace.discover(
|
||||
resolved, repository=workspace.repository
|
||||
)
|
||||
if additional_workspace.root != workspace.root:
|
||||
raise CliError(
|
||||
"All targets must belong to the same evaluation",
|
||||
ExitCode.INVALID_ARGUMENTS,
|
||||
)
|
||||
targets.append(resolved)
|
||||
return run(workspace, targets, overwrite=args.overwrite)
|
||||
|
||||
return execute(parser, argv, handle)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
|
||||
Reference in New Issue
Block a user