Compare commits
5
Commits
bcba5facc8
...
3a8d0fe3ff
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
3a8d0fe3ff | ||
|
|
b19d3b0db6 | ||
|
|
18d1e5e2bb | ||
|
|
aa40e58dd1 | ||
|
|
644e287586 |
+56
@@ -171,6 +171,12 @@ scripts migrés vers cette convention sont actuellement :
|
||||
|
||||
- =copies_tools.py=, =grouping.py= et =verify_groups.py= ;
|
||||
- =post-correction.py= et =resolve_manual.py= ;
|
||||
- =page_splitter.py=, =cutleft.py=, =plotting.py= et
|
||||
=splitting_int.py= ;
|
||||
- =gemini_for_labels.py= ;
|
||||
- =gemini_for_enonce.py= et =enonce_info.py= ;
|
||||
- =correction.py=, =submit_batches.py=, =batch_status.py= et
|
||||
=fetch_batched_results.py= ;
|
||||
- =annotating.py=, =annotating_with_checks.py= et
|
||||
=annotating_by_label.py= ;
|
||||
- =reading_annotations.py= et =reading_grouped_annotations.py= ;
|
||||
@@ -230,6 +236,12 @@ Dans le dossier de l'évaluation, mettre les fichiers suivants de l'évaluation
|
||||
|
||||
- Alternative personnelle : `python enonce_info.py Interro`
|
||||
|
||||
Ces deux commandes suivent la convention des scripts standardisés.
|
||||
Leur import ne lance aucun traitement et les erreurs partielles sont
|
||||
distinguées des échecs. Les réponses d'extraction mises en cache par
|
||||
=gemini_for_enonce.py= et le fichier =labels= sont publiés
|
||||
atomiquement.
|
||||
|
||||
** Prétraitement des copies
|
||||
|
||||
Mettre les copies scannées au format pdf dans =Interro=.
|
||||
@@ -251,6 +263,12 @@ Mettre les copies scannées au format pdf dans =Interro=.
|
||||
+ de déplacer la délimitation à droite/gauche
|
||||
|
||||
Fix issues with =python page_splitter.py Interro14/Copies/Copie01.pdf=
|
||||
|
||||
Le PDF transformé est construit dans un dossier temporaire. La
|
||||
copie produite et la sauvegarde dans =Copies Originales= sont
|
||||
ensuite installées avec rollback : une erreur conserve les deux
|
||||
versions précédentes. Une relance ciblée lit directement la
|
||||
sauvegarde originale sans la déplacer au préalable.
|
||||
4. =python cutleft.py Interro=
|
||||
|
||||
Découpe la partie gauche des copies, là où il devrait y avoir les
|
||||
@@ -260,6 +278,10 @@ Mettre les copies scannées au format pdf dans =Interro=.
|
||||
|
||||
Rerun on a single file with =python cutleft.py Interro/Copies/Copie01.pdf=
|
||||
|
||||
Les images et le fichier =_schema.json= d'une copie sont remplacés
|
||||
ensemble. Fermer l'outil juste après la dernière validation ne peut
|
||||
donc plus interrompre un thread de sauvegarde en arrière-plan.
|
||||
|
||||
** Labelisation et regroupement
|
||||
|
||||
Optional : Set proxy with ~export HTTPS_PROXY="http://10.0.0.1:3128"~
|
||||
@@ -268,6 +290,14 @@ Optional : Set proxy with ~export HTTPS_PROXY="http://10.0.0.1:3128"~
|
||||
|
||||
Fait des requêtes à Gemini pour identifier les labels des
|
||||
questions dans images générées à partir des parties gauches des copies.
|
||||
|
||||
Une copie PDF ou une image précise de =Cutleft= peut également être
|
||||
ciblée. Plusieurs cibles de la même évaluation sont acceptées. Les
|
||||
parties d'une copie restent traitées séquentiellement afin de
|
||||
conserver les labels précédents comme contexte, tandis que les
|
||||
copies différentes sont traitées en parallèle. Chaque réponse JSON
|
||||
validée est écrite atomiquement. Une cible sans image correspondante
|
||||
produit le code de sortie 4.
|
||||
2. =python plotting.py Interro=
|
||||
|
||||
Permet de vérifier visuellement les labels trouvés.
|
||||
@@ -282,6 +312,10 @@ Optional : Set proxy with ~export HTTPS_PROXY="http://10.0.0.1:3128"~
|
||||
Pour modifier une seule copie :
|
||||
=python plotting.py Interro/Copies/Copie01.pdf=
|
||||
|
||||
Les coordonnées agrégées sont écrites atomiquement dans le JSON de
|
||||
la copie. Fermer la fenêtre avant la fin d'une copie ne remplace pas
|
||||
son JSON par un résultat incomplet.
|
||||
|
||||
It also generates les =Copie01.json=, à partir des =Copie01_01.json=
|
||||
En cas de soucis, (par exemple les pages ne sont pas dans le bon ordre)
|
||||
- Réordonner les pages du fichier pdf
|
||||
@@ -291,6 +325,11 @@ Optional : Set proxy with ~export HTTPS_PROXY="http://10.0.0.1:3128"~
|
||||
|
||||
Découpe les copies suivant les exercices
|
||||
Peut-être appelé avec une seule copie.
|
||||
|
||||
Les réponses d'une copie sont préparées dans un dossier temporaire,
|
||||
puis remplacent ensemble le dossier précédent. En cas d'erreur,
|
||||
l'ancienne version est conservée. Les réponses devenues obsolètes
|
||||
restent archivées dans le sous-dossier =Missing=.
|
||||
4. =python grouping.py Interro=
|
||||
|
||||
Regroupe les mêmes questions de différentes copies en groupes de
|
||||
@@ -313,6 +352,13 @@ Optional : Set proxy with ~export HTTPS_PROXY="http://10.0.0.1:3128"~
|
||||
|
||||
Fais les requêtes de correction à Gemini.
|
||||
|
||||
=correction.py= peut être relancé sans supprimer son état. Les
|
||||
fichiers =correction.json= et =correction_progress.json= sont mis à
|
||||
jour atomiquement. Avec =--overwrite=, leur version précédente reste
|
||||
en place jusqu'à la première écriture réussie de la nouvelle
|
||||
exécution. =--reset= est la seule option qui supprime explicitement
|
||||
cet état et restaure les fichiers =*_old.pdf=.
|
||||
|
||||
L'argument =limit= limite le nombre de requêtes à Gemini Pro
|
||||
(chères), pour une version low cost, passer =--limit 0=, toutes
|
||||
les requêtes seront sur Gemini Flash.
|
||||
@@ -327,6 +373,16 @@ Optional : Set proxy with ~export HTTPS_PROXY="http://10.0.0.1:3128"~
|
||||
+ =python batch_status.py=
|
||||
+ =python fetch_batched_results.py Interro=
|
||||
+ =python correction.py Interro --deal-with-batched=
|
||||
|
||||
Les quatre commandes de ce flux suivent la convention des scripts
|
||||
standardisés. Les fichiers de requêtes et le résultat JSONL combiné
|
||||
sont publiés atomiquement : une interruption ne laisse pas de fichier
|
||||
final partiellement écrit. =submit_batches.py= conserve aussi les
|
||||
identifiants distants dans =batch_jobs.json= ; la récupération les
|
||||
utilise en priorité et garde la recherche par nom pour les anciens
|
||||
batchs. =batch_status.py --download JOB --output
|
||||
resultat.jsonl= permet aussi de télécharger atomiquement le résultat
|
||||
d'un job particulier.
|
||||
3. =python post-correction.py Interro=
|
||||
|
||||
- Essaye de corriger des erreurs d'encodage/d'accents dans
|
||||
|
||||
+79
-74
@@ -1,88 +1,93 @@
|
||||
import os
|
||||
import sys
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
from collections.abc import Sequence
|
||||
from pathlib import Path
|
||||
|
||||
from google import genai
|
||||
|
||||
if "GEMINI_API_KEY" not in os.environ:
|
||||
sys.exit("Error: GEMINI_API_KEY environment variable not set.")
|
||||
|
||||
client = genai.Client()
|
||||
|
||||
def list_jobs():
|
||||
print("Fetching recent batch jobs...\n")
|
||||
try:
|
||||
batch_jobs = client.batches.list()
|
||||
jobs_found = False
|
||||
|
||||
for job in batch_jobs:
|
||||
jobs_found = True
|
||||
state = job.state.name if hasattr(job.state, 'name') else job.state
|
||||
|
||||
print("-" * 60)
|
||||
print(f"Job Name: {job.name}")
|
||||
|
||||
if hasattr(job, 'display_name') and job.display_name:
|
||||
print(f"Display Name: {job.display_name}")
|
||||
|
||||
print(f"State: {state}")
|
||||
|
||||
if state == 'JOB_STATE_FAILED' and hasattr(job, 'error'):
|
||||
print(f"Error: {job.error}")
|
||||
|
||||
if state == 'JOB_STATE_SUCCEEDED' and hasattr(job, 'dest') and job.dest:
|
||||
if hasattr(job.dest, 'file_name') and job.dest.file_name:
|
||||
print(f"Output File: {job.dest.file_name}")
|
||||
|
||||
if not jobs_found:
|
||||
print("No batch jobs found.")
|
||||
else:
|
||||
print("-" * 60)
|
||||
print("\nTo download a completed job, run:")
|
||||
print("python batch_status.py --download batches/<YOUR_BATCH_ID>")
|
||||
|
||||
except Exception as e:
|
||||
sys.exit(f"An error occurred while listing jobs: {e}")
|
||||
import config
|
||||
from copienator import (
|
||||
CliError,
|
||||
ExitCode,
|
||||
atomic_write_bytes,
|
||||
execute,
|
||||
standard_parser,
|
||||
)
|
||||
|
||||
|
||||
def download_job(job_name):
|
||||
print(f"Checking status for {job_name}...\n")
|
||||
try:
|
||||
job = client.batches.get(name=job_name)
|
||||
state = job.state.name if hasattr(job.state, 'name') else job.state
|
||||
def _client():
|
||||
if not config.API_KEY:
|
||||
raise CliError("GEMINI_API_KEY is not configured")
|
||||
return genai.Client(api_key=config.API_KEY)
|
||||
|
||||
print(f"State: {state}")
|
||||
|
||||
if state != 'JOB_STATE_SUCCEEDED':
|
||||
print("Job is not ready yet or has failed.")
|
||||
if state == 'JOB_STATE_FAILED' and hasattr(job, 'error'):
|
||||
print(f"Error: {job.error}")
|
||||
return
|
||||
def list_jobs(*, client=None) -> ExitCode:
|
||||
client = client or _client()
|
||||
print("Fetching recent batch jobs...")
|
||||
jobs = list(client.batches.list())
|
||||
for job in jobs:
|
||||
state = job.state.name if hasattr(job.state, "name") else job.state
|
||||
print(f"{job.name}: {state}")
|
||||
if getattr(job, "display_name", None):
|
||||
print(f" Display name: {job.display_name}")
|
||||
if state == "JOB_STATE_FAILED" and getattr(job, "error", None):
|
||||
print(f" Error: {job.error}")
|
||||
destination = getattr(job, "dest", None)
|
||||
if state == "JOB_STATE_SUCCEEDED" and getattr(
|
||||
destination, "file_name", None
|
||||
):
|
||||
print(f" Output file: {destination.file_name}")
|
||||
if not jobs:
|
||||
print("No batch jobs found.")
|
||||
return ExitCode.SUCCESS
|
||||
|
||||
if hasattr(job, 'dest') and job.dest and hasattr(job.dest, 'file_name') and job.dest.file_name:
|
||||
result_file_name = job.dest.file_name
|
||||
print(f"Downloading results from {result_file_name}...")
|
||||
|
||||
file_content_bytes = client.files.download(file=result_file_name)
|
||||
output_path = f"results_{job_name.replace('/', '_')}.jsonl"
|
||||
def download_job(
|
||||
job_name: str,
|
||||
*,
|
||||
output: Path | None = None,
|
||||
client=None,
|
||||
) -> ExitCode:
|
||||
client = client or _client()
|
||||
job = client.batches.get(name=job_name)
|
||||
state = job.state.name if hasattr(job.state, "name") else job.state
|
||||
print(f"State: {state}")
|
||||
if state != "JOB_STATE_SUCCEEDED":
|
||||
if state == "JOB_STATE_FAILED" and getattr(job, "error", None):
|
||||
print(f"Error: {job.error}")
|
||||
return ExitCode.PARTIAL
|
||||
destination = getattr(job, "dest", None)
|
||||
file_name = getattr(destination, "file_name", None)
|
||||
if not file_name:
|
||||
print("Job succeeded but no output file was found.")
|
||||
return ExitCode.PARTIAL
|
||||
payload = client.files.download(file=file_name)
|
||||
output_path = output or Path(f"results_{job_name.replace('/', '_')}.jsonl")
|
||||
atomic_write_bytes(output_path, payload)
|
||||
print(f"Saved batch results to {output_path}")
|
||||
return ExitCode.SUCCESS
|
||||
|
||||
with open(output_path, "wb") as f:
|
||||
f.write(file_content_bytes)
|
||||
|
||||
print(f"Success! Saved to {output_path}")
|
||||
print(f"You can now feed this to your correction script using: --deal-with-batched {output_path}")
|
||||
else:
|
||||
print("Job succeeded but no output file was found.")
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
parser = standard_parser("List or download Gemini correction batch jobs")
|
||||
parser.add_argument("--download", metavar="JOB_NAME")
|
||||
parser.add_argument("--output", type=Path, help="Downloaded JSONL destination")
|
||||
return parser
|
||||
|
||||
|
||||
def main(argv: Sequence[str] | None = None) -> int:
|
||||
parser = build_parser()
|
||||
|
||||
def handle(args: argparse.Namespace) -> ExitCode:
|
||||
if args.output is not None and not args.download:
|
||||
raise CliError("--output requires --download", ExitCode.INVALID_ARGUMENTS)
|
||||
if args.download:
|
||||
return download_job(args.download, output=args.output)
|
||||
return list_jobs()
|
||||
|
||||
return execute(parser, argv, handle)
|
||||
|
||||
except Exception as e:
|
||||
sys.exit(f"An error occurred while fetching the job: {e}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description="Manage Gemini Batch Jobs")
|
||||
parser.add_argument("--download", type=str, metavar="JOB_NAME",
|
||||
help="Download the results for a specific batch job (e.g. batches/123456)")
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.download:
|
||||
download_job(args.download)
|
||||
else:
|
||||
list_jobs()
|
||||
raise SystemExit(main())
|
||||
|
||||
@@ -14,6 +14,7 @@ from .cli import (
|
||||
from .json_io import (
|
||||
JsonLockTimeout,
|
||||
atomic_update_json,
|
||||
atomic_write_bytes,
|
||||
atomic_write_json,
|
||||
atomic_write_text,
|
||||
read_json,
|
||||
@@ -32,6 +33,7 @@ __all__ = [
|
||||
"WorkspaceNotFoundError",
|
||||
"WorkspaceValidationError",
|
||||
"atomic_update_json",
|
||||
"atomic_write_bytes",
|
||||
"atomic_write_json",
|
||||
"atomic_write_text",
|
||||
"evaluation_parser",
|
||||
|
||||
@@ -2,6 +2,7 @@ from __future__ import annotations
|
||||
|
||||
import shutil
|
||||
import uuid
|
||||
from collections.abc import Iterable
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
|
||||
@@ -44,7 +45,11 @@ def staged_directory(destination: str | Path):
|
||||
|
||||
|
||||
@contextmanager
|
||||
def staged_files(destination: str | Path):
|
||||
def staged_files(
|
||||
destination: str | Path,
|
||||
*,
|
||||
remove: Iterable[str] = (),
|
||||
):
|
||||
"""Stage a set of files and merge them into a directory with rollback."""
|
||||
target = Path(destination)
|
||||
target.parent.mkdir(parents=True, exist_ok=True)
|
||||
@@ -56,9 +61,15 @@ def staged_files(destination: str | Path):
|
||||
try:
|
||||
yield staging
|
||||
staged = sorted(path for path in staging.iterdir() if path.is_file())
|
||||
staged_names = {path.name for path in staged}
|
||||
removed_names = set(remove) - staged_names
|
||||
target.mkdir(parents=True, exist_ok=True)
|
||||
backup.mkdir()
|
||||
try:
|
||||
for name in sorted(removed_names):
|
||||
destination_path = target / name
|
||||
if destination_path.is_file() or destination_path.is_symlink():
|
||||
destination_path.replace(backup / name)
|
||||
for source in staged:
|
||||
destination_path = target / source.name
|
||||
if destination_path.exists() or destination_path.is_symlink():
|
||||
|
||||
@@ -78,6 +78,10 @@ def atomic_write_text(
|
||||
_atomic_write(Path(path), text.encode(encoding))
|
||||
|
||||
|
||||
def atomic_write_bytes(path: str | Path, payload: bytes) -> None:
|
||||
_atomic_write(Path(path), payload)
|
||||
|
||||
|
||||
def atomic_write_json(
|
||||
path: str | Path,
|
||||
value: JsonValue,
|
||||
|
||||
@@ -101,6 +101,14 @@ class EvaluationWorkspace:
|
||||
def correction_progress_file(self) -> Path:
|
||||
return self.root / "correction_progress.json"
|
||||
|
||||
@property
|
||||
def batch_jobs_file(self) -> Path:
|
||||
return self.root / "batch_jobs.json"
|
||||
|
||||
@property
|
||||
def batched_correction_result_file(self) -> Path:
|
||||
return self.root / "batched_correction_result.jsonl"
|
||||
|
||||
@property
|
||||
def manual_resolutions_file(self) -> Path:
|
||||
return self.root / "manual_resolutions.txt"
|
||||
|
||||
+336
-208
@@ -1,91 +1,59 @@
|
||||
import sys
|
||||
import os
|
||||
import time
|
||||
from pathlib import Path
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import prompting
|
||||
import signal
|
||||
from google import genai
|
||||
import base64
|
||||
import shlex
|
||||
import json
|
||||
import threading
|
||||
import concurrent.futures
|
||||
import json
|
||||
import os
|
||||
import shlex
|
||||
import shutil
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
from collections.abc import Sequence
|
||||
from pathlib import Path
|
||||
|
||||
from copienator import atomic_write_json
|
||||
|
||||
if len(sys.argv) < 2:
|
||||
sys.exit("Usage: python script.py 'InterroTest/Ex 2/Group_1.jpg' OR <InputDir> OR 'file1' 'file2'")
|
||||
|
||||
# Parse Arguments
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("paths", nargs="+", help="List of images or directories")
|
||||
parser.add_argument("--overwrite", action="store_true",
|
||||
help="Force redo requests even if output exists")
|
||||
parser.add_argument("--limit", type=int, help="limit calls to gemini rpo integer")
|
||||
parser.add_argument("--refaire", action="store_true",
|
||||
help="Redo specific copies/labels defined in refaire.json")
|
||||
parser.add_argument("--batch", action="store_true",
|
||||
help="Generate a JSONL file of requests to send to the Gemini Batch API")
|
||||
parser.add_argument("--batch-from", type=str, metavar="LABEL",
|
||||
help="Do live requests before LABEL, and batch requests from LABEL onwards")
|
||||
parser.add_argument("--deal-with-batched", action="store_true",
|
||||
help="Process a JSONL file containing completed batch results")
|
||||
parser.add_argument("--reset", action="store_true",
|
||||
help="Remove correction.json, revert _old.pdf, delete _new.pdf, then exit")
|
||||
args, _ = parser.parse_known_args()
|
||||
|
||||
tasks = [] # List of tuples: (filepath_str, label_str)
|
||||
results = {}
|
||||
|
||||
|
||||
for path_str in args.paths:
|
||||
arg_path = Path(path_str)
|
||||
|
||||
if not arg_path.exists():
|
||||
print(f"Warning: {path_str} not found. Skipping.")
|
||||
continue
|
||||
|
||||
if arg_path.is_file() and arg_path.suffix.lower() == ".jpg":
|
||||
# Handle individual file
|
||||
# Note: assumes structure InterroTest/Ex 2/Group_1.jpg
|
||||
label = arg_path.parent.name
|
||||
INPUT_DIR = arg_path.parent.parent.parent
|
||||
COPIES_DIR = INPUT_DIR / "Copies"
|
||||
GROUPS_DIR = INPUT_DIR / "Par label"
|
||||
tasks.append((str(arg_path), label))
|
||||
if label not in results:
|
||||
results[label] = []
|
||||
|
||||
elif arg_path.is_dir():
|
||||
INPUT_DIR = arg_path
|
||||
COPIES_DIR = INPUT_DIR / "Copies"
|
||||
GROUPS_DIR = INPUT_DIR / "Par label"
|
||||
# Handle directory (original behavior)
|
||||
for sub in GROUPS_DIR.iterdir():
|
||||
if sub.is_dir():
|
||||
label = sub.name
|
||||
if label not in results:
|
||||
results[label] = []
|
||||
for img in sub.glob("*.jpg"):
|
||||
tasks.append((str(img), label))
|
||||
from google import genai
|
||||
|
||||
import config
|
||||
import grouping
|
||||
import prompting
|
||||
from copienator import (
|
||||
CliError,
|
||||
EvaluationWorkspace,
|
||||
ExitCode,
|
||||
atomic_write_json,
|
||||
atomic_write_text,
|
||||
execute,
|
||||
read_json,
|
||||
target_parser,
|
||||
workspace_from_target,
|
||||
)
|
||||
from utils import enonce_total, read_all_labels
|
||||
|
||||
NB_THREADS = 12
|
||||
|
||||
# PROXY_URL = "http://192.168.241.1:3128"
|
||||
PROXY_URL = None
|
||||
|
||||
if PROXY_URL:
|
||||
os.environ["http_proxy"] = PROXY_URL
|
||||
os.environ["https_proxy"] = PROXY_URL
|
||||
|
||||
import config
|
||||
|
||||
MODEL_ID_pro = config.MODEL_PRO_ID
|
||||
MODEL_ID_flash = config.MODEL_FLASH_ID
|
||||
api_key = config.API_KEY
|
||||
|
||||
# Runtime globals retained while the processing helpers are migrated incrementally.
|
||||
INPUT_DIR = Path()
|
||||
COPIES_DIR = Path()
|
||||
GROUPS_DIR = Path()
|
||||
output_path = Path()
|
||||
progress_path = Path()
|
||||
tasks: list[tuple] = []
|
||||
tasks_to_process: list[tuple] = []
|
||||
results: dict = {}
|
||||
completed_tasks: list = []
|
||||
errors_summary: list = []
|
||||
overwrite = False
|
||||
limit = None
|
||||
client = None
|
||||
start_time = 0.0
|
||||
|
||||
# --- Thread-safe Logging ---
|
||||
log_lock = threading.Lock()
|
||||
thread_logs = {}
|
||||
@@ -113,57 +81,6 @@ def flush_thread_log(tid=None):
|
||||
f.write("\n".join(thread_logs[tid]) + "\n\n")
|
||||
thread_logs[tid].clear()
|
||||
|
||||
def handle_interrupt(sig, frame):
|
||||
"""Flush all partial/unfinished logs if program is interrupted."""
|
||||
print("\nInterrupt received. Flushing partial logs...", file=sys.stderr)
|
||||
for tid in list(thread_logs.keys()):
|
||||
flush_thread_log(tid)
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
signal.signal(signal.SIGINT, handle_interrupt)
|
||||
signal.signal(signal.SIGTERM, handle_interrupt)
|
||||
# ---------------------------
|
||||
|
||||
client = genai.Client(api_key=api_key)
|
||||
output_path = INPUT_DIR / "correction.json"
|
||||
progress_path = INPUT_DIR / "correction_progress.json"
|
||||
|
||||
if args.reset:
|
||||
print("--- Running Reset ---")
|
||||
if output_path.exists():
|
||||
output_path.unlink()
|
||||
print(f"Deleted: {output_path}")
|
||||
if progress_path.exists():
|
||||
progress_path.unlink()
|
||||
print(f"Deleted: {progress_path}")
|
||||
|
||||
if COPIES_DIR.exists():
|
||||
for copie_dir in COPIES_DIR.iterdir():
|
||||
if not copie_dir.is_dir():
|
||||
continue
|
||||
|
||||
# Revert _old.pdf files
|
||||
for old_pdf in copie_dir.glob("*_old.pdf"):
|
||||
orig_pdf = copie_dir / old_pdf.name.replace("_old.pdf", ".pdf")
|
||||
if orig_pdf.exists():
|
||||
orig_pdf.unlink() # Prevent FileExistsError on Windows
|
||||
old_pdf.rename(orig_pdf)
|
||||
print(f"Moved: {copie_dir.name}/{old_pdf.name} -> {orig_pdf.name}")
|
||||
|
||||
# Delete _new.pdf files
|
||||
for new_pdf in copie_dir.glob("*_new.pdf"):
|
||||
new_pdf.unlink()
|
||||
print(f"Deleted: {copie_dir.name}/{new_pdf.name}")
|
||||
|
||||
sys.exit("Reset almost complete. For each deleted `_new`, you should manually delete the group in `Par label`")
|
||||
|
||||
start_time = time.time()
|
||||
overwrite = args.overwrite
|
||||
limit = args.limit
|
||||
completed_tasks = []
|
||||
errors_summary = []
|
||||
|
||||
# --- Lock for thread-safe file writing ---
|
||||
io_lock = threading.Lock()
|
||||
pro_lock = threading.Lock()
|
||||
@@ -171,21 +88,119 @@ pro_count = 0
|
||||
flash_count = 0
|
||||
pro_quota_exhausted = False
|
||||
|
||||
if overwrite:
|
||||
if output_path.exists():
|
||||
output_path.unlink()
|
||||
if progress_path.exists():
|
||||
progress_path.unlink()
|
||||
else:
|
||||
if progress_path.exists():
|
||||
with open(progress_path, "r", encoding="utf-8") as f:
|
||||
completed_tasks = json.load(f)
|
||||
if output_path.exists():
|
||||
with open(output_path, "r", encoding="utf-8") as f:
|
||||
results = json.load(f)
|
||||
|
||||
completed_set = set((str(f), l) for f, l in completed_tasks)
|
||||
tasks_to_process = [t for t in tasks if (str(t[0]), t[1]) not in completed_set]
|
||||
def discover_tasks(
|
||||
workspace: EvaluationWorkspace,
|
||||
targets: list[Path],
|
||||
) -> tuple[list[tuple[str, str]], list[str]]:
|
||||
workspace.require_directories("Copies", "Par label")
|
||||
discovered: list[tuple[str, str]] = []
|
||||
warnings: list[str] = []
|
||||
for target in targets:
|
||||
if target.is_file():
|
||||
if target.suffix.casefold() != ".jpg":
|
||||
raise CliError(
|
||||
f"Correction target is not a group JPG: {target}",
|
||||
ExitCode.INVALID_ARGUMENTS,
|
||||
)
|
||||
try:
|
||||
target.relative_to(workspace.groups_dir)
|
||||
except ValueError as exc:
|
||||
raise CliError(
|
||||
f"Group image is not inside {workspace.groups_dir}: {target}",
|
||||
ExitCode.INVALID_ARGUMENTS,
|
||||
) from exc
|
||||
discovered.append((str(target), target.parent.name))
|
||||
continue
|
||||
|
||||
group_directories = sorted(
|
||||
(path for path in workspace.groups_dir.iterdir() if path.is_dir()),
|
||||
key=lambda path: path.name.casefold(),
|
||||
)
|
||||
for group_directory in group_directories:
|
||||
images = sorted(
|
||||
group_directory.glob("*.jpg"), key=lambda path: path.name.casefold()
|
||||
)
|
||||
discovered.extend(
|
||||
(str(image), group_directory.name) for image in images
|
||||
)
|
||||
if not group_directories:
|
||||
warnings.append(f"No label groups found in {workspace.groups_dir}")
|
||||
return list(dict.fromkeys(discovered)), warnings
|
||||
|
||||
|
||||
def configure_runtime(
|
||||
workspace: EvaluationWorkspace,
|
||||
discovered_tasks: list[tuple[str, str]],
|
||||
args: argparse.Namespace,
|
||||
*,
|
||||
api_client=None,
|
||||
) -> None:
|
||||
global INPUT_DIR, COPIES_DIR, GROUPS_DIR, output_path, progress_path
|
||||
global tasks, tasks_to_process, results, completed_tasks, errors_summary
|
||||
global overwrite, limit, client, start_time
|
||||
global pro_count, flash_count, pro_quota_exhausted
|
||||
|
||||
INPUT_DIR = workspace.root
|
||||
COPIES_DIR = workspace.copies_dir
|
||||
GROUPS_DIR = workspace.groups_dir
|
||||
output_path = workspace.correction_file
|
||||
progress_path = workspace.correction_progress_file
|
||||
tasks = list(discovered_tasks)
|
||||
overwrite = bool(args.overwrite)
|
||||
limit = args.limit
|
||||
start_time = time.time()
|
||||
errors_summary = []
|
||||
completed_tasks = []
|
||||
results = {label: [] for _file, label in tasks}
|
||||
thread_logs.clear()
|
||||
pro_count = 0
|
||||
flash_count = 0
|
||||
pro_quota_exhausted = False
|
||||
|
||||
if not overwrite:
|
||||
if progress_path.is_file():
|
||||
loaded_progress = read_json(progress_path)
|
||||
if not isinstance(loaded_progress, list):
|
||||
raise TypeError("correction_progress.json must contain a JSON array")
|
||||
completed_tasks = loaded_progress
|
||||
if output_path.is_file():
|
||||
loaded_results = read_json(output_path)
|
||||
if not isinstance(loaded_results, dict):
|
||||
raise TypeError("correction.json must contain a JSON object")
|
||||
results = loaded_results
|
||||
|
||||
completed_set = {(str(file_path), label) for file_path, label in completed_tasks}
|
||||
tasks_to_process = [
|
||||
task for task in tasks if (str(task[0]), task[1]) not in completed_set
|
||||
]
|
||||
client = api_client
|
||||
|
||||
|
||||
def reset_workspace(workspace: EvaluationWorkspace) -> None:
|
||||
"""Apply the explicitly requested correction reset."""
|
||||
print("--- Running Reset ---")
|
||||
for path in (workspace.correction_file, workspace.correction_progress_file):
|
||||
if path.exists():
|
||||
path.unlink()
|
||||
print(f"Deleted: {path}")
|
||||
if workspace.copies_dir.is_dir():
|
||||
for copy_directory in workspace.copies_dir.iterdir():
|
||||
if not copy_directory.is_dir():
|
||||
continue
|
||||
for old_pdf in copy_directory.glob("*_old.pdf"):
|
||||
original = old_pdf.with_name(old_pdf.name.replace("_old.pdf", ".pdf"))
|
||||
if original.exists():
|
||||
original.unlink()
|
||||
old_pdf.replace(original)
|
||||
print(f"Moved: {copy_directory.name}/{old_pdf.name} -> {original.name}")
|
||||
for new_pdf in copy_directory.glob("*_new.pdf"):
|
||||
new_pdf.unlink()
|
||||
print(f"Deleted: {copy_directory.name}/{new_pdf.name}")
|
||||
print(
|
||||
"Reset almost complete. Manually remove groups associated with deleted "
|
||||
"_new PDFs from 'Par label'."
|
||||
)
|
||||
|
||||
def call_gemini_with_retries(model_id, contents, config,
|
||||
fallback_model_id=MODEL_ID_flash):
|
||||
@@ -267,9 +282,6 @@ def correct_boxes_with_gemini(pid, label, pdf_path, original_feedbacks,
|
||||
|
||||
return global_feedbacks + corrected_feedbacks
|
||||
|
||||
import shutil
|
||||
import grouping
|
||||
|
||||
def get_next_group_idx(label):
|
||||
"""Finds the next available Group index for a given label."""
|
||||
target_folder = GROUPS_DIR / label
|
||||
@@ -278,8 +290,6 @@ def get_next_group_idx(label):
|
||||
if not existing: return 0
|
||||
return max([int(f.stem.split("_")[1]) for f in existing])
|
||||
|
||||
from utils import read_all_labels, enonce_total
|
||||
|
||||
def handle_label_errors(pid, label, res, pdf_path):
|
||||
"""Handles Gemini labeling errors, moves/copies files, and returns new tasks."""
|
||||
new_tasks = []
|
||||
@@ -332,7 +342,7 @@ def handle_label_errors(pid, label, res, pdf_path):
|
||||
tprint(f"\tHandling additional-answer for {pid} {label}")
|
||||
try:
|
||||
add_labels = json.loads(call_gemini_with_retries(MODEL_ID_flash, contents, config))
|
||||
except Exception:
|
||||
except Exception: # noqa: BLE001 - invalid auxiliary model response
|
||||
add_labels = []
|
||||
|
||||
keep_error = False
|
||||
@@ -374,7 +384,7 @@ def handle_label_errors(pid, label, res, pdf_path):
|
||||
|
||||
def process_single_task(task_tuple, precomputed_response=None):
|
||||
try:
|
||||
global pro_count, flash_count, pro_quota_exhausted
|
||||
global pro_count, flash_count
|
||||
file_path = task_tuple[0]
|
||||
label = task_tuple[1]
|
||||
can_spawn_tasks = task_tuple[2] if len(task_tuple) > 2 else True
|
||||
@@ -383,8 +393,7 @@ def process_single_task(task_tuple, precomputed_response=None):
|
||||
json_path = group_name + '.json'
|
||||
new_tasks = []
|
||||
|
||||
with open(json_path, 'r') as f:
|
||||
group_data = json.load(f)
|
||||
group_data = read_json(json_path)
|
||||
|
||||
n = len(group_data)
|
||||
d_data = {l[0]: (l[1], l[2], l[3]) for l in group_data}
|
||||
@@ -466,7 +475,7 @@ def process_single_task(task_tuple, precomputed_response=None):
|
||||
for (i,f) in enumerate(res["feedback"]):
|
||||
b = f.get("box_2d")
|
||||
if b:
|
||||
ymin, xmin, ymax, xmax = b
|
||||
ymin, _xmin, ymax, xmax = b
|
||||
ymin = ymin * total_height // 1000
|
||||
ymax = ymax * total_height // 1000
|
||||
|
||||
@@ -493,7 +502,7 @@ def process_single_task(task_tuple, precomputed_response=None):
|
||||
res["feedback"] = correct_boxes_with_gemini(
|
||||
pid, label, pdf_path, res["feedback"],
|
||||
yming, ymaxg, width_r, total_height)
|
||||
except Exception as e:
|
||||
except Exception as e: # noqa: BLE001 - correction fallback
|
||||
tprint(f"\tCorrection failed for Copie {pid}, {group_name} : {e}\n\tRemoving the boxes")
|
||||
# Fallback if the second request fails entirely
|
||||
for (i, f) in enumerate(res["feedback"]):
|
||||
@@ -516,7 +525,7 @@ def process_single_task(task_tuple, precomputed_response=None):
|
||||
tprint(f"Error decoding JSON for {file_path}", file=sys.stderr)
|
||||
with io_lock:
|
||||
errors_summary.append(("Error decoding JSON response", file_path))
|
||||
except Exception as e:
|
||||
except Exception as e: # noqa: BLE001 - per-task processing boundary
|
||||
error_msg = f"Exception processing {file_path}: {e}"
|
||||
print(error_msg, file=sys.stderr)
|
||||
with io_lock:
|
||||
@@ -602,19 +611,20 @@ def resolve_delayed_moves():
|
||||
|
||||
return new_tasks
|
||||
|
||||
if __name__ == "__main__":
|
||||
def run_configured(args: argparse.Namespace) -> ExitCode:
|
||||
global client, tasks_to_process
|
||||
if client is None:
|
||||
client = genai.Client(api_key=api_key)
|
||||
if args.refaire:
|
||||
refaire_path = INPUT_DIR / "refaire.json"
|
||||
overwritten_path = INPUT_DIR / "overwritten_correction.json"
|
||||
|
||||
if refaire_path.exists():
|
||||
with open(refaire_path, "r", encoding="utf-8") as f:
|
||||
refaire_list = json.load(f)
|
||||
refaire_list = read_json(refaire_path)
|
||||
|
||||
overwritten_data = []
|
||||
if overwritten_path.exists():
|
||||
with open(overwritten_path, "r", encoding="utf-8") as f:
|
||||
overwritten_data = json.load(f)
|
||||
overwritten_data = read_json(overwritten_path)
|
||||
|
||||
dirty_results = False
|
||||
|
||||
@@ -659,10 +669,12 @@ if __name__ == "__main__":
|
||||
pdf_path = copie_dir / f"{label}.pdf"
|
||||
is_new = False
|
||||
|
||||
if not pdf_path.exists():
|
||||
if (copie_dir / f"{label}_new.pdf").exists():
|
||||
pdf_path = copie_dir / f"{label}_new.pdf"
|
||||
is_new = True
|
||||
if (
|
||||
not pdf_path.exists()
|
||||
and (copie_dir / f"{label}_new.pdf").exists()
|
||||
):
|
||||
pdf_path = copie_dir / f"{label}_new.pdf"
|
||||
is_new = True
|
||||
# elif (copie_dir / f"{label}_old.pdf").exists():
|
||||
# pdf_path = copie_dir / f"{label}_old.pdf"
|
||||
|
||||
@@ -690,7 +702,11 @@ if __name__ == "__main__":
|
||||
input(f"About to batch from: {args.batch_from}. Press Enter to confirm...")
|
||||
break
|
||||
if args.batch_from not in all_labels:
|
||||
sys.exit(f"Error: Label '{args.batch_from}' not found. Available labels: {all_labels}")
|
||||
raise CliError(
|
||||
f"Label '{args.batch_from}' not found. Available labels: "
|
||||
f"{all_labels}",
|
||||
ExitCode.INVALID_ARGUMENTS,
|
||||
)
|
||||
|
||||
target_idx = all_labels.index(args.batch_from)
|
||||
live_tasks = []
|
||||
@@ -714,58 +730,66 @@ if __name__ == "__main__":
|
||||
|
||||
count_flash = 0
|
||||
count_pro = 0
|
||||
|
||||
with open(batch_flash_file, "w", encoding="utf-8") as f_flash, \
|
||||
open(batch_pro_file, "w", encoding="utf-8") as f_pro:
|
||||
|
||||
for task in batch_tasks:
|
||||
file_path, label = task[0], task[1]
|
||||
group_name = os.path.splitext(file_path)[0]
|
||||
json_path = group_name + '.json'
|
||||
|
||||
with open(json_path, 'r') as jf:
|
||||
group_data = json.load(jf)
|
||||
use_flash = len(group_data) >= 4 or group_data[-1][2] <= 500
|
||||
|
||||
image_data = Path(file_path).read_bytes()
|
||||
b64_img = base64.b64encode(image_data).decode("utf-8")
|
||||
|
||||
# Format payload matching Gemini Batch API file requirements
|
||||
req = {
|
||||
"key": file_path, # The ID returned in the output file
|
||||
"request": {
|
||||
"contents": [{
|
||||
flash_lines = []
|
||||
pro_lines = []
|
||||
for task in batch_tasks:
|
||||
file_path, label = task[0], task[1]
|
||||
json_path = Path(file_path).with_suffix(".json")
|
||||
group_data = read_json(json_path)
|
||||
use_flash = len(group_data) >= 4 or group_data[-1][2] <= 500
|
||||
b64_img = base64.b64encode(Path(file_path).read_bytes()).decode(
|
||||
"utf-8"
|
||||
)
|
||||
request = {
|
||||
"key": file_path,
|
||||
"request": {
|
||||
"contents": [
|
||||
{
|
||||
"role": "user",
|
||||
"parts": [
|
||||
{"inlineData": {"mimeType": "image/jpeg", "data": b64_img}},
|
||||
{"text": prompting.make_prompt(INPUT_DIR,label)}
|
||||
]
|
||||
}],
|
||||
"generation_config": {
|
||||
"temperature": 1.0,
|
||||
"topP": 0.95,
|
||||
"maxOutputTokens": 65535,
|
||||
"responseMimeType": "application/json",
|
||||
"responseSchema": prompting.UNROLLED_SCHEMA
|
||||
{
|
||||
"inlineData": {
|
||||
"mimeType": "image/jpeg",
|
||||
"data": b64_img,
|
||||
}
|
||||
},
|
||||
{"text": prompting.make_prompt(INPUT_DIR, label)},
|
||||
],
|
||||
}
|
||||
}
|
||||
}
|
||||
],
|
||||
"generation_config": {
|
||||
"temperature": 1.0,
|
||||
"topP": 0.95,
|
||||
"maxOutputTokens": 65535,
|
||||
"responseMimeType": "application/json",
|
||||
"responseSchema": prompting.UNROLLED_SCHEMA,
|
||||
},
|
||||
},
|
||||
}
|
||||
line = json.dumps(request)
|
||||
if use_flash:
|
||||
flash_lines.append(line)
|
||||
count_flash += 1
|
||||
else:
|
||||
pro_lines.append(line)
|
||||
count_pro += 1
|
||||
atomic_write_text(
|
||||
batch_flash_file,
|
||||
"\n".join(flash_lines) + ("\n" if flash_lines else ""),
|
||||
)
|
||||
atomic_write_text(
|
||||
batch_pro_file,
|
||||
"\n".join(pro_lines) + ("\n" if pro_lines else ""),
|
||||
)
|
||||
|
||||
if use_flash:
|
||||
f_flash.write(json.dumps(req) + "\n")
|
||||
count_flash += 1
|
||||
else:
|
||||
f_pro.write(json.dumps(req) + "\n")
|
||||
count_pro += 1
|
||||
|
||||
print(f"Batch generation complete.")
|
||||
print("Batch generation complete.")
|
||||
print(f" - {count_flash} requests saved to {batch_flash_file} (for {MODEL_ID_flash})")
|
||||
print(f" - {count_pro} requests saved to {batch_pro_file} (for {MODEL_ID_pro})")
|
||||
print("Upload these files via the File API and create two separate batch jobs.")
|
||||
|
||||
# If there's no live tasks to do, and we aren't doing a batched ingestion, exit right away
|
||||
if not tasks_to_process and not args.deal_with_batched:
|
||||
sys.exit(0)
|
||||
return ExitCode.SUCCESS
|
||||
|
||||
batched_responses = {}
|
||||
if args.deal_with_batched:
|
||||
@@ -807,8 +831,11 @@ if __name__ == "__main__":
|
||||
if new_generated_tasks:
|
||||
for new_task in new_generated_tasks:
|
||||
futures[executor.submit(process_single_task, new_task)] = new_task
|
||||
except Exception as e:
|
||||
except Exception as e: # noqa: BLE001 - future boundary
|
||||
print(f"Exception during task execution: {e}", file=sys.stderr)
|
||||
failed_task = futures[future]
|
||||
with io_lock:
|
||||
errors_summary.append((str(e), failed_task[0]))
|
||||
|
||||
tasks_to_process = [] # Vider la liste une fois traitée
|
||||
|
||||
@@ -840,9 +867,12 @@ if __name__ == "__main__":
|
||||
|
||||
if unresolved_delayed:
|
||||
manual_path = INPUT_DIR / "manual_resolutions.txt"
|
||||
with open(manual_path, "w", encoding="utf-8") as f:
|
||||
f.write("### Use -> x>, -x, ss, sx, xx, xs\n")
|
||||
f.write("\n".join(unresolved_delayed) + "\n")
|
||||
atomic_write_text(
|
||||
manual_path,
|
||||
"### Use -> x>, -x, ss, sx, xx, xs\n"
|
||||
+ "\n".join(unresolved_delayed)
|
||||
+ "\n",
|
||||
)
|
||||
print(f"\n[!] Unresolved delayed tasks found! Wrote to {manual_path}.")
|
||||
print(" Please edit it manually, then run `python resolve_manual.py <InputDir>`")
|
||||
|
||||
@@ -855,3 +885,101 @@ if __name__ == "__main__":
|
||||
print(err, file=sys.stderr)
|
||||
escaped_path = shlex.quote(str(file))
|
||||
print(f"Run : python correction.py {escaped_path}")
|
||||
return ExitCode.PARTIAL if errors_summary else ExitCode.SUCCESS
|
||||
|
||||
|
||||
def run(
|
||||
workspace: EvaluationWorkspace,
|
||||
targets: list[Path],
|
||||
args: argparse.Namespace,
|
||||
*,
|
||||
api_client=None,
|
||||
) -> ExitCode:
|
||||
if args.reset:
|
||||
workspace.require_directories("Copies")
|
||||
reset_workspace(workspace)
|
||||
return ExitCode.SUCCESS
|
||||
workspace.require_directories("Copies", "Par label")
|
||||
workspace.require_files("labels")
|
||||
if args.refaire:
|
||||
workspace.require_files("refaire.json")
|
||||
discovered, warnings = discover_tasks(workspace, targets)
|
||||
for warning in warnings:
|
||||
print(f"Warning: {warning}")
|
||||
configure_runtime(workspace, discovered, args, api_client=api_client)
|
||||
if not discovered and not args.refaire:
|
||||
return ExitCode.PARTIAL
|
||||
try:
|
||||
status = run_configured(args)
|
||||
finally:
|
||||
for thread_id in list(thread_logs):
|
||||
flush_thread_log(thread_id)
|
||||
if warnings and status == ExitCode.SUCCESS:
|
||||
return ExitCode.PARTIAL
|
||||
return status
|
||||
|
||||
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
parser = target_parser("Correct grouped answers with Gemini")
|
||||
parser.add_argument(
|
||||
"additional_targets",
|
||||
nargs="*",
|
||||
type=Path,
|
||||
help="Additional group JPG files from the same evaluation",
|
||||
)
|
||||
parser.add_argument("--overwrite", action="store_true", help="Redo requests")
|
||||
parser.add_argument("--limit", type=int, help="Maximum Gemini Pro calls")
|
||||
parser.add_argument(
|
||||
"--refaire",
|
||||
action="store_true",
|
||||
help="Redo copies and labels listed in refaire.json",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--batch",
|
||||
action="store_true",
|
||||
help="Generate Gemini batch request JSONL files",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--batch-from",
|
||||
metavar="LABEL",
|
||||
help="Process earlier labels live and batch from LABEL onward",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--deal-with-batched",
|
||||
action="store_true",
|
||||
help="Consume batched_correction_result.jsonl",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--reset",
|
||||
action="store_true",
|
||||
help="Delete correction state, restore _old PDFs, and delete _new PDFs",
|
||||
)
|
||||
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)
|
||||
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,
|
||||
)
|
||||
if EvaluationWorkspace.discover(resolved).root != workspace.root:
|
||||
raise CliError(
|
||||
"All targets must belong to the same evaluation",
|
||||
ExitCode.INVALID_ARGUMENTS,
|
||||
)
|
||||
targets.append(resolved)
|
||||
return run(workspace, targets, args)
|
||||
|
||||
return execute(parser, argv, handle)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
|
||||
+251
-274
@@ -1,366 +1,343 @@
|
||||
import sys
|
||||
from functools import lru_cache
|
||||
import os
|
||||
import time
|
||||
import json # Added for schema output
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import threading
|
||||
import time
|
||||
import tkinter as tk
|
||||
from collections.abc import Sequence
|
||||
from functools import lru_cache
|
||||
from pathlib import Path
|
||||
from queue import Empty, Queue
|
||||
from threading import Thread
|
||||
from queue import Queue, Empty
|
||||
|
||||
from pdf2image import convert_from_path
|
||||
from PIL import Image, ImageTk
|
||||
|
||||
# --- Configuration ---
|
||||
from copienator import (
|
||||
CliError,
|
||||
EvaluationWorkspace,
|
||||
ExitCode,
|
||||
atomic_write_json,
|
||||
execute,
|
||||
target_parser,
|
||||
workspace_from_target,
|
||||
)
|
||||
from copienator.filesystem import staged_files
|
||||
|
||||
DELIMITER_WIDTH = 5
|
||||
DELIMITER_COLOR = (0, 0, 0)
|
||||
OUTPUT_SIZE = (1800, 1000)
|
||||
|
||||
parser = argparse.ArgumentParser(description="PDF Cropper")
|
||||
parser.add_argument("path", help="Directory path or PDF file path")
|
||||
parser.add_argument("--fullpage", action="store_true", help="Process all files in full page mode (1 page per output file)")
|
||||
args = parser.parse_args()
|
||||
|
||||
path_arg = args.path
|
||||
fullpage_mode = args.fullpage
|
||||
files = []
|
||||
INPUT_DIR = ""
|
||||
COPIES_DIR = ""
|
||||
|
||||
if os.path.isfile(path_arg) and path_arg.lower().endswith('.pdf'):
|
||||
COPIES_DIR = os.path.abspath(os.path.dirname(path_arg))
|
||||
# If the file is inside a "Copies" folder, set INPUT_DIR to the parent
|
||||
if os.path.basename(COPIES_DIR).lower() == 'copies':
|
||||
INPUT_DIR = os.path.dirname(COPIES_DIR)
|
||||
else:
|
||||
INPUT_DIR = COPIES_DIR
|
||||
files = [os.path.basename(path_arg)]
|
||||
elif os.path.isdir(path_arg):
|
||||
# Support passing either the base dir or the Copies dir directly
|
||||
abs_path = os.path.abspath(path_arg)
|
||||
if os.path.basename(abs_path).lower() == 'copies':
|
||||
COPIES_DIR = abs_path
|
||||
INPUT_DIR = os.path.dirname(abs_path)
|
||||
else:
|
||||
INPUT_DIR = abs_path
|
||||
COPIES_DIR = os.path.join(INPUT_DIR, 'Copies')
|
||||
|
||||
if os.path.exists(COPIES_DIR):
|
||||
files = sorted([f for f in os.listdir(COPIES_DIR) if f.lower().endswith('.pdf') and
|
||||
"nonc" not in f.lower()])
|
||||
else:
|
||||
sys.exit(f"Error: Could not find 'Copies' directory inside {INPUT_DIR}")
|
||||
else:
|
||||
sys.exit("Error: Input must be a directory or a PDF file.")
|
||||
|
||||
|
||||
OUTPUT_DIR = os.path.join(INPUT_DIR, 'Cutleft')
|
||||
|
||||
if not os.path.exists(OUTPUT_DIR):
|
||||
os.makedirs(OUTPUT_DIR)
|
||||
|
||||
# --- Processing Logic ---
|
||||
|
||||
def distribute_pages(total_pages, max_per_file=5):
|
||||
"""
|
||||
Calculates how to split pages into chunks <= max_per_file,
|
||||
balancing the number of columns per file.
|
||||
Example: 12 pages, max 5 -> [4, 4, 4]
|
||||
"""
|
||||
if total_pages == 0:
|
||||
return []
|
||||
|
||||
# Calculate minimum number of files needed
|
||||
num_files = (total_pages + max_per_file - 1) // max_per_file
|
||||
|
||||
# Calculate base size and remainder
|
||||
base_count = total_pages // num_files
|
||||
remainder = total_pages % num_files
|
||||
|
||||
distribution = []
|
||||
for i in range(num_files):
|
||||
# Distribute remainder to the first few files
|
||||
count = base_count + (1 if i < remainder else 0)
|
||||
distribution.append(count)
|
||||
|
||||
return distribution
|
||||
|
||||
def stitch_images(image_list):
|
||||
"""Helper to stitch a list of images horizontally with delimiters."""
|
||||
if not image_list:
|
||||
return None
|
||||
|
||||
num_images = len(image_list)
|
||||
total_width = sum(img.width for img in image_list) + (num_images - 1) * DELIMITER_WIDTH
|
||||
max_height = max(img.height for img in image_list)
|
||||
|
||||
combined = Image.new('RGB', (total_width, max_height), color=(255, 255, 255))
|
||||
|
||||
x_offset = 0
|
||||
for idx, img in enumerate(image_list):
|
||||
combined.paste(img, (x_offset, 0))
|
||||
x_offset += img.width
|
||||
if idx < num_images - 1:
|
||||
delimiter = Image.new('RGB', (DELIMITER_WIDTH, max_height), color=DELIMITER_COLOR)
|
||||
combined.paste(delimiter, (x_offset, 0))
|
||||
x_offset += DELIMITER_WIDTH
|
||||
|
||||
return combined
|
||||
|
||||
import threading
|
||||
pdf_cache_lock = threading.Lock()
|
||||
|
||||
|
||||
def distribute_pages(total_pages: int, max_per_file: int = 5) -> list[int]:
|
||||
"""Distribute pages into balanced chunks no larger than max_per_file."""
|
||||
if total_pages == 0:
|
||||
return []
|
||||
number_of_files = (total_pages + max_per_file - 1) // max_per_file
|
||||
base_count, remainder = divmod(total_pages, number_of_files)
|
||||
return [
|
||||
base_count + (1 if index < remainder else 0)
|
||||
for index in range(number_of_files)
|
||||
]
|
||||
|
||||
|
||||
def stitch_images(image_list: list[Image.Image]) -> Image.Image | None:
|
||||
if not image_list:
|
||||
return None
|
||||
total_width = sum(image.width for image in image_list)
|
||||
total_width += (len(image_list) - 1) * DELIMITER_WIDTH
|
||||
max_height = max(image.height for image in image_list)
|
||||
combined = Image.new("RGB", (total_width, max_height), color="white")
|
||||
x_offset = 0
|
||||
for index, image in enumerate(image_list):
|
||||
combined.paste(image, (x_offset, 0))
|
||||
x_offset += image.width
|
||||
if index < len(image_list) - 1:
|
||||
delimiter = Image.new(
|
||||
"RGB", (DELIMITER_WIDTH, max_height), color=DELIMITER_COLOR
|
||||
)
|
||||
combined.paste(delimiter, (x_offset, 0))
|
||||
x_offset += DELIMITER_WIDTH
|
||||
return combined
|
||||
|
||||
|
||||
@lru_cache(maxsize=3)
|
||||
def _get_pdf_pages_cached(filename):
|
||||
pdf_path = os.path.join(COPIES_DIR, filename)
|
||||
def _get_pdf_pages_cached(pdf_path: Path) -> list[Image.Image]:
|
||||
return convert_from_path(pdf_path)
|
||||
|
||||
def get_pdf_pages(filename):
|
||||
"""Thread-safe wrapper for the cached PDF conversion."""
|
||||
|
||||
def get_pdf_pages(pdf_path: Path) -> list[Image.Image]:
|
||||
"""Thread-safe wrapper around the small PDF conversion cache."""
|
||||
with pdf_cache_lock:
|
||||
return _get_pdf_pages_cached(filename)
|
||||
return _get_pdf_pages_cached(pdf_path)
|
||||
|
||||
def process_single_pdf(filename, shift_offset=0, max_per_file=5):
|
||||
"""
|
||||
Converts PDF to stitched images.
|
||||
Returns a tuple: (preview_image_resized, list_of_split_images, schema_dict)
|
||||
"""
|
||||
|
||||
def process_single_pdf(
|
||||
pdf_path: Path,
|
||||
shift_offset: int = 0,
|
||||
max_per_file: int = 5,
|
||||
) -> tuple[Image.Image, list[Image.Image], dict[str, object]] | None:
|
||||
"""Convert one PDF into a preview, full-resolution splits and metadata."""
|
||||
try:
|
||||
pages = get_pdf_pages(filename)
|
||||
cropped_images = []
|
||||
|
||||
for img in pages:
|
||||
width, height = img.size
|
||||
|
||||
for image in get_pdf_pages(pdf_path):
|
||||
width, height = image.size
|
||||
if max_per_file == 1:
|
||||
# If Single Page mode, take the full width (ignore shift/crop)
|
||||
left = 0
|
||||
right = width
|
||||
left, right = 0, width
|
||||
else:
|
||||
# Original "Cutleft" logic (approx 1/3 width)
|
||||
left = 100 + shift_offset
|
||||
right = (width // 3) + 100 + shift_offset
|
||||
|
||||
# Ensure crop box is valid
|
||||
left = max(0, left)
|
||||
right = min(width, right)
|
||||
|
||||
left = max(0, 100 + shift_offset)
|
||||
right = min(width, width // 3 + 100 + shift_offset)
|
||||
if right > left:
|
||||
crop_box = (left, 0, right, height)
|
||||
cropped = img.crop(crop_box)
|
||||
cropped_images.append(cropped)
|
||||
|
||||
cropped_images.append(image.crop((left, 0, right, height)))
|
||||
if not cropped_images:
|
||||
return None
|
||||
|
||||
# 1. Generate Schema / Distribution
|
||||
col_distribution = distribute_pages(len(cropped_images), max_per_file=max_per_file)
|
||||
|
||||
# 2. Generate Split Images (Full Resolution)
|
||||
distribution = distribute_pages(len(cropped_images), max_per_file)
|
||||
split_images = []
|
||||
current_idx = 0
|
||||
for count in col_distribution:
|
||||
chunk = cropped_images[current_idx : current_idx + count]
|
||||
stitched_chunk = stitch_images(chunk)
|
||||
split_images.append(stitched_chunk)
|
||||
current_idx += count
|
||||
|
||||
# 3. Generate Preview (All stitched together, Resized)
|
||||
current_index = 0
|
||||
for count in distribution:
|
||||
stitched = stitch_images(cropped_images[current_index : current_index + count])
|
||||
if stitched is not None:
|
||||
split_images.append(stitched)
|
||||
current_index += count
|
||||
full_stitch = stitch_images(cropped_images)
|
||||
preview_resized = full_stitch.resize(OUTPUT_SIZE, Image.BILINEAR)
|
||||
|
||||
schema = {
|
||||
"original_filename": filename,
|
||||
if full_stitch is None:
|
||||
return None
|
||||
preview = full_stitch.resize(OUTPUT_SIZE, Image.Resampling.BILINEAR)
|
||||
schema: dict[str, object] = {
|
||||
"original_filename": pdf_path.name,
|
||||
"total_pages": len(cropped_images),
|
||||
"number_of_files": len(split_images),
|
||||
"columns_per_file": col_distribution
|
||||
"columns_per_file": distribution,
|
||||
}
|
||||
|
||||
return (preview_resized, split_images, schema)
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error processing {filename}: {e}")
|
||||
return preview, split_images, schema
|
||||
except Exception as exc: # noqa: BLE001 - interactive item failure
|
||||
print(f"Error processing {pdf_path.name}: {exc}")
|
||||
return None
|
||||
|
||||
def save_results(result_tuple, filename):
|
||||
"""
|
||||
Saves the split images and the schema JSON.
|
||||
"""
|
||||
_, splits, schema = result_tuple
|
||||
base_name = os.path.splitext(filename)[0]
|
||||
|
||||
# --- Cleanup: Delete existing files for this PDF ---
|
||||
for f in os.listdir(OUTPUT_DIR):
|
||||
file_path = os.path.join(OUTPUT_DIR, f)
|
||||
if f == f"{base_name}_schema.json":
|
||||
os.remove(file_path)
|
||||
elif f.startswith(f"{base_name}_") and f.endswith(".jpg"):
|
||||
suffix = f[len(base_name)+1:-4]
|
||||
if suffix.isdigit():
|
||||
os.remove(file_path)
|
||||
|
||||
# Save Images
|
||||
for i, img in enumerate(splits):
|
||||
suffix = f"_{i+1:02d}"
|
||||
output_filename = f"{base_name}{suffix}.jpg"
|
||||
output_path = os.path.join(OUTPUT_DIR, output_filename)
|
||||
img.save(output_path, "JPEG", quality=95)
|
||||
print(f"Saved: {output_filename}")
|
||||
|
||||
# Save Schema
|
||||
json_filename = f"{base_name}_schema.json"
|
||||
json_path = os.path.join(OUTPUT_DIR, json_filename)
|
||||
with open(json_path, 'w') as f:
|
||||
json.dump(schema, f, indent=4)
|
||||
print(f"Saved schema: {json_filename}")
|
||||
def _previous_cutleft_outputs(output_dir: Path, base_name: str) -> set[str]:
|
||||
if not output_dir.is_dir():
|
||||
return set()
|
||||
result = {f"{base_name}_schema.json"}
|
||||
for path in output_dir.glob(f"{base_name}_*.jpg"):
|
||||
suffix = path.stem.removeprefix(f"{base_name}_")
|
||||
if suffix.isdigit():
|
||||
result.add(path.name)
|
||||
return result
|
||||
|
||||
|
||||
# --- GUI Application ---
|
||||
def save_results(
|
||||
result: tuple[Image.Image, list[Image.Image], dict[str, object]],
|
||||
pdf_path: Path,
|
||||
output_dir: Path,
|
||||
) -> None:
|
||||
"""Atomically replace every Cutleft output associated with one copy."""
|
||||
_, splits, schema = result
|
||||
base_name = pdf_path.stem
|
||||
previous = _previous_cutleft_outputs(output_dir, base_name)
|
||||
with staged_files(output_dir, remove=previous) as staging:
|
||||
for index, image in enumerate(splits, start=1):
|
||||
filename = f"{base_name}_{index:02d}.jpg"
|
||||
image.save(staging / filename, "JPEG", quality=95)
|
||||
atomic_write_json(staging / f"{base_name}_schema.json", schema)
|
||||
for index in range(1, len(splits) + 1):
|
||||
print(f"Saved: {base_name}_{index:02d}.jpg")
|
||||
print(f"Saved schema: {base_name}_schema.json")
|
||||
|
||||
|
||||
class ImageReviewer:
|
||||
def __init__(self, file_list, default_max_per_file=5):
|
||||
self.files = file_list
|
||||
def __init__(
|
||||
self,
|
||||
files: list[Path],
|
||||
output_dir: Path,
|
||||
default_max_per_file: int = 5,
|
||||
) -> None:
|
||||
self.files = files
|
||||
self.output_dir = output_dir
|
||||
self.index = 0
|
||||
self.current_shift = 0
|
||||
self.default_max_per_file = default_max_per_file
|
||||
self.current_max_per_file = default_max_per_file
|
||||
self.current_preview = None # Only stores the resized preview for GUI
|
||||
self.current_preview: Image.Image | None = None
|
||||
self.is_processing = False
|
||||
self.manual_queue: Queue[
|
||||
tuple[Image.Image, list[Image.Image], dict[str, object]] | None
|
||||
] = Queue()
|
||||
|
||||
# Queue for manual re-processing results
|
||||
self.manual_queue = Queue()
|
||||
|
||||
# Setup GUI
|
||||
self.root = tk.Tk()
|
||||
self.root.title("PDF Cropper")
|
||||
self.root.geometry("+100+100")
|
||||
|
||||
self.label_img = tk.Label(self.root)
|
||||
self.label_img.pack()
|
||||
|
||||
self.label_info = tk.Label(self.root, text="", font=("Arial", 12, "bold"))
|
||||
self.label_info.pack(pady=5)
|
||||
self.root.bind("<Return>", self.on_next)
|
||||
self.root.bind("n", lambda _event: self.on_shift(50))
|
||||
self.root.bind("N", lambda _event: self.on_shift(100))
|
||||
self.root.bind("t", lambda _event: self.on_shift(-50))
|
||||
self.root.bind("1", lambda _event: self.on_set_max_pages(1))
|
||||
|
||||
# Bindings
|
||||
self.root.bind('<Return>', self.on_next)
|
||||
self.root.bind('n', lambda e: self.on_shift(50))
|
||||
self.root.bind('N', lambda e: self.on_shift(100))
|
||||
self.root.bind('t', lambda e: self.on_shift(-50))
|
||||
self.root.bind('1', lambda e: self.on_set_max_pages(1))
|
||||
|
||||
# Start background pre-fetcher
|
||||
self.bg_thread = Thread(target=self.prefetch_worker, daemon=True)
|
||||
self.bg_thread.start()
|
||||
|
||||
# Load first image
|
||||
Thread(target=self.prefetch_worker, daemon=True).start()
|
||||
self.load_current_image()
|
||||
|
||||
self.root.lift()
|
||||
self.root.focus_force()
|
||||
self.root.mainloop()
|
||||
|
||||
def on_set_max_pages(self, count):
|
||||
def on_set_max_pages(self, count: int) -> None:
|
||||
if self.is_processing:
|
||||
return
|
||||
self.current_max_per_file = count
|
||||
print(f"Setting max pages per file: {count}")
|
||||
self.trigger_processing(self.files[self.index], self.current_shift)
|
||||
|
||||
def prefetch_worker(self):
|
||||
"""Background thread to load the NEXT file's PDF pages into RAM."""
|
||||
idx_to_process = -1
|
||||
def prefetch_worker(self) -> None:
|
||||
processed_index = -1
|
||||
while True:
|
||||
target = self.index + 1
|
||||
if target < len(self.files) and target != idx_to_process:
|
||||
fname = self.files[target]
|
||||
get_pdf_pages(fname)
|
||||
idx_to_process = target
|
||||
if target < len(self.files) and target != processed_index:
|
||||
get_pdf_pages(self.files[target])
|
||||
processed_index = target
|
||||
time.sleep(0.05)
|
||||
|
||||
def load_current_image(self, use_prefetch=False):
|
||||
def load_current_image(self) -> None:
|
||||
if self.index >= len(self.files):
|
||||
print("All files processed.")
|
||||
self.root.destroy()
|
||||
return
|
||||
|
||||
filename = self.files[self.index]
|
||||
self.is_processing = False
|
||||
self.current_shift = 0
|
||||
self.trigger_processing(self.files[self.index], self.current_shift)
|
||||
|
||||
self.trigger_processing(filename, self.current_shift)
|
||||
|
||||
def trigger_processing(self, filename, shift):
|
||||
"""Starts a thread to process image so GUI doesn't freeze."""
|
||||
def trigger_processing(self, pdf_path: Path, shift: int) -> None:
|
||||
self.is_processing = True
|
||||
self.label_info.configure(text=f"Processing {filename} (Shift {shift})... Please wait.", fg="red")
|
||||
self.label_info.configure(
|
||||
text=f"Processing {pdf_path.name} (Shift {shift})... Please wait.",
|
||||
fg="red",
|
||||
)
|
||||
|
||||
def worker():
|
||||
res = process_single_pdf(filename, shift, self.current_max_per_file)
|
||||
self.manual_queue.put(res)
|
||||
|
||||
Thread(target=worker, daemon=True).start()
|
||||
self.check_manual_queue(filename)
|
||||
|
||||
def check_manual_queue(self, filename):
|
||||
"""Polls the manual queue for result."""
|
||||
try:
|
||||
result = self.manual_queue.get_nowait()
|
||||
if result:
|
||||
self.handle_processing_result(result, filename)
|
||||
else:
|
||||
print(f"Failed to process {filename}, skipping.")
|
||||
self.index += 1
|
||||
self.load_current_image(use_prefetch=True)
|
||||
self.is_processing = False
|
||||
except Empty:
|
||||
self.root.after(100, lambda: self.check_manual_queue(filename))
|
||||
|
||||
def handle_processing_result(self, result, filename):
|
||||
"""Unpacks result, saves files, and updates display."""
|
||||
preview, splits, schema = result
|
||||
self.current_preview = preview
|
||||
|
||||
Thread(target=save_results, args=(result, filename), daemon=True).start()
|
||||
|
||||
self.update_display(filename, schema)
|
||||
|
||||
def update_display(self, filename, schema=None):
|
||||
if self.current_preview:
|
||||
tk_image = ImageTk.PhotoImage(self.current_preview)
|
||||
self.label_img.configure(image=tk_image)
|
||||
self.label_img.image = tk_image
|
||||
|
||||
schema_info = ""
|
||||
if schema:
|
||||
cols = str(schema['columns_per_file'])
|
||||
schema_info = f"\nFiles: {schema['number_of_files']} | Cols: {cols}"
|
||||
|
||||
self.label_info.configure(
|
||||
text=f"[{self.index+1}/{len(self.files)}] {filename} | Shift: {self.current_shift}px"
|
||||
f"{schema_info}\n"
|
||||
f"Enter: Next | n: +50 | N: +100 | t: -50 | 1: use single column",
|
||||
fg="black"
|
||||
def worker() -> None:
|
||||
self.manual_queue.put(
|
||||
process_single_pdf(pdf_path, shift, self.current_max_per_file)
|
||||
)
|
||||
|
||||
def on_shift(self, amount):
|
||||
Thread(target=worker, daemon=True).start()
|
||||
self.check_manual_queue(pdf_path)
|
||||
|
||||
def check_manual_queue(self, pdf_path: Path) -> None:
|
||||
try:
|
||||
result = self.manual_queue.get_nowait()
|
||||
if result is None:
|
||||
print(f"Failed to process {pdf_path.name}, skipping.")
|
||||
self.index += 1
|
||||
self.load_current_image()
|
||||
else:
|
||||
self.handle_processing_result(result, pdf_path)
|
||||
self.is_processing = False
|
||||
except Empty:
|
||||
self.root.after(100, lambda: self.check_manual_queue(pdf_path))
|
||||
|
||||
def handle_processing_result(
|
||||
self,
|
||||
result: tuple[Image.Image, list[Image.Image], dict[str, object]],
|
||||
pdf_path: Path,
|
||||
) -> None:
|
||||
self.current_preview = result[0]
|
||||
save_results(result, pdf_path, self.output_dir)
|
||||
self.update_display(pdf_path.name, result[2])
|
||||
|
||||
def update_display(self, filename: str, schema: dict[str, object]) -> None:
|
||||
if self.current_preview is None:
|
||||
return
|
||||
tk_image = ImageTk.PhotoImage(self.current_preview)
|
||||
self.label_img.configure(image=tk_image)
|
||||
self.label_img.image = tk_image
|
||||
self.label_info.configure(
|
||||
text=(
|
||||
f"[{self.index + 1}/{len(self.files)}] {filename} | "
|
||||
f"Shift: {self.current_shift}px\nFiles: {schema['number_of_files']} | "
|
||||
f"Cols: {schema['columns_per_file']}\n"
|
||||
"Enter: Next | n: +50 | N: +100 | t: -50 | "
|
||||
"1: use single column"
|
||||
),
|
||||
fg="black",
|
||||
)
|
||||
|
||||
def on_shift(self, amount: int) -> None:
|
||||
if self.is_processing:
|
||||
return
|
||||
self.current_shift += amount
|
||||
print(f"Applying shift: {self.current_shift}")
|
||||
self.trigger_processing(self.files[self.index], self.current_shift)
|
||||
|
||||
def on_next(self, event):
|
||||
def on_next(self, _event: object) -> None:
|
||||
if self.is_processing:
|
||||
return
|
||||
self.index += 1
|
||||
self.current_shift = 0
|
||||
self.current_max_per_file = self.default_max_per_file
|
||||
self.load_current_image(use_prefetch=True)
|
||||
self.load_current_image()
|
||||
|
||||
|
||||
# --- Entry Point ---
|
||||
if __name__ == "__main__":
|
||||
def _selected_files(
|
||||
workspace: EvaluationWorkspace,
|
||||
target: Path,
|
||||
) -> list[Path]:
|
||||
workspace.require_directories("Copies")
|
||||
if target.is_file():
|
||||
if target.suffix.casefold() != ".pdf":
|
||||
raise CliError(f"Target is not a PDF: {target}", ExitCode.INVALID_ARGUMENTS)
|
||||
return [target]
|
||||
return sorted(
|
||||
(
|
||||
path
|
||||
for path in workspace.copies_dir.glob("*.pdf")
|
||||
if "nonc" not in path.name.casefold()
|
||||
),
|
||||
key=lambda path: path.name.casefold(),
|
||||
)
|
||||
|
||||
|
||||
def run(
|
||||
workspace: EvaluationWorkspace,
|
||||
target: Path,
|
||||
*,
|
||||
fullpage: bool = False,
|
||||
) -> ExitCode:
|
||||
files = _selected_files(workspace, target)
|
||||
if not files:
|
||||
print("No PDF files found.")
|
||||
else:
|
||||
app = ImageReviewer(files, default_max_per_file=1 if fullpage_mode else 5)
|
||||
return ExitCode.SUCCESS
|
||||
workspace.cutleft_dir.mkdir(parents=True, exist_ok=True)
|
||||
_get_pdf_pages_cached.cache_clear()
|
||||
ImageReviewer(
|
||||
files,
|
||||
workspace.cutleft_dir,
|
||||
default_max_per_file=1 if fullpage else 5,
|
||||
)
|
||||
return ExitCode.SUCCESS
|
||||
|
||||
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
parser = target_parser("Interactively crop the label margin from PDF copies")
|
||||
parser.add_argument(
|
||||
"--fullpage",
|
||||
action="store_true",
|
||||
help="Use each complete page instead of cropping the label margin",
|
||||
)
|
||||
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)
|
||||
return run(workspace, target, fullpage=args.fullpage)
|
||||
|
||||
return execute(parser, argv, handle)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
|
||||
+54
-20
@@ -1,16 +1,26 @@
|
||||
import sys
|
||||
import os
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import glob
|
||||
import json
|
||||
import urllib.request
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import tempfile
|
||||
import shutil
|
||||
import urllib.request
|
||||
from collections.abc import Sequence
|
||||
from uuid import uuid4
|
||||
|
||||
from copienator import (
|
||||
CliError,
|
||||
EvaluationWorkspace,
|
||||
ExitCode,
|
||||
evaluation_parser,
|
||||
execute,
|
||||
workspace_from_args,
|
||||
)
|
||||
from platform_utils import WindowsLabelError, validate_windows_labels
|
||||
from utils import compile_to_pdf
|
||||
|
||||
|
||||
def fetch_and_save_sub_text(ex_id, indices, label, text_path):
|
||||
"""Fetches text for a specific sub-question and saves it to Text/{label}.tex"""
|
||||
qinds = ",".join(map(str, indices))
|
||||
@@ -29,6 +39,7 @@ def fetch_and_save_sub_text(ex_id, indices, label, text_path):
|
||||
compile_to_pdf(content, pdf_file)
|
||||
except Exception as e:
|
||||
print(f"Error fetching sub-text from {url}: {e}")
|
||||
raise
|
||||
|
||||
def fetch_and_save_sub_sol(ex_id, indices, label, sol_path):
|
||||
"""Fetches text for a specific sub-question and saves it to Text/{label}.tex"""
|
||||
@@ -48,6 +59,7 @@ def fetch_and_save_sub_sol(ex_id, indices, label, sol_path):
|
||||
compile_to_pdf(content, pdf_file)
|
||||
except Exception as e:
|
||||
print(f"Error fetching sub-text from {url}: {e}")
|
||||
raise
|
||||
|
||||
|
||||
ROMANS_CAP = ["", "I", "II", "III", "IV", "V", "VI", "VII", "VIII", "IX", "X"]
|
||||
@@ -127,18 +139,21 @@ def save_split_content(text, path, base_fname, problem):
|
||||
f.write(chunk)
|
||||
|
||||
|
||||
def process_directory(directory):
|
||||
def process_directory(workspace: EvaluationWorkspace) -> ExitCode:
|
||||
directory = str(workspace.root)
|
||||
# Find the first .tex file in the directory
|
||||
tex_files = glob.glob(os.path.join(directory, "*.tex"))
|
||||
if not tex_files:
|
||||
print(f"No .tex file found in {directory}. Looking in /Staging/Interro/")
|
||||
int_name = directory[:-1] if directory.endswith("/") else directory
|
||||
int_name = directory.removesuffix("/")
|
||||
tex_path = os.path.join(os.path.expanduser("~"), "Prépa/Staging/Interro", f"{int_name}.tex")
|
||||
if os.path.exists(tex_path):
|
||||
tex_file = tex_path
|
||||
else:
|
||||
print("Not found in ", tex_path)
|
||||
return
|
||||
raise CliError(
|
||||
f"No .tex input found in {workspace.root}",
|
||||
ExitCode.INVALID_WORKSPACE,
|
||||
)
|
||||
else:
|
||||
tex_file = tex_files[0]
|
||||
|
||||
@@ -153,8 +168,10 @@ def process_directory(directory):
|
||||
for p in paths.values():
|
||||
os.makedirs(p, exist_ok=True)
|
||||
|
||||
labels_file = os.path.join(directory, "labels")
|
||||
labels_file = workspace.labels_file
|
||||
labels_staging = labels_file.with_name(f".{labels_file.name}.{uuid4().hex}.tmp")
|
||||
current_ex_num = 1
|
||||
had_errors = False
|
||||
|
||||
# Read entirely to allow chunking
|
||||
with open(tex_file, 'r', encoding='utf-8') as f_in:
|
||||
@@ -162,8 +179,11 @@ def process_directory(directory):
|
||||
|
||||
# Split by the specific SHEETINFO tag
|
||||
blocks = content.split("%%SHEETINFO :")
|
||||
if len(blocks) == 1:
|
||||
print(f"No SHEETINFO blocks found in {tex_file}")
|
||||
return ExitCode.PARTIAL
|
||||
|
||||
with open(labels_file, 'w', encoding='utf-8') as f_labels:
|
||||
with open(labels_staging, 'w', encoding='utf-8') as f_labels:
|
||||
# Skip blocks[0] (content before first SHEETINFO)
|
||||
for block in blocks[1:]:
|
||||
parts_line = block.split("\n", 1)
|
||||
@@ -177,6 +197,7 @@ def process_directory(directory):
|
||||
|
||||
try:
|
||||
data = json.loads(json_str)
|
||||
block_labels = []
|
||||
# Construct 'ids' parameter
|
||||
ex_id = str(data['id'])
|
||||
selection = data.get('select')
|
||||
@@ -193,7 +214,7 @@ def process_directory(directory):
|
||||
if not indexes:
|
||||
label = f"Ex {current_ex_num}"
|
||||
validate_windows_labels([label])
|
||||
f_labels.write(f"{label}\n")
|
||||
block_labels.append(label)
|
||||
fetch_and_save_sub_text(ids, [], label, paths['Text2'])
|
||||
fetch_and_save_sub_sol(ids, [], label, paths['Sol2'])
|
||||
else:
|
||||
@@ -201,7 +222,7 @@ def process_directory(directory):
|
||||
suffix = format_indices(item['indices'], problem)
|
||||
label = f"Ex {current_ex_num}" + (f" : {suffix}" if suffix else "")
|
||||
validate_windows_labels([label])
|
||||
f_labels.write(f"{label}\n")
|
||||
block_labels.append(label)
|
||||
fetch_and_save_sub_text(ids, item['indices'], label, paths['Text2'])
|
||||
fetch_and_save_sub_sol(ids, item['indices'], label, paths['Sol2'])
|
||||
|
||||
@@ -244,18 +265,31 @@ def process_directory(directory):
|
||||
save_split_content(s_text, paths['Sol'], base_filename, problem)
|
||||
save_split_content(p_text, paths['Persp'], base_filename, problem)
|
||||
|
||||
for label in block_labels:
|
||||
f_labels.write(f"{label}\n")
|
||||
current_ex_num += 1
|
||||
|
||||
except WindowsLabelError:
|
||||
labels_staging.unlink(missing_ok=True)
|
||||
raise
|
||||
except json.JSONDecodeError:
|
||||
print(f"Error decoding JSON in block: {json_str}")
|
||||
except Exception as e:
|
||||
had_errors = True
|
||||
except Exception as e: # noqa: BLE001 - one malformed exercise is partial
|
||||
print(f"Error processing block {ex_id if 'ex_id' in locals() else 'unknown'}: {e}")
|
||||
had_errors = True
|
||||
|
||||
labels_staging.replace(labels_file)
|
||||
return ExitCode.PARTIAL if had_errors else ExitCode.SUCCESS
|
||||
|
||||
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
return evaluation_parser("Generate statement metadata from SHEETINFO blocks")
|
||||
|
||||
|
||||
def main(argv: Sequence[str] | None = None) -> int:
|
||||
parser = build_parser()
|
||||
return execute(parser, argv, lambda args: process_directory(workspace_from_args(args)))
|
||||
|
||||
if __name__ == "__main__":
|
||||
if len(sys.argv) < 2:
|
||||
print("Usage: python script.py <Dir>")
|
||||
sys.exit(1)
|
||||
|
||||
process_directory(sys.argv[1])
|
||||
raise SystemExit(main())
|
||||
|
||||
+70
-50
@@ -1,63 +1,83 @@
|
||||
import os
|
||||
import sys
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
from collections.abc import Sequence
|
||||
|
||||
from google import genai
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Download and combine completed batch jobs for a directory.")
|
||||
parser.add_argument("root_dir", type=str, help="Directory containing the original batches")
|
||||
args = parser.parse_args()
|
||||
import config
|
||||
from copienator import (
|
||||
CliError,
|
||||
EvaluationWorkspace,
|
||||
ExitCode,
|
||||
atomic_write_bytes,
|
||||
evaluation_parser,
|
||||
execute,
|
||||
read_json,
|
||||
workspace_from_args,
|
||||
)
|
||||
|
||||
target_dir = Path(args.root_dir)
|
||||
dir_name = target_dir.name
|
||||
output_path = target_dir / "batched_correction_result.jsonl"
|
||||
|
||||
if "GEMINI_API_KEY" not in os.environ:
|
||||
sys.exit("Error: GEMINI_API_KEY environment variable not set.")
|
||||
def run(workspace: EvaluationWorkspace, *, client=None) -> ExitCode:
|
||||
if client is None:
|
||||
if not config.API_KEY:
|
||||
raise CliError("GEMINI_API_KEY is not configured")
|
||||
client = genai.Client(api_key=config.API_KEY)
|
||||
matching = []
|
||||
if workspace.batch_jobs_file.is_file():
|
||||
manifest = read_json(workspace.batch_jobs_file)
|
||||
jobs = manifest.get("jobs") if isinstance(manifest, dict) else None
|
||||
if not isinstance(jobs, dict):
|
||||
raise CliError(f"Invalid batch manifest: {workspace.batch_jobs_file}")
|
||||
matching = [
|
||||
client.batches.get(name=entry["name"])
|
||||
for entry in jobs.values()
|
||||
if isinstance(entry, dict) and isinstance(entry.get("name"), str)
|
||||
]
|
||||
else:
|
||||
matching = [
|
||||
job
|
||||
for job in client.batches.list()
|
||||
if workspace.name in str(getattr(job, "display_name", ""))
|
||||
]
|
||||
if not matching:
|
||||
raise CliError(
|
||||
f"No batch jobs found for evaluation {workspace.name!r}"
|
||||
)
|
||||
for job in matching:
|
||||
state = job.state.name if hasattr(job.state, "name") else job.state
|
||||
print(f"{job.display_name}: {state}")
|
||||
if state != "JOB_STATE_SUCCEEDED":
|
||||
print("Not all matching jobs have succeeded yet.")
|
||||
return ExitCode.PARTIAL
|
||||
|
||||
client = genai.Client()
|
||||
chunks = []
|
||||
incomplete = False
|
||||
for job in matching:
|
||||
destination = getattr(job, "dest", None)
|
||||
file_name = getattr(destination, "file_name", None)
|
||||
if not file_name:
|
||||
print(f"Warning: {job.display_name} has no output file.")
|
||||
incomplete = True
|
||||
continue
|
||||
payload = client.files.download(file=file_name)
|
||||
chunks.append(payload.rstrip(b"\n"))
|
||||
if not chunks:
|
||||
return ExitCode.PARTIAL
|
||||
output_path = workspace.batched_correction_result_file
|
||||
atomic_write_bytes(output_path, b"\n".join(chunks) + b"\n")
|
||||
print(f"Saved combined batch results to {output_path}")
|
||||
return ExitCode.PARTIAL if incomplete else ExitCode.SUCCESS
|
||||
|
||||
print(f"Fetching jobs matching '{dir_name}'...")
|
||||
all_jobs = client.batches.list()
|
||||
matching_jobs = []
|
||||
|
||||
# 1. Find jobs associated with this directory
|
||||
for job in all_jobs:
|
||||
if hasattr(job, 'display_name') and job.display_name and dir_name in job.display_name:
|
||||
matching_jobs.append(job)
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
return evaluation_parser("Download and combine correction batch results")
|
||||
|
||||
if not matching_jobs:
|
||||
sys.exit(f"No batch jobs found containing '{dir_name}' in their display name.")
|
||||
|
||||
# 2. Check that all matching jobs are complete
|
||||
for job in matching_jobs:
|
||||
state = job.state.name if hasattr(job.state, 'name') else job.state
|
||||
print(f"Found Job: {job.display_name} | State: {state}")
|
||||
if state != 'JOB_STATE_SUCCEEDED':
|
||||
sys.exit(f"Error: Job '{job.display_name}' has not succeeded yet. Try again later.")
|
||||
def main(argv: Sequence[str] | None = None) -> int:
|
||||
parser = build_parser()
|
||||
return execute(parser, argv, lambda args: run(workspace_from_args(args)))
|
||||
|
||||
# 3. Download and concatenate
|
||||
print("\nAll jobs succeeded. Downloading results...")
|
||||
combined_data = b""
|
||||
|
||||
for job in matching_jobs:
|
||||
if hasattr(job, 'dest') and job.dest and hasattr(job.dest, 'file_name') and job.dest.file_name:
|
||||
print(f"Downloading output for {job.display_name}...")
|
||||
file_content_bytes = client.files.download(file=job.dest.file_name)
|
||||
|
||||
combined_data += file_content_bytes
|
||||
# Ensure proper line separation between files in JSONL
|
||||
if combined_data and not combined_data.endswith(b'\n'):
|
||||
combined_data += b'\n'
|
||||
else:
|
||||
print(f"Warning: Job {job.display_name} succeeded but has no output file.")
|
||||
|
||||
# 4. Save to destination
|
||||
with open(output_path, "wb") as f:
|
||||
f.write(combined_data)
|
||||
|
||||
print(f"\nSuccess! All results concatenated and saved to:\n{output_path}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
raise SystemExit(main())
|
||||
|
||||
+91
-51
@@ -1,18 +1,30 @@
|
||||
import re
|
||||
import os
|
||||
import utils
|
||||
import subprocess
|
||||
import sys
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import re
|
||||
from collections.abc import Sequence
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from pathlib import Path
|
||||
from pydantic import BaseModel, Field
|
||||
from typing import List, Union
|
||||
|
||||
from google import genai
|
||||
from google.genai import types
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
import config
|
||||
import utils
|
||||
from copienator import (
|
||||
CliError,
|
||||
EvaluationWorkspace,
|
||||
ExitCode,
|
||||
atomic_write_text,
|
||||
evaluation_parser,
|
||||
execute,
|
||||
workspace_from_args,
|
||||
)
|
||||
from platform_utils import validate_windows_labels
|
||||
from utils import compile_to_pdf
|
||||
|
||||
|
||||
def get_lcp(s1: str, s2: str) -> str:
|
||||
i = 0
|
||||
while i < len(s1) and i < len(s2) and s1[i] == s2[i]:
|
||||
@@ -25,8 +37,6 @@ def get_lcp(s1: str, s2: str) -> str:
|
||||
return lcp
|
||||
|
||||
|
||||
import config
|
||||
|
||||
MODEL_ID = config.MODEL_LITE_ID
|
||||
api_key = config.API_KEY
|
||||
|
||||
@@ -36,7 +46,7 @@ class QuestionOnlyItem(BaseModel):
|
||||
question_content: str = Field(description="The source text of the question, strictly extracted from the enonce file, EXCLUDING the label itself.")
|
||||
|
||||
class ExamQuestions(BaseModel):
|
||||
questions: List[QuestionOnlyItem]
|
||||
questions: list[QuestionOnlyItem]
|
||||
|
||||
# --- Modèles pour la Requête 2 ---
|
||||
class SolutionOnlyItem(BaseModel):
|
||||
@@ -44,7 +54,7 @@ class SolutionOnlyItem(BaseModel):
|
||||
solution_content: str = Field(description="The source text of the solution, strictly extracted from the correction file.")
|
||||
|
||||
class ExamSolutions(BaseModel):
|
||||
solutions: List[SolutionOnlyItem]
|
||||
solutions: list[SolutionOnlyItem]
|
||||
|
||||
# --- Modèles pour la Requête 3 ---
|
||||
class ExtractedContext(BaseModel):
|
||||
@@ -53,7 +63,7 @@ class ExtractedContext(BaseModel):
|
||||
context_content: str = Field(description="The source text of the definitions, notations, or hypotheses, extracted from the enonce.")
|
||||
|
||||
class ExamContext(BaseModel):
|
||||
contexts: List[ExtractedContext]
|
||||
contexts: list[ExtractedContext]
|
||||
|
||||
# --- Modèles pour la Requête 4 (Barèmes) ---
|
||||
class RubricItem(BaseModel):
|
||||
@@ -61,7 +71,7 @@ class RubricItem(BaseModel):
|
||||
rubric_content: str = Field(description="Le barème détaillé en français.")
|
||||
|
||||
class GroupRubrics(BaseModel):
|
||||
rubrics: List[RubricItem]
|
||||
rubrics: list[RubricItem]
|
||||
|
||||
PROMPT_4 = """Je te fournis les questions, le contexte éventuel, et les corrections pour un groupe de questions d'un examen.
|
||||
Ta tâche :
|
||||
@@ -87,10 +97,10 @@ class ContextItem(BaseModel):
|
||||
content: str # Juste une string encapsulée pour le différencier facilement
|
||||
|
||||
class ExamExtraction(BaseModel):
|
||||
items: List[Union[QuestionItem, ContextItem]] # Liste mixte
|
||||
items: list[QuestionItem | ContextItem] # Liste mixte
|
||||
|
||||
class GroupedExamExtraction(BaseModel):
|
||||
groups: List[List[Union[QuestionItem, ContextItem]]]
|
||||
groups: list[list[QuestionItem | ContextItem]]
|
||||
|
||||
PROMPT_1 = """I am providing:
|
||||
1. A PDF of an exam (`enonce.pdf`)
|
||||
@@ -138,15 +148,20 @@ the `Let N, M be two commutating matrices` part is not a question itself, and is
|
||||
Return the result as a JSON list.
|
||||
"""
|
||||
|
||||
def find_file(folder: Path, base_name: str) -> Path:
|
||||
def find_file(folder: Path, base_name: str) -> Path | None:
|
||||
for ext in [".org", ".tex"]:
|
||||
path = folder / f"{base_name}{ext}"
|
||||
if path.is_file():
|
||||
return path
|
||||
return None
|
||||
|
||||
def process_exam(folder_path: str, restart: bool = False):
|
||||
folder = Path(folder_path)
|
||||
def process_exam(
|
||||
workspace: EvaluationWorkspace,
|
||||
restart: bool = False,
|
||||
*,
|
||||
api_client=None,
|
||||
) -> ExitCode:
|
||||
folder = workspace.root
|
||||
|
||||
cache_dir = folder / "Cache"
|
||||
tmp_dir = folder / "Tmp"
|
||||
@@ -168,15 +183,21 @@ def process_exam(folder_path: str, restart: bool = False):
|
||||
if not correction_path: missing.append("correction.org or correction.tex")
|
||||
|
||||
if missing:
|
||||
print(f"Error: Missing files in {folder}: {', '.join(missing)}")
|
||||
sys.exit(1)
|
||||
raise CliError(
|
||||
f"Missing files in {folder}: {', '.join(missing)}",
|
||||
ExitCode.INVALID_WORKSPACE,
|
||||
)
|
||||
|
||||
print("Reading files...")
|
||||
pdf_bytes = pdf_path.read_bytes()
|
||||
enonce_text = enonce_path.read_text(encoding="utf-8")
|
||||
correction_text = correction_path.read_text(encoding="utf-8")
|
||||
|
||||
client = genai.Client(api_key=api_key)
|
||||
if api_client is None:
|
||||
if not api_key:
|
||||
raise CliError("GEMINI_API_KEY is not configured")
|
||||
api_client = genai.Client(api_key=api_key)
|
||||
client = api_client
|
||||
|
||||
# ==========================================
|
||||
# REQUÊTE 1 : Extraction des Énoncés
|
||||
@@ -210,7 +231,7 @@ def process_exam(folder_path: str, restart: bool = False):
|
||||
)
|
||||
response_q_text = response_q.text
|
||||
print("Saving questions to cache...")
|
||||
cache_q_file.write_text(response_q_text, encoding="utf-8")
|
||||
atomic_write_text(cache_q_file, response_q_text)
|
||||
|
||||
questions_data = ExamQuestions.model_validate_json(response_q_text)
|
||||
|
||||
@@ -248,7 +269,7 @@ def process_exam(folder_path: str, restart: bool = False):
|
||||
)
|
||||
response_s_text = response_s.text
|
||||
print("Saving solutions to cache...")
|
||||
cache_s_file.write_text(response_s_text, encoding="utf-8")
|
||||
atomic_write_text(cache_s_file, response_s_text)
|
||||
|
||||
solutions_data = ExamSolutions.model_validate_json(response_s_text)
|
||||
|
||||
@@ -284,7 +305,7 @@ def process_exam(folder_path: str, restart: bool = False):
|
||||
)
|
||||
response_c_text = response_c.text
|
||||
print("Saving context to cache...")
|
||||
cache_c_file.write_text(response_c_text, encoding="utf-8")
|
||||
atomic_write_text(cache_c_file, response_c_text)
|
||||
|
||||
context_data = ExamContext.model_validate_json(response_c_text)
|
||||
|
||||
@@ -388,7 +409,6 @@ def process_exam(folder_path: str, restart: bool = False):
|
||||
for g_indices in q_group_indices:
|
||||
group_items = []
|
||||
first_q_idx = g_indices[0]
|
||||
last_q_idx = g_indices[-1]
|
||||
|
||||
for q_idx in g_indices:
|
||||
q_item = questions_only[q_idx]
|
||||
@@ -497,14 +517,13 @@ def process_exam(folder_path: str, restart: bool = False):
|
||||
|
||||
# 2. Actual Parsing
|
||||
grouped_items = []
|
||||
current_raw_group = [] # Stores (is_context, label_or_flag, content)
|
||||
all_new_q_labels = []
|
||||
|
||||
# Pass 1: Read all edited lines and collect question labels in sequence
|
||||
for line in edited_lines:
|
||||
if line == "---" or " ### " not in line:
|
||||
continue
|
||||
lbl, content_raw = line.split(" ### ", 1)
|
||||
lbl, _content_raw = line.split(" ### ", 1)
|
||||
lbl = lbl.strip()
|
||||
if lbl != "CONTEXT":
|
||||
all_new_q_labels.append(lbl)
|
||||
@@ -606,10 +625,6 @@ def process_exam(folder_path: str, restart: bool = False):
|
||||
validate_windows_labels(labels_list)
|
||||
|
||||
# Save labels and proceed
|
||||
with open(folder / "labels", 'w', encoding='utf-8') as f_labels:
|
||||
for label in labels_list:
|
||||
f_labels.write(f"{label}\n")
|
||||
|
||||
grouped_extraction = GroupedExamExtraction(groups=grouped_items)
|
||||
|
||||
# 2. Setup output directories
|
||||
@@ -628,8 +643,7 @@ def process_exam(folder_path: str, restart: bool = False):
|
||||
).strip().lower()
|
||||
|
||||
if answer not in ("y", "yes"):
|
||||
print("Aborted.")
|
||||
sys.exit(1)
|
||||
raise CliError("Output replacement aborted", ExitCode.INVALID_ARGUMENTS)
|
||||
# Empty each directory
|
||||
for d in dirs:
|
||||
if d.exists():
|
||||
@@ -645,6 +659,7 @@ def process_exam(folder_path: str, restart: bool = False):
|
||||
|
||||
|
||||
print("Writing grouped question and solution files...")
|
||||
processing_errors = []
|
||||
|
||||
for group in grouped_extraction.groups:
|
||||
q_items = [item for item in group if isinstance(item, QuestionItem)]
|
||||
@@ -690,8 +705,9 @@ def process_exam(folder_path: str, restart: bool = False):
|
||||
)
|
||||
rubrics_data = GroupRubrics.model_validate_json(response_r.text)
|
||||
rubrics_map = {r.label: r.rubric_content for r in rubrics_data.rubrics}
|
||||
except Exception as e:
|
||||
except Exception as e: # noqa: BLE001 - remote API boundary
|
||||
print(f"Error generating rubric for group {labels[0]}: {e}")
|
||||
processing_errors.append(str(e))
|
||||
rubrics_map = {}
|
||||
|
||||
# 1. Compute the common prefix for the group
|
||||
@@ -737,7 +753,7 @@ def process_exam(folder_path: str, restart: bool = False):
|
||||
elif isinstance(item, ContextItem):
|
||||
raw_ctx = item.content.strip()
|
||||
tabulated_ctx = "\t" + re.sub(r'\n\s*', '\n\t', raw_ctx)
|
||||
text_content_lines.append(f"CONTEXT :")
|
||||
text_content_lines.append("CONTEXT :")
|
||||
text_content_lines.append(tabulated_ctx)
|
||||
|
||||
# --- Save context to Text2 (Concatenating if exists) ---
|
||||
@@ -762,29 +778,53 @@ def process_exam(folder_path: str, restart: bool = False):
|
||||
# ==========================================
|
||||
all_tex_files = list(text2_dir.glob("*.tex")) + list(sol2_dir.glob("*.tex"))
|
||||
|
||||
def compile_worker(tex_file: Path):
|
||||
def compile_worker(tex_file: Path) -> str | None:
|
||||
"""Helper to read content and call the utility function."""
|
||||
try:
|
||||
content = tex_file.read_text(encoding="utf-8")
|
||||
pdf_path = tex_file.with_suffix(".pdf")
|
||||
compile_to_pdf(content, pdf_path)
|
||||
except Exception as e:
|
||||
print(f"Error compiling {tex_file.name}: {e}")
|
||||
except Exception as e: # noqa: BLE001 - compiler worker boundary
|
||||
return f"Error compiling {tex_file.name}: {e}"
|
||||
return None
|
||||
|
||||
print(f"Compiling {len(all_tex_files)} files to PDF using 4 threads...")
|
||||
with ThreadPoolExecutor(max_workers=4) as executor:
|
||||
executor.map(compile_worker, all_tex_files)
|
||||
compile_errors = [
|
||||
error for error in executor.map(compile_worker, all_tex_files) if error
|
||||
]
|
||||
for error in compile_errors:
|
||||
print(error)
|
||||
processing_errors.extend(compile_errors)
|
||||
atomic_write_text(
|
||||
workspace.labels_file,
|
||||
"".join(f"{label}\n" for label in labels_list),
|
||||
)
|
||||
|
||||
return ExitCode.PARTIAL if processing_errors else ExitCode.SUCCESS
|
||||
|
||||
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
parser = evaluation_parser("Extract exam and solution code via Gemini")
|
||||
parser.add_argument(
|
||||
"--restart",
|
||||
action="store_true",
|
||||
help="Ignore cached Gemini extraction responses",
|
||||
)
|
||||
return parser
|
||||
|
||||
|
||||
def main(argv: Sequence[str] | None = None) -> int:
|
||||
parser = build_parser()
|
||||
return execute(
|
||||
parser,
|
||||
argv,
|
||||
lambda args: process_exam(
|
||||
workspace_from_args(args),
|
||||
restart=args.restart,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if not api_key:
|
||||
print("Error: GEMINI_API_KEY environment variable is not set.")
|
||||
sys.exit(1)
|
||||
|
||||
parser = argparse.ArgumentParser(description="Extract exam and solution code via Gemini.")
|
||||
parser.add_argument("folder", help="Directory containing the exam files")
|
||||
parser.add_argument("--restart", action="store_true", help="Ignore cache files and re-run extraction requests.")
|
||||
|
||||
|
||||
args = parser.parse_args()
|
||||
process_exam(args.folder, restart=args.restart)
|
||||
raise SystemExit(main())
|
||||
|
||||
+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())
|
||||
|
||||
+171
-110
@@ -1,9 +1,15 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import glob
|
||||
import os
|
||||
import re
|
||||
import shutil
|
||||
import sys
|
||||
import tempfile
|
||||
import tkinter as tk
|
||||
import uuid
|
||||
from collections.abc import Sequence
|
||||
from pathlib import Path
|
||||
from tkinter import messagebox
|
||||
|
||||
import fitz # PyMuPDF
|
||||
@@ -11,16 +17,78 @@ from PIL import Image, ImageDraw, ImageTk
|
||||
from pypdf import PdfReader, PdfWriter
|
||||
|
||||
from config import PAGE_SPLITTER_KB
|
||||
from copienator import (
|
||||
CliError,
|
||||
EvaluationWorkspace,
|
||||
ExitCode,
|
||||
execute,
|
||||
target_parser,
|
||||
workspace_from_target,
|
||||
)
|
||||
from platform_utils import launch_pdf_arranger
|
||||
|
||||
|
||||
# --- Constants ---
|
||||
# Conversion factor: 1 cm to points (1 inch = 2.54 cm, 72 points = 1 inch)
|
||||
CM_TO_POINTS = (1 / 2.54) * 72
|
||||
|
||||
def list_pdf_files(directory):
|
||||
l = list(reversed(sorted(glob.glob(os.path.join(directory, "*.pdf")))))
|
||||
return [u for u in l if "enonce" not in u]
|
||||
def list_pdf_files(directory: str | Path) -> list[Path]:
|
||||
paths = sorted(Path(directory).glob("*.pdf"), key=lambda path: path.name.casefold())
|
||||
return [path for path in paths if "enonce" not in path.name.casefold()]
|
||||
|
||||
|
||||
def _temporary_sibling(path: Path, purpose: str) -> Path:
|
||||
return path.with_name(f".{path.name}.{purpose}.{uuid.uuid4().hex}.tmp")
|
||||
|
||||
|
||||
def commit_processed_pdf(
|
||||
workspace: EvaluationWorkspace,
|
||||
original_path: Path,
|
||||
generated_path: Path,
|
||||
) -> Path:
|
||||
"""Commit a processed copy and its original backup with rollback."""
|
||||
backup_path = workspace.original_copies_dir / original_path.name
|
||||
output_path = workspace.copies_dir / original_path.name
|
||||
workspace.original_copies_dir.mkdir(parents=True, exist_ok=True)
|
||||
workspace.copies_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
staged_backup = None
|
||||
if original_path.resolve() != backup_path.resolve():
|
||||
staged_backup = _temporary_sibling(backup_path, "new-original")
|
||||
shutil.copy2(original_path, staged_backup)
|
||||
saved_backup = _temporary_sibling(backup_path, "old-original")
|
||||
saved_output = _temporary_sibling(output_path, "old-output")
|
||||
backup_replaced = False
|
||||
output_replaced = False
|
||||
try:
|
||||
if staged_backup is not None:
|
||||
if backup_path.exists():
|
||||
backup_path.replace(saved_backup)
|
||||
staged_backup.replace(backup_path)
|
||||
backup_replaced = True
|
||||
if output_path.exists():
|
||||
output_path.replace(saved_output)
|
||||
generated_path.replace(output_path)
|
||||
output_replaced = True
|
||||
if original_path.resolve() not in {
|
||||
backup_path.resolve(),
|
||||
output_path.resolve(),
|
||||
}:
|
||||
original_path.unlink()
|
||||
except Exception:
|
||||
if output_replaced and output_path.exists():
|
||||
output_path.unlink()
|
||||
if saved_output.exists():
|
||||
saved_output.replace(output_path)
|
||||
if backup_replaced and backup_path.exists():
|
||||
backup_path.unlink()
|
||||
if saved_backup.exists():
|
||||
saved_backup.replace(backup_path)
|
||||
raise
|
||||
finally:
|
||||
for temporary in (staged_backup, saved_backup, saved_output):
|
||||
if temporary is not None and temporary.exists():
|
||||
temporary.unlink()
|
||||
return output_path
|
||||
|
||||
class PDFPreviewer:
|
||||
|
||||
@@ -30,26 +98,36 @@ class PDFPreviewer:
|
||||
return False
|
||||
self.pdf_path = self.inputs.pop()
|
||||
self.file_rotation = 0
|
||||
self.base_name = os.path.splitext(os.path.basename(self.pdf_path))[0]
|
||||
self.split_dir = f"{self.base_name}_split"
|
||||
self.reorder_dir = f"{self.base_name}_reorder"
|
||||
|
||||
# Create a temporary output file
|
||||
self.final_file = f"{self.base_name}_temp.pdf"
|
||||
self.base_name = self.pdf_path.stem
|
||||
self._temporary_directory = tempfile.TemporaryDirectory(
|
||||
prefix=f".{self.base_name}.page-splitter.",
|
||||
dir=self.workspace.root,
|
||||
)
|
||||
working_dir = Path(self._temporary_directory.name)
|
||||
self.split_dir = working_dir / "split"
|
||||
self.reorder_dir = working_dir / "reorder"
|
||||
self.final_file = working_dir / f"{self.base_name}.pdf"
|
||||
|
||||
self.current_page_index = 0
|
||||
self.page_settings = []
|
||||
self.processing = False # Flag to prevent multiple finish calls
|
||||
try:
|
||||
self.doc = fitz.open(self.pdf_path)
|
||||
except Exception as e:
|
||||
except (OSError, RuntimeError, ValueError) as e:
|
||||
self.failed = True
|
||||
self._temporary_directory.cleanup()
|
||||
messagebox.showerror("Error", f"Failed to open PDF file: {e}")
|
||||
self.master.destroy()
|
||||
return
|
||||
self.master.title(f"PDF Splitter - {os.path.basename(self.pdf_path)}")
|
||||
self.master.title(f"PDF Splitter - {self.pdf_path.name}")
|
||||
return True
|
||||
|
||||
def __init__(self, master, path):
|
||||
def __init__(
|
||||
self,
|
||||
master: tk.Tk,
|
||||
workspace: EvaluationWorkspace,
|
||||
inputs: list[Path],
|
||||
) -> None:
|
||||
"""
|
||||
Initializes the application.
|
||||
|
||||
@@ -57,40 +135,16 @@ class PDFPreviewer:
|
||||
master (tk.Tk): The root Tkinter window.
|
||||
pdf_path (str): The path to the input PDF file.
|
||||
"""
|
||||
if not os.path.exists(path):
|
||||
messagebox.showerror("Error", f"File not found: {path}")
|
||||
master.destroy()
|
||||
return
|
||||
|
||||
if os.path.isdir(path):
|
||||
self.inputs = list_pdf_files(path)
|
||||
else:
|
||||
# Check for existing original in backup and restore if found
|
||||
dir_name = os.path.dirname(os.path.abspath(path))
|
||||
file_name = os.path.basename(path)
|
||||
if os.path.basename(dir_name) == "Copies":
|
||||
dir_name = os.path.dirname(dir_name)
|
||||
path = os.path.join(dir_name, file_name)
|
||||
backup_path = os.path.join(dir_name, "Copies Originales", file_name)
|
||||
|
||||
if os.path.exists(backup_path):
|
||||
try:
|
||||
shutil.move(backup_path, path)
|
||||
print(f"Restored original file from: {backup_path}")
|
||||
except Exception as e:
|
||||
messagebox.showerror("Error", f"Failed to restore original file: {e}")
|
||||
master.destroy()
|
||||
return
|
||||
|
||||
self.inputs = [path]
|
||||
|
||||
self.workspace = workspace
|
||||
self.inputs = inputs
|
||||
self.output_dir = None
|
||||
self.master = master
|
||||
self.num = 0
|
||||
self.global_rotation = 0 # Rotation appliquée à tous les fichiers
|
||||
self.history = []
|
||||
self.failed = False
|
||||
if not self.setup_next_file():
|
||||
print(f"Aucun fichier PDF valide trouvé dans : {path}")
|
||||
print(f"No PDF files found in {workspace.root}")
|
||||
master.destroy()
|
||||
return
|
||||
|
||||
@@ -249,7 +303,7 @@ class PDFPreviewer:
|
||||
# Re-open the file from disk to reset changes (like moved pages)
|
||||
try:
|
||||
self.doc = fitz.open(self.pdf_path)
|
||||
except Exception as e:
|
||||
except (OSError, RuntimeError, ValueError) as e:
|
||||
messagebox.showerror("Error", f"Failed to reopen PDF file: {e}")
|
||||
self.master.destroy()
|
||||
return
|
||||
@@ -328,7 +382,8 @@ class PDFPreviewer:
|
||||
self._initialize_current_page_settings()
|
||||
self.load_page()
|
||||
else:
|
||||
self.finish_and_process()
|
||||
if not self.finish_and_process():
|
||||
return
|
||||
self.history.append(self.pdf_path)
|
||||
if self.setup_next_file():
|
||||
self._initialize_current_page_settings()
|
||||
@@ -336,59 +391,24 @@ class PDFPreviewer:
|
||||
else:
|
||||
self.master.destroy()
|
||||
|
||||
def finish_and_process(self):
|
||||
"""Starts the PDF splitting process and moves files."""
|
||||
self.split_pdf()
|
||||
# print("Debug : ", self.page_settings)
|
||||
# input("Splitting done. Continue ?")
|
||||
self.reorder_pdfs()
|
||||
# input("Reorder done. Continue ?")
|
||||
self.concate_files()
|
||||
|
||||
# Logic to move original to backup and replace with new file
|
||||
def finish_and_process(self) -> bool:
|
||||
"""Render and transactionally install the processed PDF."""
|
||||
try:
|
||||
abs_path = os.path.abspath(self.pdf_path)
|
||||
dir_name = os.path.dirname(abs_path)
|
||||
file_name = os.path.basename(abs_path)
|
||||
|
||||
backup_dir = os.path.join(dir_name, "Copies Originales")
|
||||
copies_dir = os.path.join(dir_name, "Copies")
|
||||
os.makedirs(backup_dir, exist_ok=True)
|
||||
os.makedirs(copies_dir, exist_ok=True)
|
||||
|
||||
backup_path = os.path.join(backup_dir, file_name)
|
||||
copies_path = os.path.join(copies_dir, file_name)
|
||||
|
||||
# Remove backup if it already exists (overwrite)
|
||||
if os.path.exists(backup_path):
|
||||
os.remove(backup_path)
|
||||
|
||||
# Move the original file to "Copies Originales"
|
||||
shutil.move(self.pdf_path, backup_path)
|
||||
|
||||
# Move the temp output file to replace the original
|
||||
shutil.move(self.final_file, copies_path)
|
||||
|
||||
# print(f"Original moved to {backup_path}, new file saved at {self.pdf_path}")
|
||||
|
||||
except Exception as e:
|
||||
messagebox.showerror("Error", f"Failed to move/replace files: {e}")
|
||||
|
||||
self.remove_dirs()
|
||||
|
||||
def _restore_original(self, path):
|
||||
"""Restores the original file from the 'Copies Originales' backup."""
|
||||
dir_name = os.path.dirname(os.path.abspath(path))
|
||||
file_name = os.path.basename(path)
|
||||
backup_path = os.path.join(dir_name, "Copies Originales", file_name)
|
||||
|
||||
if os.path.exists(backup_path):
|
||||
try:
|
||||
# Moving overwrites the generated PDF with the original backup
|
||||
shutil.move(backup_path, path)
|
||||
print(f"Restored original file from: {backup_path}")
|
||||
except Exception as e:
|
||||
print(f"Failed to restore original file: {e}")
|
||||
self.split_pdf()
|
||||
self.reorder_pdfs()
|
||||
self.concate_files()
|
||||
commit_processed_pdf(self.workspace, self.pdf_path, self.final_file)
|
||||
except Exception as exc: # noqa: BLE001 - interactive boundary
|
||||
self.failed = True
|
||||
self.processing = False
|
||||
print(f"Failed to process {self.pdf_path}: {exc}")
|
||||
messagebox.showerror("Error", f"Failed to process PDF: {exc}")
|
||||
self._temporary_directory.cleanup()
|
||||
self.master.destroy()
|
||||
return False
|
||||
else:
|
||||
self._temporary_directory.cleanup()
|
||||
return True
|
||||
|
||||
def go_to_previous_file(self, event=None):
|
||||
"""Goes back to the beginning of the previously completed file."""
|
||||
@@ -398,14 +418,16 @@ class PDFPreviewer:
|
||||
# Close the currently open document to avoid lock issues
|
||||
if hasattr(self, 'doc'):
|
||||
self.doc.close()
|
||||
if hasattr(self, "_temporary_directory"):
|
||||
self._temporary_directory.cleanup()
|
||||
|
||||
# 1. Push current file back onto the stack so it processes next
|
||||
self.inputs.append(self.pdf_path)
|
||||
|
||||
# 2. Get the previous file, restore its original state, and push to stack
|
||||
# 2. Reprocess the previous file from its preserved original backup
|
||||
prev_file = self.history.pop()
|
||||
self._restore_original(prev_file)
|
||||
self.inputs.append(prev_file)
|
||||
backup = self.workspace.original_copies_dir / Path(prev_file).name
|
||||
self.inputs.append(backup if backup.is_file() else Path(prev_file))
|
||||
|
||||
# 3. Reload environment (setup_next_file will pop prev_file back off the stack)
|
||||
self.setup_next_file()
|
||||
@@ -426,13 +448,9 @@ class PDFPreviewer:
|
||||
for pdf in pdf_files:
|
||||
try:
|
||||
os.remove(pdf)
|
||||
except Exception as e:
|
||||
except OSError as e:
|
||||
print(f"Error deleting {pdf}: {e}")
|
||||
|
||||
def remove_dirs(self):
|
||||
shutil.rmtree(self.split_dir)
|
||||
shutil.rmtree(self.reorder_dir)
|
||||
|
||||
def split_pdf(self):
|
||||
"""Splits each page of the PDF according to the saved settings."""
|
||||
print("Starting PDF processing...")
|
||||
@@ -588,13 +606,56 @@ class PDFPreviewer:
|
||||
print(f"Created merged PDF: {self.final_file}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if len(sys.argv) != 2:
|
||||
print("Usage: python script_name.py <path_to_pdf_file>")
|
||||
sys.exit(1)
|
||||
def _selected_inputs(
|
||||
workspace: EvaluationWorkspace,
|
||||
target: Path,
|
||||
) -> list[Path]:
|
||||
if target.is_file():
|
||||
if target.suffix.casefold() != ".pdf":
|
||||
raise CliError(f"Target is not a PDF: {target}", ExitCode.INVALID_ARGUMENTS)
|
||||
backup = workspace.original_copies_dir / target.name
|
||||
return [backup if backup.is_file() else target]
|
||||
|
||||
pdf_file_path = sys.argv[1]
|
||||
directory = target
|
||||
if target == workspace.copies_dir:
|
||||
candidates = list_pdf_files(workspace.copies_dir)
|
||||
candidates = [
|
||||
(
|
||||
workspace.original_copies_dir / path.name
|
||||
if (workspace.original_copies_dir / path.name).is_file()
|
||||
else path
|
||||
)
|
||||
for path in candidates
|
||||
]
|
||||
else:
|
||||
candidates = list_pdf_files(directory)
|
||||
return list(reversed(candidates))
|
||||
|
||||
|
||||
def run(workspace: EvaluationWorkspace, target: Path) -> ExitCode:
|
||||
inputs = _selected_inputs(workspace, target)
|
||||
if not inputs:
|
||||
print(f"No PDF files found in {target}")
|
||||
return ExitCode.SUCCESS
|
||||
root = tk.Tk()
|
||||
app = PDFPreviewer(root, pdf_file_path)
|
||||
application = PDFPreviewer(root, workspace, inputs)
|
||||
root.mainloop()
|
||||
return ExitCode.FAILURE if application.failed else ExitCode.SUCCESS
|
||||
|
||||
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
return target_parser("Interactively split and reorder scanned PDF pages")
|
||||
|
||||
|
||||
def main(argv: Sequence[str] | None = None) -> int:
|
||||
parser = build_parser()
|
||||
|
||||
def handle(args: argparse.Namespace) -> ExitCode:
|
||||
workspace, target = workspace_from_target(args)
|
||||
return run(workspace, target)
|
||||
|
||||
return execute(parser, argv, handle)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
|
||||
+117
-80
@@ -1,24 +1,29 @@
|
||||
import json
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import queue
|
||||
import sys
|
||||
import threading
|
||||
import tkinter as tk
|
||||
from collections.abc import Sequence
|
||||
from pathlib import Path
|
||||
from tkinter import messagebox
|
||||
|
||||
from PIL import Image, ImageDraw, ImageFont, ImageTk
|
||||
|
||||
from copienator import (
|
||||
EvaluationWorkspace,
|
||||
ExitCode,
|
||||
atomic_write_json,
|
||||
execute,
|
||||
read_json,
|
||||
target_parser,
|
||||
workspace_from_target,
|
||||
)
|
||||
from platform_utils import open_path
|
||||
|
||||
print("o to open pdf, O original pdf, e to emacs part, p to go back, i to interro, click for coordinates")
|
||||
from utils import natural_key, read_all_labels
|
||||
|
||||
# --- Configuration & Globals ---
|
||||
padding = 60
|
||||
valid_labels_set = None
|
||||
|
||||
# Queue payload: (pil_image, json_path, metadata)
|
||||
# metadata is a dict: {'copie': str, 'part': int, 'schema': dict}
|
||||
image_queue = queue.Queue(maxsize=5)
|
||||
|
||||
try:
|
||||
font = ImageFont.truetype("DejaVuSans.ttf", size=30)
|
||||
@@ -55,6 +60,14 @@ def convert_list(l, group_id, json_schema):
|
||||
ll.append(ee)
|
||||
return ll
|
||||
|
||||
|
||||
def normalized_labels(entries):
|
||||
return [
|
||||
str(value["label"]).removeprefix("|").removesuffix("|")
|
||||
for value in entries
|
||||
if str(value["label"]) != "_"
|
||||
]
|
||||
|
||||
def prepare_image(image_path: str, bounding_boxes, all_labels, nb_pages, last_label_index):
|
||||
im = Image.open(image_path)
|
||||
im.load()
|
||||
@@ -94,7 +107,7 @@ def prepare_image(image_path: str, bounding_boxes, all_labels, nb_pages, last_la
|
||||
|
||||
# --- Processing Logic (Worker Thread) ---
|
||||
|
||||
def worker_thread(base_dir, files_to_process, all_labels):
|
||||
def _worker_items(base_dir, files_to_process, all_labels, output_queue):
|
||||
"""
|
||||
Iterates through files, prepares VISUALS only, and puts metadata in queue.
|
||||
Does NOT write final JSON files anymore.
|
||||
@@ -111,9 +124,8 @@ def worker_thread(base_dir, files_to_process, all_labels):
|
||||
json_schema_path = base_dir / 'Cutleft' / f"{copie}_schema.json"
|
||||
|
||||
try:
|
||||
with open(json_schema_path, 'r') as f:
|
||||
json_schema = json.load(f)
|
||||
except:
|
||||
json_schema = read_json(json_schema_path)
|
||||
except (OSError, TypeError, ValueError):
|
||||
print("No json_schema : ", json_schema_path)
|
||||
continue
|
||||
|
||||
@@ -124,11 +136,10 @@ def worker_thread(base_dir, files_to_process, all_labels):
|
||||
bb_list = []
|
||||
json_name = ""
|
||||
try:
|
||||
with open(json_path, 'r') as f:
|
||||
json_result = json.load(f)
|
||||
json_result = read_json(json_path)
|
||||
bb_list = json_result.get("list", [])
|
||||
json_name = json_result.get("name", "")
|
||||
except Exception as e:
|
||||
except Exception as e: # noqa: BLE001 - malformed user-editable JSON
|
||||
print(f"Warning: {json_path.name} is malformed! Loading blank. {e}")
|
||||
# We do NOT skip; we continue so the user can fix it in the GUI
|
||||
|
||||
@@ -138,7 +149,7 @@ def worker_thread(base_dir, files_to_process, all_labels):
|
||||
prepare_image(str(img_path), bb_list, all_labels, nb_pages, last_label_index)
|
||||
error_msg = None
|
||||
|
||||
except Exception as e:
|
||||
except Exception as e: # noqa: BLE001 - keep the item editable in the GUI
|
||||
print(f"Error processing {img_path.name}: {e}")
|
||||
pil_image = Image.open(str(img_path))
|
||||
error_msg = str(e)
|
||||
@@ -151,15 +162,24 @@ def worker_thread(base_dir, files_to_process, all_labels):
|
||||
"error": error_msg
|
||||
}
|
||||
|
||||
image_queue.put((pil_image, json_path, metadata))
|
||||
output_queue.put((pil_image, json_path, metadata))
|
||||
|
||||
# Sentinel to indicate finished
|
||||
image_queue.put((None, None, None))
|
||||
def worker_thread(base_dir, files_to_process, all_labels, output_queue):
|
||||
"""Prepare queue items and always terminate the GUI stream."""
|
||||
failure = None
|
||||
try:
|
||||
_worker_items(base_dir, files_to_process, all_labels, output_queue)
|
||||
except Exception as exc: # noqa: BLE001 - worker boundary
|
||||
failure = str(exc)
|
||||
print(f"Plotting worker failed: {exc}")
|
||||
finally:
|
||||
metadata = {"worker_error": failure} if failure else None
|
||||
output_queue.put((None, None, metadata))
|
||||
|
||||
# --- GUI Logic (Main Thread) ---
|
||||
|
||||
class ImageViewer:
|
||||
def __init__(self, root, base_dir):
|
||||
def __init__(self, root, workspace, valid_labels, input_queue):
|
||||
self.root = root
|
||||
self.root.resizable(False, False) # If you resize, coordinates will be wrong
|
||||
|
||||
@@ -171,7 +191,10 @@ class ImageViewer:
|
||||
|
||||
root.geometry(f"+{x}+{y}")
|
||||
|
||||
self.base_dir = base_dir
|
||||
self.workspace = workspace
|
||||
self.base_dir = workspace.root
|
||||
self.valid_labels = valid_labels
|
||||
self.image_queue = input_queue
|
||||
self.root.title("Bounding Box Viewer")
|
||||
self.label = tk.Label(root, text="Waiting for images...")
|
||||
self.label.pack(expand=True, fill="both")
|
||||
@@ -192,6 +215,7 @@ class ImageViewer:
|
||||
self.history = []
|
||||
self.forward_stack = []
|
||||
self.current_pil_image = None
|
||||
self.failed = False
|
||||
|
||||
from config import PLOTTING_KB
|
||||
|
||||
@@ -202,7 +226,8 @@ class ImageViewer:
|
||||
self.root.bind(PLOTTING_KB["open pdf"], self.on_open_pdf)
|
||||
self.root.bind(PLOTTING_KB["open original pdf"], self.on_open_ori_pdf)
|
||||
self.root.bind(PLOTTING_KB["open eval"], self.on_open_interro)
|
||||
self.root.bind('<Escape>', lambda e: self.root.quit())
|
||||
self.root.bind('<Escape>', lambda _event: self.close())
|
||||
self.root.protocol("WM_DELETE_WINDOW", self.close)
|
||||
self.label.bind('<Button-1>', self.on_click)
|
||||
|
||||
self.poll_queue()
|
||||
@@ -214,10 +239,15 @@ class ImageViewer:
|
||||
if self.forward_stack:
|
||||
pil_image, json_path, metadata = self.forward_stack.pop()
|
||||
else:
|
||||
pil_image, json_path, metadata = image_queue.get_nowait()
|
||||
pil_image, json_path, metadata = self.image_queue.get_nowait()
|
||||
|
||||
# Handle End of Stream
|
||||
if pil_image is None:
|
||||
if metadata and metadata.get("worker_error"):
|
||||
self.failed = True
|
||||
messagebox.showerror(
|
||||
"Processing Error", metadata["worker_error"]
|
||||
)
|
||||
self.save_current_batch() # Save any remaining data
|
||||
print("All images processed.")
|
||||
self.root.quit()
|
||||
@@ -241,10 +271,12 @@ class ImageViewer:
|
||||
if self.active_copie_name and self.accumulated_results:
|
||||
main_json_path = self.base_dir / "Copies" / f"{self.active_copie_name}.json"
|
||||
print(f"Writing aggregated result to {main_json_path}")
|
||||
with open(main_json_path, 'w') as f:
|
||||
json.dump(self.accumulated_results, f)
|
||||
atomic_write_json(main_json_path, self.accumulated_results)
|
||||
self.accumulated_results = None
|
||||
|
||||
def close(self):
|
||||
self.root.quit()
|
||||
|
||||
|
||||
def on_previous(self, event):
|
||||
if self.is_viewing and self.history:
|
||||
@@ -288,8 +320,7 @@ class ImageViewer:
|
||||
num_added = 0 # ADD THIS LINE
|
||||
|
||||
try:
|
||||
with open(self.current_json_path, 'r') as f:
|
||||
current_data = json.load(f)
|
||||
current_data = read_json(self.current_json_path)
|
||||
|
||||
# Perform the conversion now, post-edit
|
||||
converted_items = convert_list(
|
||||
@@ -298,11 +329,10 @@ class ImageViewer:
|
||||
self.current_meta["schema"]
|
||||
)
|
||||
|
||||
labels = [v["label"] for v in current_data["list"]]
|
||||
labels = [label for label in labels if label != "_"]
|
||||
labels = [label[1:] for label in labels if label[0] == "|"]
|
||||
labels = [label[:-1] for label in labels if label[-1] == "|"]
|
||||
false_labels = [label for label in labels if label not in valid_labels_set]
|
||||
labels = normalized_labels(current_data["list"])
|
||||
false_labels = [
|
||||
label for label in labels if label not in self.valid_labels
|
||||
]
|
||||
|
||||
if false_labels:
|
||||
msg = f"Wrong label in {self.current_json_path.name}: {false_labels}\n\n\tPlease press 'e' to fix it, then press Enter again."
|
||||
@@ -318,7 +348,7 @@ class ImageViewer:
|
||||
if "name" in current_data and current_data["name"] != "Continued":
|
||||
self.accumulated_results["name"] = current_data["name"]
|
||||
|
||||
except Exception as e:
|
||||
except Exception as e: # noqa: BLE001 - interactive validation boundary
|
||||
# Warn user and STOP (do not advance to next image)
|
||||
msg = f"Error reading {self.current_json_path.name}:\n\n{e}\n\nPlease press 'e' to fix it, then press Enter again."
|
||||
print(msg)
|
||||
@@ -387,51 +417,58 @@ class ImageViewer:
|
||||
self.root.clipboard_clear()
|
||||
self.root.clipboard_append(box_str)
|
||||
|
||||
from utils import natural_key, read_all_labels
|
||||
def _selected_images(workspace: EvaluationWorkspace, target: Path) -> list[Path]:
|
||||
workspace.require_directories("Cutleft", "Copies")
|
||||
if target.is_file():
|
||||
stem = target.stem
|
||||
exact = workspace.cutleft_dir / f"{stem}.jpg"
|
||||
if exact.is_file():
|
||||
return [exact]
|
||||
return sorted(
|
||||
workspace.cutleft_dir.glob(f"{stem}_*.jpg"),
|
||||
key=natural_key,
|
||||
)
|
||||
return sorted(workspace.cutleft_dir.glob("*.jpg"), key=natural_key)
|
||||
|
||||
|
||||
def run(workspace: EvaluationWorkspace, target: Path) -> ExitCode:
|
||||
workspace.require_files("labels")
|
||||
all_labels = read_all_labels(workspace.root)
|
||||
files_to_process = _selected_images(workspace, target)
|
||||
if not files_to_process:
|
||||
print(f"No Cutleft images found for {target}")
|
||||
return ExitCode.PARTIAL
|
||||
|
||||
print(
|
||||
"o to open pdf, O original pdf, e to edit part, p to go back, "
|
||||
"i to open the statement, click for coordinates"
|
||||
)
|
||||
input_queue = queue.Queue(maxsize=5)
|
||||
worker = threading.Thread(
|
||||
target=worker_thread,
|
||||
args=(workspace.root, files_to_process, all_labels, input_queue),
|
||||
daemon=True,
|
||||
)
|
||||
worker.start()
|
||||
root = tk.Tk()
|
||||
application = ImageViewer(root, workspace, set(all_labels), input_queue)
|
||||
root.mainloop()
|
||||
return ExitCode.PARTIAL if application.failed else ExitCode.SUCCESS
|
||||
|
||||
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
return target_parser("Interactively verify detected label coordinates")
|
||||
|
||||
|
||||
def main(argv: Sequence[str] | None = None) -> int:
|
||||
parser = build_parser()
|
||||
|
||||
def handle(args: argparse.Namespace) -> ExitCode:
|
||||
workspace, target = workspace_from_target(args)
|
||||
return run(workspace, target)
|
||||
|
||||
return execute(parser, argv, handle)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if len(sys.argv) < 2:
|
||||
print("Usage: python plotting.py <directory_or_file>")
|
||||
sys.exit(1)
|
||||
|
||||
input_path = Path(sys.argv[1])
|
||||
files_to_process = []
|
||||
|
||||
|
||||
if input_path.is_file():
|
||||
# Correctly identify base_dir if we are in 'Copies' or 'Cutleft'
|
||||
if input_path.parent.name in ["Copies", "Cutleft"]:
|
||||
base_dir = input_path.parent.parent
|
||||
else:
|
||||
base_dir = input_path.parent
|
||||
|
||||
stem = input_path.stem
|
||||
cutleft_dir = base_dir / "Cutleft"
|
||||
img_path = cutleft_dir / f"{stem}.jpg"
|
||||
|
||||
if img_path.exists():
|
||||
files_to_process = [img_path]
|
||||
else:
|
||||
# We're given something like Copie01.pdf, look for its split image parts
|
||||
files_to_process = sorted(list(cutleft_dir.glob(f"{stem}_*.jpg")), key=natural_key)
|
||||
else:
|
||||
base_dir = input_path
|
||||
cutleft_dir = base_dir / "Cutleft"
|
||||
if not cutleft_dir.exists():
|
||||
print(f"Error: {cutleft_dir} does not exist.")
|
||||
sys.exit(1)
|
||||
files_to_process = sorted(cutleft_dir.glob("*.jpg"))
|
||||
|
||||
try:
|
||||
all_labels = read_all_labels(base_dir)
|
||||
except FileNotFoundError:
|
||||
all_labels = []
|
||||
valid_labels_set = set(all_labels)
|
||||
|
||||
t = threading.Thread(target=worker_thread, args=(base_dir, files_to_process, all_labels))
|
||||
t.daemon = True
|
||||
t.start()
|
||||
|
||||
root = tk.Tk()
|
||||
app = ImageViewer(root, base_dir)
|
||||
root.mainloop()
|
||||
raise SystemExit(main())
|
||||
|
||||
+234
-191
@@ -1,223 +1,266 @@
|
||||
import fitz # PyMuPDF
|
||||
from pypdf import PdfWriter
|
||||
from pypdf import PdfReader
|
||||
import os
|
||||
import sys
|
||||
import json
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import shutil
|
||||
from pathlib import Path
|
||||
import tempfile
|
||||
from collections import defaultdict
|
||||
from collections.abc import Sequence
|
||||
from pathlib import Path
|
||||
|
||||
from utils import read_all_labels
|
||||
import fitz
|
||||
from pypdf import PdfReader, PdfWriter
|
||||
|
||||
carreau = 1000 // 38
|
||||
import utils
|
||||
from copienator import (
|
||||
CliError,
|
||||
EvaluationWorkspace,
|
||||
ExitCode,
|
||||
execute,
|
||||
read_json,
|
||||
target_parser,
|
||||
workspace_from_target,
|
||||
)
|
||||
from copienator.filesystem import staged_directory
|
||||
|
||||
SQUARE = 1000 // 38
|
||||
Coordinate = tuple[str, int, int, int, int, int]
|
||||
ParsedCoordinate = tuple[str, str, int, int, int, int, int]
|
||||
|
||||
|
||||
def decode_json(pdf_file):
|
||||
file_path = Path(pdf_file)
|
||||
with open(file_path.with_suffix(".json"), "r") as f:
|
||||
json_result = json.load(f)
|
||||
|
||||
nb_pages = len(PdfReader(file_path).pages)
|
||||
|
||||
bb_list = json_result["list"]
|
||||
name = json_result["name"]
|
||||
column_width = 1000 // nb_pages
|
||||
|
||||
def page_number(b):
|
||||
return ((b[1] + b[3]) // 2) // column_width
|
||||
|
||||
result = []
|
||||
for d in bb_list:
|
||||
(b, label) = d["box_2d"], d["label"]
|
||||
pn = page_number(b)
|
||||
result.append((label, pn, b[0] - carreau, b[2]-carreau, b[1], b[3]))
|
||||
result.sort(key=lambda x: (x[1], x[2]))
|
||||
return (name, result)
|
||||
def decode_json(pdf_file: str | Path) -> tuple[str, list[Coordinate]]:
|
||||
"""Read verified label coordinates associated with one copy PDF."""
|
||||
pdf_path = Path(pdf_file)
|
||||
loaded = read_json(pdf_path.with_suffix(".json"))
|
||||
if not isinstance(loaded, dict):
|
||||
raise TypeError(f"Expected a JSON object for {pdf_path}")
|
||||
boxes = loaded.get("list")
|
||||
if not isinstance(boxes, list):
|
||||
raise TypeError(f"Expected a list of labels for {pdf_path}")
|
||||
page_count = len(PdfReader(pdf_path).pages)
|
||||
if page_count == 0:
|
||||
raise ValueError(f"PDF contains no pages: {pdf_path}")
|
||||
column_width = 1000 // page_count
|
||||
result: list[Coordinate] = []
|
||||
for entry in boxes:
|
||||
if not isinstance(entry, dict):
|
||||
raise TypeError(f"Malformed label entry for {pdf_path}: {entry!r}")
|
||||
box = entry["box_2d"]
|
||||
label = str(entry["label"])
|
||||
page_number = ((box[1] + box[3]) // 2) // column_width
|
||||
result.append(
|
||||
(label, page_number, box[0] - SQUARE, box[2] - SQUARE, box[1], box[3])
|
||||
)
|
||||
result.sort(key=lambda item: (item[1], item[2]))
|
||||
return str(loaded.get("name", "")), result
|
||||
|
||||
|
||||
def split_an_interro(base_dir, input_pdf, coords_list):
|
||||
doc = fitz.open(input_pdf)
|
||||
|
||||
output_dir = base_dir / "Copies" / input_pdf.stem
|
||||
generated_files = set()
|
||||
parts_by_label = defaultdict(list)
|
||||
|
||||
# 1. Parse labels to strip '|' and determine type: L (Left), R (Right), N (Normal)
|
||||
parsed_coords = []
|
||||
for item in coords_list:
|
||||
label, pn, y0, y1, x0, x1 = item
|
||||
def _parse_coordinates(coords_list: list[Coordinate]) -> list[ParsedCoordinate]:
|
||||
parsed: list[ParsedCoordinate] = []
|
||||
for label, page, y0, y1, x0, x1 in coords_list:
|
||||
if label.startswith("|"):
|
||||
c_type, clean_label = "L", label[1:]
|
||||
kind, clean_label = "L", label[1:]
|
||||
elif label.endswith("|"):
|
||||
c_type, clean_label = "R", label[:-1]
|
||||
kind, clean_label = "R", label[:-1]
|
||||
else:
|
||||
c_type, clean_label = "N", label
|
||||
parsed_coords.append((clean_label, c_type, pn, y0, y1, x0, x1))
|
||||
kind, clean_label = "N", label
|
||||
parsed.append((clean_label, kind, page, y0, y1, x0, x1))
|
||||
filtered: list[ParsedCoordinate] = []
|
||||
for item in parsed:
|
||||
if not filtered or item[0] != filtered[-1][0]:
|
||||
filtered.append(item)
|
||||
return filtered
|
||||
|
||||
# 2. Filter consecutive duplicate labels based on the cleaned name
|
||||
filtered_coords = []
|
||||
if parsed_coords:
|
||||
filtered_coords.append(parsed_coords[0])
|
||||
for item in parsed_coords[1:]:
|
||||
if item[0] != filtered_coords[-1][0]:
|
||||
filtered_coords.append(item)
|
||||
coords_list = filtered_coords
|
||||
|
||||
def scale_coord(y, page):
|
||||
"""Scale y from 0–1000 range to PDF points."""
|
||||
page_height = page.rect.height
|
||||
return (y / 1000) * page_height
|
||||
|
||||
def save_cropped_page(doc, page_num, x0, y0, x1, y1, out_path):
|
||||
"""Saves a cropped portion of a page as a new PDF."""
|
||||
page = doc[page_num]
|
||||
rotated_rect = page.rect * page.transformation_matrix
|
||||
visual_crop_rect = fitz.Rect(rotated_rect.x0 + x0, y0, rotated_rect.x0 + x1, y1)
|
||||
unrotated_clip_rect = visual_crop_rect * page.derotation_matrix
|
||||
|
||||
temp_doc = fitz.open()
|
||||
temp_page = temp_doc.new_page(
|
||||
width=visual_crop_rect.width,
|
||||
height=visual_crop_rect.height
|
||||
)
|
||||
temp_page.show_pdf_page(
|
||||
temp_page.rect,
|
||||
doc,
|
||||
page_num,
|
||||
def _save_cropped_page(
|
||||
document: fitz.Document,
|
||||
page_number: int,
|
||||
x0: float,
|
||||
y0: float,
|
||||
x1: float,
|
||||
y1: float,
|
||||
output_path: Path,
|
||||
) -> None:
|
||||
page = document[page_number]
|
||||
rotated_rectangle = page.rect * page.transformation_matrix
|
||||
visual_crop = fitz.Rect(
|
||||
rotated_rectangle.x0 + x0,
|
||||
y0,
|
||||
rotated_rectangle.x0 + x1,
|
||||
y1,
|
||||
)
|
||||
unrotated_clip = visual_crop * page.derotation_matrix
|
||||
cropped = fitz.open()
|
||||
try:
|
||||
target_page = cropped.new_page(width=visual_crop.width, height=visual_crop.height)
|
||||
target_page.show_pdf_page(
|
||||
target_page.rect,
|
||||
document,
|
||||
page_number,
|
||||
rotate=-page.rotation,
|
||||
clip=unrotated_clip_rect
|
||||
clip=unrotated_clip,
|
||||
)
|
||||
temp_doc.save(out_path)
|
||||
temp_doc.close()
|
||||
cropped.save(output_path)
|
||||
finally:
|
||||
cropped.close()
|
||||
|
||||
# Iterate through all labels
|
||||
for idx, (clean_label, c_type, start_page, y_start_raw, y_end_box, x0_raw, x1_raw) in enumerate(coords_list):
|
||||
if clean_label == "_":
|
||||
|
||||
def _render_split_outputs(
|
||||
input_pdf: Path,
|
||||
coords_list: list[Coordinate],
|
||||
staging: Path,
|
||||
) -> set[str]:
|
||||
"""Render every current answer into an otherwise empty staging directory."""
|
||||
document = fitz.open(input_pdf)
|
||||
try:
|
||||
parsed = _parse_coordinates(coords_list)
|
||||
parts_by_label: defaultdict[str, list[Path]] = defaultdict(list)
|
||||
with tempfile.TemporaryDirectory(prefix="copienator-split-") as temp_directory:
|
||||
temporary = Path(temp_directory)
|
||||
for index, item in enumerate(parsed):
|
||||
clean_label, kind, start_page, y_start, _y_end, x0_raw, _x1_raw = item
|
||||
if clean_label == "_":
|
||||
continue
|
||||
if not 0 <= start_page < document.page_count:
|
||||
raise ValueError(
|
||||
f"Invalid page {start_page} for {input_pdf.name}"
|
||||
)
|
||||
end_page = document.page_count - 1
|
||||
end_y = 1000
|
||||
for next_item in parsed[index + 1 :]:
|
||||
_next_label, next_kind, next_page, next_y, *_rest = next_item
|
||||
if (
|
||||
(kind == "L" and next_kind in {"L", "N"})
|
||||
or (kind == "R" and next_kind in {"R", "N"})
|
||||
or kind == "N"
|
||||
):
|
||||
end_page = next_page
|
||||
end_y = min(next_y + int(1.5 * SQUARE), 1000)
|
||||
break
|
||||
|
||||
column_width = 1000 / document.page_count
|
||||
if kind == "L":
|
||||
fraction_x0 = (x0_raw % column_width) / column_width
|
||||
fraction_x1 = 1.0
|
||||
end_y = min(1000, end_y + 40)
|
||||
elif kind == "R":
|
||||
fraction_x0 = 0.0
|
||||
left_labels = [entry for entry in parsed if entry[1] == "L"]
|
||||
if left_labels:
|
||||
closest = min(left_labels, key=lambda entry: abs(entry[3] - y_start))
|
||||
center = (closest[5] + closest[6]) / 2.0
|
||||
fraction_x1 = (center % column_width) / column_width
|
||||
if fraction_x1 <= fraction_x0:
|
||||
fraction_x1 = 1.0
|
||||
else:
|
||||
fraction_x1 = 1.0
|
||||
else:
|
||||
fraction_x0, fraction_x1 = 0.0, 1.0
|
||||
|
||||
for page_number in range(start_page, end_page + 1):
|
||||
page = document[page_number]
|
||||
y0 = (y_start / 1000) * page.rect.height if page_number == start_page else 0
|
||||
y1 = (end_y / 1000) * page.rect.height if page_number == end_page else page.rect.height
|
||||
if y1 <= y0 + 1:
|
||||
continue
|
||||
part_path = temporary / f"part-{index}-{page_number}.pdf"
|
||||
_save_cropped_page(
|
||||
document,
|
||||
page_number,
|
||||
fraction_x0 * page.rect.width,
|
||||
y0,
|
||||
fraction_x1 * page.rect.width,
|
||||
y1,
|
||||
part_path,
|
||||
)
|
||||
parts_by_label[clean_label].append(part_path)
|
||||
|
||||
generated: set[str] = set()
|
||||
for label, parts in parts_by_label.items():
|
||||
filename = f"{label}.pdf"
|
||||
merger = PdfWriter()
|
||||
try:
|
||||
for part in parts:
|
||||
merger.append(part)
|
||||
merger.write(staging / filename)
|
||||
finally:
|
||||
merger.close()
|
||||
generated.add(filename)
|
||||
return generated
|
||||
finally:
|
||||
document.close()
|
||||
|
||||
|
||||
def _preserve_previous_outputs(
|
||||
output_dir: Path,
|
||||
staging: Path,
|
||||
generated_files: set[str],
|
||||
) -> None:
|
||||
if not output_dir.is_dir():
|
||||
return
|
||||
for directory in (path for path in output_dir.iterdir() if path.is_dir()):
|
||||
shutil.copytree(directory, staging / directory.name, dirs_exist_ok=True)
|
||||
missing_dir = staging / "Missing"
|
||||
for item in (path for path in output_dir.iterdir() if path.is_file()):
|
||||
if item.name in generated_files:
|
||||
continue
|
||||
print(f"ALERT: File '{item.name}' not generated. Moving to {missing_dir}")
|
||||
missing_dir.mkdir(exist_ok=True)
|
||||
shutil.copy2(item, missing_dir / item.name)
|
||||
|
||||
temp_parts = []
|
||||
end_page = doc.page_count - 1
|
||||
end_y_target_raw = 1000
|
||||
|
||||
# RULE 2: Determine stopping label
|
||||
for next_item in coords_list[idx + 1:]:
|
||||
n_clean, n_type, n_pn, n_y_start, n_y_end, _, _ = next_item
|
||||
def split_an_interro(
|
||||
workspace: EvaluationWorkspace,
|
||||
input_pdf: Path,
|
||||
coords_list: list[Coordinate],
|
||||
) -> None:
|
||||
"""Regenerate one copy's answers and preserve obsolete ones under Missing."""
|
||||
output_dir = workspace.copies_dir / input_pdf.stem
|
||||
with staged_directory(output_dir) as staging:
|
||||
generated = _render_split_outputs(input_pdf, coords_list, staging)
|
||||
_preserve_previous_outputs(output_dir, staging, generated)
|
||||
|
||||
if c_type == "L":
|
||||
is_stop = (n_type in ("L", "N"))
|
||||
elif c_type == "R":
|
||||
is_stop = (n_type in ("R", "N"))
|
||||
else:
|
||||
is_stop = True # Normal labels stop at anything
|
||||
|
||||
if is_stop:
|
||||
end_page = n_pn
|
||||
# end_y_target_raw = n_y_start
|
||||
# On avait retiré un carreau précédemment inutilement, on le rajoute, plus un demi carreau
|
||||
end_y_target_raw = min(n_y_start + int(1.5 * carreau), 1000)
|
||||
break
|
||||
def _selected_pdfs(workspace: EvaluationWorkspace, target: Path) -> list[Path]:
|
||||
workspace.require_directories("Copies")
|
||||
if target.is_file():
|
||||
if target.suffix.casefold() != ".pdf":
|
||||
raise CliError(f"Target is not a PDF: {target}", ExitCode.INVALID_ARGUMENTS)
|
||||
return [target]
|
||||
return sorted(workspace.copies_dir.glob("*.pdf"), key=lambda path: path.name.casefold())
|
||||
|
||||
# RULES 3 & 4: Calculate horizontal boundaries (0.0 to 1.0 fraction of local page width)
|
||||
col_w = 1000 / doc.page_count
|
||||
if c_type == "L": # |name
|
||||
fraction_x0 = (x0_raw % col_w) / col_w
|
||||
fraction_x1 = 1.0
|
||||
end_y_target_raw = min(1000, end_y_target_raw + 40)
|
||||
elif c_type == "R": # name|
|
||||
fraction_x0 = 0.0
|
||||
# Find the closest 'L' label in y-distance
|
||||
L_labels = [it for it in parsed_coords if it[1] == "L"]
|
||||
if L_labels:
|
||||
closest_L = min(L_labels, key=lambda it: abs(it[3] - y_start_raw))
|
||||
closest_L_x_center = (closest_L[5] + closest_L[6]) / 2.0
|
||||
fraction_x1 = (closest_L_x_center % col_w) / col_w
|
||||
if fraction_x1 <= fraction_x0: fraction_x1 = 1.0 # Fallback
|
||||
else:
|
||||
fraction_x1 = 1.0
|
||||
else: # Normal
|
||||
fraction_x0 = 0.0
|
||||
fraction_x1 = 1.0
|
||||
|
||||
current_p = start_page
|
||||
while current_p <= end_page:
|
||||
page = doc[current_p]
|
||||
def run(workspace: EvaluationWorkspace, target: Path) -> ExitCode:
|
||||
workspace.require_files("labels")
|
||||
utils.read_all_labels(workspace.root)
|
||||
pdf_files = _selected_pdfs(workspace, target)
|
||||
status = ExitCode.SUCCESS
|
||||
for pdf_path in pdf_files:
|
||||
json_path = pdf_path.with_suffix(".json")
|
||||
if not json_path.is_file():
|
||||
print(f"Warning: No JSON found for {pdf_path.name}")
|
||||
status = ExitCode.PARTIAL
|
||||
continue
|
||||
name, coordinates = decode_json(pdf_path)
|
||||
print(f"Decoded name: {name}")
|
||||
split_an_interro(workspace, pdf_path, coordinates)
|
||||
if not pdf_files:
|
||||
print("No PDF copies found.")
|
||||
return status
|
||||
|
||||
y0 = scale_coord(y_start_raw, page) if current_p == start_page else 0
|
||||
y1 = scale_coord(end_y_target_raw, page) if current_p == end_page else page.rect.height
|
||||
|
||||
if y1 > y0 + 1:
|
||||
# Convert fractions to absolute PDF points
|
||||
x0_pdf = fraction_x0 * page.rect.width
|
||||
x1_pdf = fraction_x1 * page.rect.width
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
return target_parser("Split verified PDF copies into answers by label")
|
||||
|
||||
temp_path = f"_part_{idx}_{current_p}.pdf"
|
||||
save_cropped_page(doc, current_p, x0_pdf, y0, x1_pdf, y1, temp_path)
|
||||
temp_parts.append(temp_path)
|
||||
|
||||
current_p += 1
|
||||
def main(argv: Sequence[str] | None = None) -> int:
|
||||
parser = build_parser()
|
||||
|
||||
parts_by_label[clean_label].extend(temp_parts)
|
||||
def handle(args: argparse.Namespace) -> ExitCode:
|
||||
workspace, target = workspace_from_target(args)
|
||||
return run(workspace, target)
|
||||
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Process aggregated parts by label
|
||||
for title, parts in parts_by_label.items():
|
||||
merger = PdfWriter()
|
||||
for part in parts:
|
||||
if os.path.exists(part):
|
||||
merger.append(part)
|
||||
|
||||
filename = f"{title}.pdf"
|
||||
merger.write(output_dir / filename)
|
||||
merger.close()
|
||||
generated_files.add(filename)
|
||||
|
||||
# Cleanup
|
||||
for part in parts:
|
||||
if os.path.exists(part):
|
||||
os.remove(part)
|
||||
|
||||
doc.close()
|
||||
|
||||
# Move files not generated in this run to 'Missing' folder
|
||||
if output_dir.exists():
|
||||
missing_dir = output_dir / "Missing"
|
||||
for item in output_dir.iterdir():
|
||||
if item.is_file() and item.name not in generated_files:
|
||||
print(f"ALERT: File '{item.name}' not generated. Moving to {missing_dir}")
|
||||
missing_dir.mkdir(exist_ok=True)
|
||||
item.rename(missing_dir / item.name)
|
||||
return execute(parser, argv, handle)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if len(sys.argv) < 2:
|
||||
print("Usage: python script.py <directory or pdf_file>")
|
||||
sys.exit(1)
|
||||
|
||||
input_arg = Path(sys.argv[1])
|
||||
|
||||
if input_arg.is_file():
|
||||
base_dir = input_arg.parent
|
||||
if base_dir.name == "Copies":
|
||||
base_dir = base_dir.parent
|
||||
pdf_files = [input_arg]
|
||||
elif input_arg.is_dir():
|
||||
base_dir = input_arg
|
||||
copies_dir = base_dir / "Copies"
|
||||
pdf_files = sorted(copies_dir.glob("*.pdf"))
|
||||
else:
|
||||
print(f"Error: {input_arg} is not a valid file or directory.")
|
||||
sys.exit(1)
|
||||
|
||||
read_all_labels(base_dir)
|
||||
|
||||
for pdf_path in pdf_files:
|
||||
json_path = pdf_path.with_suffix(".json")
|
||||
# print("Debug :", json_path)
|
||||
if json_path.exists():
|
||||
(name, coords) = decode_json(pdf_path)
|
||||
print("Decoded name : ", name)
|
||||
split_an_interro(base_dir, pdf_path, coords)
|
||||
else:
|
||||
print(f"Warning: No JSON found for {pdf_path.name}")
|
||||
raise SystemExit(main())
|
||||
|
||||
+76
-57
@@ -1,77 +1,96 @@
|
||||
import os
|
||||
import sys
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
from pathlib import Path
|
||||
from collections.abc import Sequence
|
||||
|
||||
from google import genai
|
||||
from google.genai import types
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Upload JSONL files and create Gemini Batch jobs.")
|
||||
parser.add_argument("root_dir", type=str, help="Root directory containing the batch JSONL files")
|
||||
args = parser.parse_args()
|
||||
import config
|
||||
from copienator import (
|
||||
CliError,
|
||||
EvaluationWorkspace,
|
||||
ExitCode,
|
||||
atomic_write_json,
|
||||
evaluation_parser,
|
||||
execute,
|
||||
read_json,
|
||||
workspace_from_args,
|
||||
)
|
||||
|
||||
root_dir = Path(args.root_dir)
|
||||
|
||||
if "GEMINI_API_KEY" not in os.environ:
|
||||
sys.exit("Error: GEMINI_API_KEY environment variable not set.")
|
||||
|
||||
client = genai.Client()
|
||||
|
||||
# Define the batch files and their corresponding models
|
||||
batches_to_create = [
|
||||
{
|
||||
"file_path": root_dir / "batch_requests_flash.jsonl",
|
||||
"model_id": "gemini-3-flash-preview",
|
||||
"display_name": f"flash-correction-{root_dir.name}"
|
||||
},
|
||||
{
|
||||
"file_path": root_dir / "batch_requests_pro.jsonl",
|
||||
"model_id": "gemini-3.1-pro-preview",
|
||||
"display_name": f"pro-correction-{root_dir.name}"
|
||||
}
|
||||
]
|
||||
|
||||
for batch in batches_to_create:
|
||||
file_path = batch["file_path"]
|
||||
model_id = batch["model_id"]
|
||||
display_name = batch["display_name"]
|
||||
|
||||
# Check if the file exists
|
||||
if not file_path.exists():
|
||||
def run(workspace: EvaluationWorkspace, *, client=None) -> ExitCode:
|
||||
if client is None:
|
||||
if not config.API_KEY:
|
||||
raise CliError("GEMINI_API_KEY is not configured")
|
||||
client = genai.Client(api_key=config.API_KEY)
|
||||
batches = (
|
||||
(
|
||||
"flash",
|
||||
workspace.root / "batch_requests_flash.jsonl",
|
||||
config.MODEL_FLASH_ID,
|
||||
f"flash-correction-{workspace.name}",
|
||||
),
|
||||
(
|
||||
"pro",
|
||||
workspace.root / "batch_requests_pro.jsonl",
|
||||
config.MODEL_PRO_ID,
|
||||
f"pro-correction-{workspace.name}",
|
||||
),
|
||||
)
|
||||
manifest = {
|
||||
"version": 1,
|
||||
"evaluation": workspace.name,
|
||||
"jobs": {},
|
||||
}
|
||||
if workspace.batch_jobs_file.is_file():
|
||||
previous = read_json(workspace.batch_jobs_file)
|
||||
if isinstance(previous, dict) and isinstance(previous.get("jobs"), dict):
|
||||
manifest["jobs"] = previous["jobs"]
|
||||
started = 0
|
||||
for tier, file_path, model_id, display_name in batches:
|
||||
if not file_path.is_file():
|
||||
print(f"Skipping {model_id}: {file_path.name} does not exist.")
|
||||
continue
|
||||
|
||||
# Check if the file is empty (e.g., if all tasks went to Flash, Pro might be empty)
|
||||
if file_path.stat().st_size == 0:
|
||||
print(f"Skipping {model_id}: {file_path.name} is empty.")
|
||||
continue
|
||||
|
||||
print(f"Processing {file_path.name} for model {model_id}...")
|
||||
|
||||
# 1. Upload the file to the File API
|
||||
print(f" Uploading file...")
|
||||
uploaded_file = client.files.upload(
|
||||
print(f"Uploading {file_path.name} for model {model_id}...")
|
||||
uploaded = client.files.upload(
|
||||
file=str(file_path),
|
||||
config=types.UploadFileConfig(
|
||||
display_name=f"{display_name}-input",
|
||||
mime_type='jsonl'
|
||||
)
|
||||
mime_type="jsonl",
|
||||
),
|
||||
)
|
||||
print(f" Uploaded successfully! File ID: {uploaded_file.name}")
|
||||
|
||||
# 2. Create the batch job
|
||||
print(f" Starting batch job...")
|
||||
batch_job = client.batches.create(
|
||||
job = client.batches.create(
|
||||
model=model_id,
|
||||
src=uploaded_file.name,
|
||||
config={
|
||||
'display_name': display_name,
|
||||
},
|
||||
src=uploaded.name,
|
||||
config={"display_name": display_name},
|
||||
)
|
||||
print(f" Success! Batch Job Name: {batch_job.name}\n")
|
||||
started += 1
|
||||
manifest["jobs"][tier] = {
|
||||
"name": job.name,
|
||||
"display_name": display_name,
|
||||
"model": model_id,
|
||||
"request_file": file_path.name,
|
||||
}
|
||||
atomic_write_json(workspace.batch_jobs_file, manifest)
|
||||
print(f"Started batch job: {job.name}")
|
||||
if not started:
|
||||
print("No non-empty batch request files were found.")
|
||||
return ExitCode.PARTIAL
|
||||
return ExitCode.SUCCESS
|
||||
|
||||
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
return evaluation_parser("Upload correction JSONL files and start Gemini batches")
|
||||
|
||||
|
||||
def main(argv: Sequence[str] | None = None) -> int:
|
||||
parser = build_parser()
|
||||
return execute(parser, argv, lambda args: run(workspace_from_args(args)))
|
||||
|
||||
print("-" * 50)
|
||||
print("All batch jobs have been initiated.")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
raise SystemExit(main())
|
||||
|
||||
+583
-1
@@ -2,6 +2,7 @@ from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import queue
|
||||
import sys
|
||||
@@ -11,7 +12,8 @@ import unittest
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from contextlib import redirect_stderr
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
from PIL import Image
|
||||
from pypdf import PdfReader, PdfWriter
|
||||
@@ -21,6 +23,7 @@ from copienator import (
|
||||
WorkspaceNotFoundError,
|
||||
WorkspaceValidationError,
|
||||
atomic_update_json,
|
||||
atomic_write_bytes,
|
||||
atomic_write_json,
|
||||
read_json,
|
||||
workspace_from_target,
|
||||
@@ -114,6 +117,12 @@ class WorkspaceTests(unittest.TestCase):
|
||||
|
||||
|
||||
class AtomicJsonTests(unittest.TestCase):
|
||||
def test_atomic_binary_round_trip(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
path = Path(directory) / "result.jsonl"
|
||||
atomic_write_bytes(path, b'{"one":1}\n')
|
||||
self.assertEqual(path.read_bytes(), b'{"one":1}\n')
|
||||
|
||||
def test_atomic_round_trip_and_unicode(self) -> None:
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
path = Path(directory) / "nested" / "state.json"
|
||||
@@ -302,6 +311,29 @@ class StandardCliTests(unittest.TestCase):
|
||||
"reading_grouped_annotations": load_script_module(
|
||||
"reading_grouped_annotations.py", "reading_grouped_annotations"
|
||||
),
|
||||
"cutleft": load_script_module("cutleft.py", "cutleft"),
|
||||
"splitting_int": load_script_module(
|
||||
"splitting_int.py", "splitting_int"
|
||||
),
|
||||
"page_splitter": load_script_module(
|
||||
"page_splitter.py", "page_splitter"
|
||||
),
|
||||
"plotting": load_script_module("plotting.py", "plotting"),
|
||||
"gemini_for_labels": load_script_module(
|
||||
"gemini_for_labels.py", "gemini_for_labels"
|
||||
),
|
||||
"gemini_for_enonce": load_script_module(
|
||||
"gemini_for_enonce.py", "gemini_for_enonce"
|
||||
),
|
||||
"enonce_info": load_script_module("enonce_info.py", "enonce_info"),
|
||||
"correction": load_script_module("correction.py", "correction"),
|
||||
"submit_batches": load_script_module(
|
||||
"submit_batches.py", "submit_batches"
|
||||
),
|
||||
"batch_status": load_script_module("batch_status.py", "batch_status"),
|
||||
"fetch_batched_results": load_script_module(
|
||||
"fetch_batched_results.py", "fetch_batched_results"
|
||||
),
|
||||
"copies_tools": load_script_module(
|
||||
"copies_tools.py", "copienator_copies_tools_test"
|
||||
),
|
||||
@@ -339,6 +371,16 @@ class StandardCliTests(unittest.TestCase):
|
||||
"annotating_by_label": [missing],
|
||||
"reading_annotations": [missing],
|
||||
"reading_grouped_annotations": [missing],
|
||||
"cutleft": [missing],
|
||||
"splitting_int": [missing],
|
||||
"page_splitter": [missing],
|
||||
"plotting": [missing],
|
||||
"gemini_for_labels": [missing],
|
||||
"gemini_for_enonce": [missing],
|
||||
"enonce_info": [missing],
|
||||
"correction": [missing],
|
||||
"submit_batches": [missing],
|
||||
"fetch_batched_results": [missing],
|
||||
}
|
||||
for name, arguments in invocations.items():
|
||||
with self.subTest(script=name), redirect_stderr(io.StringIO()):
|
||||
@@ -383,6 +425,16 @@ class StandardCliTests(unittest.TestCase):
|
||||
steps = {step.id: step for step in build_workflow(True)}
|
||||
evaluation = "Evaluation with spaces"
|
||||
cases = {
|
||||
"gemini_for_enonce": (
|
||||
"statement",
|
||||
"gemini",
|
||||
{"target": evaluation, "restart": True},
|
||||
),
|
||||
"enonce_info": (
|
||||
"statement",
|
||||
"personal",
|
||||
{"target": evaluation},
|
||||
),
|
||||
"export": ("export", "default", {"target": evaluation, "refaire": True}),
|
||||
"import": ("import", "default", {"target": evaluation, "refaire": True}),
|
||||
"giving_names": (
|
||||
@@ -431,6 +483,46 @@ class StandardCliTests(unittest.TestCase):
|
||||
"grouped",
|
||||
{"target": evaluation, "update_score": True, "refaire": True},
|
||||
),
|
||||
"cutleft": (
|
||||
"cutleft",
|
||||
"default",
|
||||
{"target": evaluation, "fullpage": True},
|
||||
),
|
||||
"splitting_int": (
|
||||
"splitting",
|
||||
"default",
|
||||
{"target": evaluation},
|
||||
),
|
||||
"page_splitter": (
|
||||
"page_splitter",
|
||||
"default",
|
||||
{"target": evaluation},
|
||||
),
|
||||
"plotting": (
|
||||
"plotting",
|
||||
"default",
|
||||
{"target": evaluation},
|
||||
),
|
||||
"gemini_for_labels": (
|
||||
"labels",
|
||||
"default",
|
||||
{"target": evaluation, "overwrite": True},
|
||||
),
|
||||
"correction": (
|
||||
"correction",
|
||||
"live",
|
||||
{"target": evaluation, "overwrite": True, "limit": 5},
|
||||
),
|
||||
"submit_batches": (
|
||||
"submit_batches",
|
||||
"default",
|
||||
{"target": evaluation},
|
||||
),
|
||||
"fetch_batched_results": (
|
||||
"fetch_batches",
|
||||
"default",
|
||||
{"target": evaluation},
|
||||
),
|
||||
}
|
||||
for module_name, (step_id, variant_id, values) in cases.items():
|
||||
step = steps[step_id]
|
||||
@@ -458,6 +550,21 @@ class StandardCliTests(unittest.TestCase):
|
||||
self.assertEqual(parsed.operation, step_id)
|
||||
self.assertEqual(str(parsed.evaluation), evaluation)
|
||||
|
||||
def test_enonce_info_preserves_labels_when_no_blocks_are_found(self) -> None:
|
||||
module = self.modules["enonce_info"]
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
evaluation = Path(directory) / "Exam"
|
||||
evaluation.mkdir()
|
||||
(evaluation / "source.tex").write_text(
|
||||
"No SHEETINFO blocks\n", encoding="utf-8"
|
||||
)
|
||||
(evaluation / "labels").write_text("Existing\n", encoding="utf-8")
|
||||
self.assertEqual(
|
||||
module.process_directory(EvaluationWorkspace(evaluation)), 4
|
||||
)
|
||||
self.assertEqual((evaluation / "labels").read_text(), "Existing\n")
|
||||
self.assertFalse(list(evaluation.glob(".labels.*.tmp")))
|
||||
|
||||
def test_export_main_copies_outputs(self) -> None:
|
||||
module = self.modules["export"]
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
@@ -513,6 +620,481 @@ class StandardCliTests(unittest.TestCase):
|
||||
self.assertEqual(module.main([str(evaluation)]), 0)
|
||||
self.assertTrue((evaluation / "Par label").is_dir())
|
||||
|
||||
def test_cutleft_can_target_one_copy_without_rendering_at_import(self) -> None:
|
||||
module = self.modules["cutleft"]
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
evaluation = Path(directory) / "Exam"
|
||||
copy_pdf = evaluation / "Copies" / "Copie01.pdf"
|
||||
copy_pdf.parent.mkdir(parents=True)
|
||||
copy_pdf.write_bytes(b"pdf")
|
||||
with patch.object(module, "ImageReviewer") as reviewer:
|
||||
self.assertEqual(module.main([str(copy_pdf), "--fullpage"]), 0)
|
||||
files, output_dir = reviewer.call_args.args[:2]
|
||||
self.assertEqual(files, [copy_pdf])
|
||||
self.assertEqual(output_dir, evaluation / "Cutleft")
|
||||
self.assertEqual(reviewer.call_args.kwargs["default_max_per_file"], 1)
|
||||
|
||||
def test_cutleft_atomic_save_removes_only_obsolete_copy_outputs(self) -> None:
|
||||
module = self.modules["cutleft"]
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
output_dir = Path(directory) / "Cutleft"
|
||||
output_dir.mkdir()
|
||||
(output_dir / "Copie01_01.jpg").write_bytes(b"old-one")
|
||||
(output_dir / "Copie01_02.jpg").write_bytes(b"old-two")
|
||||
(output_dir / "Copie02_01.jpg").write_bytes(b"other-copy")
|
||||
(output_dir / "Copie01_schema.json").write_text(
|
||||
"{}", encoding="utf-8"
|
||||
)
|
||||
image = Image.new("RGB", (5, 5), "white")
|
||||
result = (
|
||||
image,
|
||||
[image],
|
||||
{
|
||||
"original_filename": "Copie01.pdf",
|
||||
"total_pages": 1,
|
||||
"number_of_files": 1,
|
||||
"columns_per_file": [1],
|
||||
},
|
||||
)
|
||||
module.save_results(result, Path("Copie01.pdf"), output_dir)
|
||||
self.assertTrue((output_dir / "Copie01_01.jpg").is_file())
|
||||
self.assertFalse((output_dir / "Copie01_02.jpg").exists())
|
||||
self.assertEqual(
|
||||
(output_dir / "Copie02_01.jpg").read_bytes(), b"other-copy"
|
||||
)
|
||||
self.assertEqual(
|
||||
read_json(output_dir / "Copie01_schema.json")["total_pages"], 1
|
||||
)
|
||||
|
||||
def test_splitting_missing_json_is_partial(self) -> None:
|
||||
module = self.modules["splitting_int"]
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
evaluation = Path(directory) / "Exam"
|
||||
copy_pdf = evaluation / "Copies" / "Copie01.pdf"
|
||||
copy_pdf.parent.mkdir(parents=True)
|
||||
copy_pdf.write_bytes(b"pdf")
|
||||
(evaluation / "labels").write_text("Ex 1\n", encoding="utf-8")
|
||||
self.assertEqual(module.main([str(evaluation)]), 4)
|
||||
|
||||
def test_splitting_failure_preserves_previous_copy_outputs(self) -> None:
|
||||
module = self.modules["splitting_int"]
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
evaluation = Path(directory) / "Exam"
|
||||
copy_pdf = evaluation / "Copies" / "Copie01.pdf"
|
||||
output_dir = evaluation / "Copies" / "Copie01"
|
||||
output_dir.mkdir(parents=True)
|
||||
copy_pdf.write_bytes(b"pdf")
|
||||
(output_dir / "sentinel.pdf").write_bytes(b"old")
|
||||
workspace = EvaluationWorkspace(evaluation)
|
||||
with patch.object(
|
||||
module,
|
||||
"_render_split_outputs",
|
||||
side_effect=RuntimeError("render failed"),
|
||||
), self.assertRaises(RuntimeError):
|
||||
module.split_an_interro(workspace, copy_pdf, [])
|
||||
self.assertEqual((output_dir / "sentinel.pdf").read_bytes(), b"old")
|
||||
|
||||
def test_splitting_moves_obsolete_outputs_to_missing_on_commit(self) -> None:
|
||||
module = self.modules["splitting_int"]
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
evaluation = Path(directory) / "Exam"
|
||||
copy_pdf = evaluation / "Copies" / "Copie01.pdf"
|
||||
output_dir = evaluation / "Copies" / "Copie01"
|
||||
missing_dir = output_dir / "Missing"
|
||||
missing_dir.mkdir(parents=True)
|
||||
copy_pdf.write_bytes(b"pdf")
|
||||
(output_dir / "Old.pdf").write_bytes(b"obsolete")
|
||||
(missing_dir / "Earlier.pdf").write_bytes(b"earlier")
|
||||
|
||||
def render(_pdf, _coordinates, staging):
|
||||
(staging / "Ex 1.pdf").write_bytes(b"new")
|
||||
return {"Ex 1.pdf"}
|
||||
|
||||
with patch.object(module, "_render_split_outputs", side_effect=render):
|
||||
module.split_an_interro(EvaluationWorkspace(evaluation), copy_pdf, [])
|
||||
self.assertEqual((output_dir / "Ex 1.pdf").read_bytes(), b"new")
|
||||
self.assertEqual(
|
||||
(output_dir / "Missing" / "Old.pdf").read_bytes(), b"obsolete"
|
||||
)
|
||||
self.assertEqual(
|
||||
(output_dir / "Missing" / "Earlier.pdf").read_bytes(), b"earlier"
|
||||
)
|
||||
|
||||
def test_splitting_renders_a_real_one_page_answer(self) -> None:
|
||||
module = self.modules["splitting_int"]
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
evaluation = Path(directory) / "Exam"
|
||||
copy_pdf = evaluation / "Copies" / "Copie01.pdf"
|
||||
copy_pdf.parent.mkdir(parents=True)
|
||||
document = module.fitz.open()
|
||||
page = document.new_page(width=600, height=800)
|
||||
page.insert_text((100, 200), "Student answer")
|
||||
document.save(copy_pdf)
|
||||
document.close()
|
||||
(evaluation / "labels").write_text("Ex 1\n", encoding="utf-8")
|
||||
atomic_write_json(
|
||||
copy_pdf.with_suffix(".json"),
|
||||
{
|
||||
"name": "Copie01",
|
||||
"list": [{"label": "Ex 1", "box_2d": [100, 100, 300, 300]}],
|
||||
},
|
||||
)
|
||||
|
||||
self.assertEqual(module.main([str(evaluation)]), 0)
|
||||
answer = evaluation / "Copies" / "Copie01" / "Ex 1.pdf"
|
||||
self.assertTrue(answer.is_file())
|
||||
self.assertEqual(len(PdfReader(answer).pages), 1)
|
||||
|
||||
def test_page_splitter_commits_original_and_generated_copy(self) -> None:
|
||||
module = self.modules["page_splitter"]
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
evaluation = Path(directory) / "Exam"
|
||||
evaluation.mkdir()
|
||||
original = evaluation / "Copie01.pdf"
|
||||
generated = evaluation / ".generated.pdf"
|
||||
original.write_bytes(b"original")
|
||||
generated.write_bytes(b"processed")
|
||||
workspace = EvaluationWorkspace(evaluation)
|
||||
|
||||
output = module.commit_processed_pdf(workspace, original, generated)
|
||||
self.assertEqual(output, evaluation / "Copies" / "Copie01.pdf")
|
||||
self.assertEqual(output.read_bytes(), b"processed")
|
||||
self.assertEqual(
|
||||
(evaluation / "Copies Originales" / "Copie01.pdf").read_bytes(),
|
||||
b"original",
|
||||
)
|
||||
self.assertFalse(original.exists())
|
||||
|
||||
def test_page_splitter_commit_failure_rolls_back_both_files(self) -> None:
|
||||
module = self.modules["page_splitter"]
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
evaluation = Path(directory) / "Exam"
|
||||
original = evaluation / "Copie01.pdf"
|
||||
backup = evaluation / "Copies Originales" / "Copie01.pdf"
|
||||
output = evaluation / "Copies" / "Copie01.pdf"
|
||||
backup.parent.mkdir(parents=True)
|
||||
output.parent.mkdir()
|
||||
original.write_bytes(b"new original")
|
||||
backup.write_bytes(b"previous original")
|
||||
output.write_bytes(b"previous output")
|
||||
missing_generated = evaluation / "missing-generated.pdf"
|
||||
|
||||
with self.assertRaises(FileNotFoundError):
|
||||
module.commit_processed_pdf(
|
||||
EvaluationWorkspace(evaluation), original, missing_generated
|
||||
)
|
||||
self.assertEqual(original.read_bytes(), b"new original")
|
||||
self.assertEqual(backup.read_bytes(), b"previous original")
|
||||
self.assertEqual(output.read_bytes(), b"previous output")
|
||||
|
||||
def test_page_splitter_reprocesses_from_preserved_original(self) -> None:
|
||||
module = self.modules["page_splitter"]
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
evaluation = Path(directory) / "Exam"
|
||||
generated = evaluation / "Copies" / "Copie01.pdf"
|
||||
original = evaluation / "Copies Originales" / "Copie01.pdf"
|
||||
generated.parent.mkdir(parents=True)
|
||||
original.parent.mkdir()
|
||||
generated.write_bytes(b"processed")
|
||||
original.write_bytes(b"original")
|
||||
workspace = EvaluationWorkspace(evaluation)
|
||||
self.assertEqual(
|
||||
module._selected_inputs(workspace, generated),
|
||||
[original],
|
||||
)
|
||||
|
||||
def test_plotting_batch_save_is_atomic(self) -> None:
|
||||
module = self.modules["plotting"]
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
evaluation = Path(directory) / "Exam"
|
||||
(evaluation / "Copies").mkdir(parents=True)
|
||||
viewer = module.ImageViewer.__new__(module.ImageViewer)
|
||||
viewer.base_dir = evaluation
|
||||
viewer.active_copie_name = "Copie01"
|
||||
viewer.accumulated_results = {
|
||||
"name": "Test",
|
||||
"list": [{"label": "Ex 1"}],
|
||||
}
|
||||
viewer.save_current_batch()
|
||||
self.assertEqual(
|
||||
read_json(evaluation / "Copies" / "Copie01.json"),
|
||||
{"name": "Test", "list": [{"label": "Ex 1"}]},
|
||||
)
|
||||
self.assertIsNone(viewer.accumulated_results)
|
||||
|
||||
def test_plotting_validates_plain_and_directional_labels(self) -> None:
|
||||
module = self.modules["plotting"]
|
||||
self.assertEqual(
|
||||
module.normalized_labels(
|
||||
[
|
||||
{"label": "Ex 1"},
|
||||
{"label": "|Ex 2"},
|
||||
{"label": "Ex 3|"},
|
||||
{"label": "_"},
|
||||
]
|
||||
),
|
||||
["Ex 1", "Ex 2", "Ex 3"],
|
||||
)
|
||||
|
||||
def test_plotting_worker_failure_still_terminates_its_queue(self) -> None:
|
||||
module = self.modules["plotting"]
|
||||
output_queue = queue.Queue()
|
||||
with patch.object(
|
||||
module, "_worker_items", side_effect=RuntimeError("worker failed")
|
||||
):
|
||||
module.worker_thread(Path("."), [], [], output_queue)
|
||||
image, json_path, metadata = output_queue.get_nowait()
|
||||
self.assertIsNone(image)
|
||||
self.assertIsNone(json_path)
|
||||
self.assertEqual(metadata, {"worker_error": "worker failed"})
|
||||
|
||||
def test_label_detection_resolves_copy_and_cutleft_targets(self) -> None:
|
||||
module = self.modules["gemini_for_labels"]
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
evaluation = Path(directory) / "Exam"
|
||||
copy_pdf = evaluation / "Copies" / "Copie01.pdf"
|
||||
image_one = evaluation / "Cutleft" / "Copie01_01.jpg"
|
||||
image_two = evaluation / "Cutleft" / "Copie01_02.jpg"
|
||||
copy_pdf.parent.mkdir(parents=True)
|
||||
image_one.parent.mkdir()
|
||||
copy_pdf.write_bytes(b"pdf")
|
||||
image_one.write_bytes(b"image")
|
||||
image_two.write_bytes(b"image")
|
||||
workspace = EvaluationWorkspace(evaluation)
|
||||
|
||||
images, warnings = module.selected_images(workspace, [copy_pdf])
|
||||
self.assertEqual(images, [image_one, image_two])
|
||||
self.assertEqual(warnings, [])
|
||||
images, warnings = module.selected_images(workspace, [image_two])
|
||||
self.assertEqual(images, [image_two])
|
||||
self.assertEqual(warnings, [])
|
||||
|
||||
def test_label_detection_preserves_context_and_writes_atomically(self) -> None:
|
||||
module = self.modules["gemini_for_labels"]
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
evaluation = Path(directory) / "Exam"
|
||||
copies = evaluation / "Copies"
|
||||
cutleft = evaluation / "Cutleft"
|
||||
copies.mkdir(parents=True)
|
||||
cutleft.mkdir()
|
||||
first = cutleft / "Copie01_01.jpg"
|
||||
second = cutleft / "Copie01_02.jpg"
|
||||
first.write_bytes(b"first image")
|
||||
second.write_bytes(b"second image")
|
||||
atomic_write_json(
|
||||
copies / "Copie01_01.json",
|
||||
{
|
||||
"name": "Student",
|
||||
"list": [{"box_2d": [1, 2, 3, 4], "label": "Ex 1"}],
|
||||
},
|
||||
)
|
||||
response = Mock(
|
||||
text=(
|
||||
'{"name":"Continued","list":'
|
||||
'[{"box_2d":[10,20,30,40],"label":"Ex 2"}]}'
|
||||
)
|
||||
)
|
||||
client = Mock()
|
||||
client.models.generate_content.return_value = response
|
||||
|
||||
generated = module.process_copy_group(
|
||||
EvaluationWorkspace(evaluation),
|
||||
"Copie01",
|
||||
[first, second],
|
||||
client=client,
|
||||
labels_text="Ex 1\nEx 2\n",
|
||||
names_text="Student\n",
|
||||
valid_labels={"Ex 1", "Ex 2"},
|
||||
valid_names={"Student", "Continued", "Unknown"},
|
||||
overwrite=False,
|
||||
sleep=lambda _seconds: None,
|
||||
target_interval=0,
|
||||
)
|
||||
self.assertEqual(generated, 1)
|
||||
self.assertEqual(client.models.generate_content.call_count, 1)
|
||||
self.assertEqual(
|
||||
read_json(copies / "Copie01_02.json"),
|
||||
{
|
||||
"name": "Continued",
|
||||
"list": [{"box_2d": [10, 20, 30, 40], "label": "Ex 2"}],
|
||||
},
|
||||
)
|
||||
|
||||
def test_label_detection_retries_unknown_labels(self) -> None:
|
||||
module = self.modules["gemini_for_labels"]
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
evaluation = Path(directory) / "Exam"
|
||||
(evaluation / "Copies").mkdir(parents=True)
|
||||
image = evaluation / "Cutleft" / "Copie01_01.jpg"
|
||||
image.parent.mkdir()
|
||||
image.write_bytes(b"image")
|
||||
client = Mock()
|
||||
client.models.generate_content.side_effect = [
|
||||
Mock(
|
||||
text=(
|
||||
'{"name":"Student","list":'
|
||||
'[{"box_2d":[1,2,3,4],"label":"Wrong"}]}'
|
||||
)
|
||||
),
|
||||
Mock(
|
||||
text=(
|
||||
'{"name":"Student","list":'
|
||||
'[{"box_2d":[1,2,3,4],"label":"Ex 1"}]}'
|
||||
)
|
||||
),
|
||||
]
|
||||
sleeps = []
|
||||
module.process_copy_group(
|
||||
EvaluationWorkspace(evaluation),
|
||||
"Copie01",
|
||||
[image],
|
||||
client=client,
|
||||
labels_text="Ex 1\n",
|
||||
names_text="Student\n",
|
||||
valid_labels={"Ex 1"},
|
||||
valid_names={"Student", "Unknown", "Continued"},
|
||||
overwrite=True,
|
||||
sleep=sleeps.append,
|
||||
target_interval=0,
|
||||
)
|
||||
self.assertEqual(client.models.generate_content.call_count, 2)
|
||||
self.assertIn(10, sleeps)
|
||||
self.assertEqual(
|
||||
read_json(evaluation / "Copies" / "Copie01_01.json")["list"][0][
|
||||
"label"
|
||||
],
|
||||
"Ex 1",
|
||||
)
|
||||
|
||||
def test_correction_overwrite_keeps_previous_state_until_a_commit(self) -> None:
|
||||
module = self.modules["correction"]
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
evaluation = Path(directory) / "Exam"
|
||||
evaluation.mkdir()
|
||||
correction = evaluation / "correction.json"
|
||||
progress = evaluation / "correction_progress.json"
|
||||
atomic_write_json(correction, {"Ex 1": [[{"id": "01"}]]})
|
||||
atomic_write_json(progress, [["old.jpg", "Ex 1"]])
|
||||
args = module.build_parser().parse_args(
|
||||
[str(evaluation), "--overwrite"]
|
||||
)
|
||||
module.configure_runtime(
|
||||
EvaluationWorkspace(evaluation),
|
||||
[("new.jpg", "Ex 1")],
|
||||
args,
|
||||
api_client=Mock(),
|
||||
)
|
||||
self.assertEqual(read_json(correction), {"Ex 1": [[{"id": "01"}]]})
|
||||
self.assertEqual(read_json(progress), [["old.jpg", "Ex 1"]])
|
||||
self.assertEqual(module.results, {"Ex 1": []})
|
||||
self.assertEqual(module.tasks_to_process, [("new.jpg", "Ex 1")])
|
||||
|
||||
def test_correction_reset_restores_old_and_deletes_new_files(self) -> None:
|
||||
module = self.modules["correction"]
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
evaluation = Path(directory) / "Exam"
|
||||
copy_dir = evaluation / "Copies" / "Copie01"
|
||||
copy_dir.mkdir(parents=True)
|
||||
atomic_write_json(evaluation / "correction.json", {"old": True})
|
||||
atomic_write_json(evaluation / "correction_progress.json", ["old"])
|
||||
(copy_dir / "Ex 1.pdf").write_bytes(b"current")
|
||||
(copy_dir / "Ex 1_old.pdf").write_bytes(b"original")
|
||||
(copy_dir / "Ex 2_new.pdf").write_bytes(b"generated")
|
||||
|
||||
self.assertEqual(module.main([str(evaluation), "--reset"]), 0)
|
||||
self.assertFalse((evaluation / "correction.json").exists())
|
||||
self.assertFalse((evaluation / "correction_progress.json").exists())
|
||||
self.assertEqual((copy_dir / "Ex 1.pdf").read_bytes(), b"original")
|
||||
self.assertFalse((copy_dir / "Ex 1_old.pdf").exists())
|
||||
self.assertFalse((copy_dir / "Ex 2_new.pdf").exists())
|
||||
|
||||
def test_correction_batch_request_files_are_written_atomically(self) -> None:
|
||||
module = self.modules["correction"]
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
evaluation = Path(directory) / "Exam"
|
||||
group_dir = evaluation / "Par label" / "Ex 1"
|
||||
(evaluation / "Copies").mkdir(parents=True)
|
||||
group_dir.mkdir(parents=True)
|
||||
(evaluation / "labels").write_text("Ex 1\n", encoding="utf-8")
|
||||
image = group_dir / "Group_1.jpg"
|
||||
image.write_bytes(b"image")
|
||||
atomic_write_json(
|
||||
image.with_suffix(".json"),
|
||||
[["01", 0, 400, 1.0, "Ex 1"]],
|
||||
)
|
||||
args = module.build_parser().parse_args([str(evaluation), "--batch"])
|
||||
module.configure_runtime(
|
||||
EvaluationWorkspace(evaluation),
|
||||
[(str(image), "Ex 1")],
|
||||
args,
|
||||
api_client=Mock(),
|
||||
)
|
||||
with patch.object(module.prompting, "make_prompt", return_value="prompt"):
|
||||
self.assertEqual(module.run_configured(args), 0)
|
||||
lines = (evaluation / "batch_requests_flash.jsonl").read_text().splitlines()
|
||||
self.assertEqual(len(lines), 1)
|
||||
self.assertEqual(json.loads(lines[0])["key"], str(image))
|
||||
self.assertEqual(
|
||||
(evaluation / "batch_requests_pro.jsonl").read_text(), ""
|
||||
)
|
||||
|
||||
def test_batch_helpers_use_mocked_api_and_atomic_combination(self) -> None:
|
||||
submit = self.modules["submit_batches"]
|
||||
fetch = self.modules["fetch_batched_results"]
|
||||
status = self.modules["batch_status"]
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
evaluation = Path(directory) / "Exam"
|
||||
evaluation.mkdir()
|
||||
(evaluation / "batch_requests_flash.jsonl").write_text(
|
||||
"{}\n", encoding="utf-8"
|
||||
)
|
||||
client = Mock()
|
||||
client.files.upload.return_value = SimpleNamespace(name="files/input")
|
||||
client.batches.create.return_value = SimpleNamespace(name="batches/1")
|
||||
self.assertEqual(
|
||||
submit.run(EvaluationWorkspace(evaluation), client=client), 0
|
||||
)
|
||||
self.assertEqual(client.batches.create.call_count, 1)
|
||||
manifest = json.loads((evaluation / "batch_jobs.json").read_text())
|
||||
self.assertEqual(manifest["jobs"]["flash"]["name"], "batches/1")
|
||||
|
||||
jobs = [
|
||||
SimpleNamespace(
|
||||
name="batches/1",
|
||||
display_name=f"flash-correction-{evaluation.name}",
|
||||
state=SimpleNamespace(name="JOB_STATE_SUCCEEDED"),
|
||||
dest=SimpleNamespace(file_name="files/flash-result"),
|
||||
),
|
||||
SimpleNamespace(
|
||||
name="batches/2",
|
||||
display_name=f"pro-correction-{evaluation.name}",
|
||||
state=SimpleNamespace(name="JOB_STATE_SUCCEEDED"),
|
||||
dest=SimpleNamespace(file_name="files/pro-result"),
|
||||
),
|
||||
]
|
||||
client.batches.get.return_value = jobs[0]
|
||||
client.files.download.side_effect = [b'{"flash":1}\n', b'{"pro":1}']
|
||||
self.assertEqual(
|
||||
fetch.run(EvaluationWorkspace(evaluation), client=client), 0
|
||||
)
|
||||
self.assertEqual(
|
||||
(evaluation / "batched_correction_result.jsonl").read_bytes(),
|
||||
b'{"flash":1}\n',
|
||||
)
|
||||
|
||||
job = SimpleNamespace(
|
||||
state=SimpleNamespace(name="JOB_STATE_SUCCEEDED"),
|
||||
dest=SimpleNamespace(file_name="files/result"),
|
||||
)
|
||||
client.batches.get.return_value = job
|
||||
client.files.download.side_effect = None
|
||||
client.files.download.return_value = b'{"downloaded":true}\n'
|
||||
output = evaluation / "one-result.jsonl"
|
||||
self.assertEqual(
|
||||
status.download_job("batches/1", output=output, client=client), 0
|
||||
)
|
||||
self.assertEqual(output.read_bytes(), b'{"downloaded":true}\n')
|
||||
|
||||
def test_post_correction_main_cleans_json_atomically(self) -> None:
|
||||
module = self.modules["post_correction"]
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
|
||||
Reference in New Issue
Block a user