995 lines
42 KiB
Python
995 lines
42 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import base64
|
|
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 google import genai
|
|
|
|
from copienator import configuration as config
|
|
from copienator.commands import grouping
|
|
from copienator import prompting
|
|
from copienator.feedback_boxes import valid_feedback_box
|
|
from copienator import (
|
|
CliError,
|
|
EvaluationWorkspace,
|
|
ExitCode,
|
|
atomic_write_json,
|
|
atomic_write_text,
|
|
execute,
|
|
read_json,
|
|
target_parser,
|
|
workspace_from_target,
|
|
)
|
|
from copienator.utils import enonce_total, read_all_labels
|
|
|
|
NB_THREADS = 12
|
|
|
|
# PROXY_URL = "http://192.168.241.1:3128"
|
|
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 = {}
|
|
|
|
def tprint(*args, **kwargs):
|
|
"""Buffer messages per thread to group them."""
|
|
tid = threading.current_thread().name
|
|
msg = " ".join(map(str, args))
|
|
|
|
with log_lock:
|
|
if tid not in thread_logs:
|
|
thread_logs[tid] = []
|
|
thread_logs[tid].append(msg)
|
|
|
|
# Optional: Keep printing to console but prefix with thread name
|
|
print(f"[{tid}] {msg}", **kwargs)
|
|
|
|
def flush_thread_log(tid=None):
|
|
"""Append a thread's buffered messages to the log file contiguously."""
|
|
tid = tid or threading.current_thread().name
|
|
with log_lock:
|
|
if thread_logs.get(tid):
|
|
with open(INPUT_DIR / "correction_log", "a", encoding="utf-8") as f:
|
|
f.write(f"--- Task Log [{tid}] ---\n")
|
|
f.write("\n".join(thread_logs[tid]) + "\n\n")
|
|
thread_logs[tid].clear()
|
|
|
|
# --- Lock for thread-safe file writing ---
|
|
io_lock = threading.Lock()
|
|
pro_lock = threading.Lock()
|
|
pro_count = 0
|
|
flash_count = 0
|
|
pro_quota_exhausted = False
|
|
|
|
|
|
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):
|
|
"""Handles requests to Gemini with a 1min and 5min retry mechanism, and quota fallback."""
|
|
global pro_quota_exhausted
|
|
delays = [60, 300]
|
|
|
|
for attempt in range(3):
|
|
# Switch to fallback immediately if quota was exhausted by another thread
|
|
if model_id == MODEL_ID_pro and pro_quota_exhausted and fallback_model_id:
|
|
model_id = fallback_model_id
|
|
|
|
try:
|
|
full_response_text = ""
|
|
for chunk in client.models.generate_content_stream(
|
|
model=model_id,
|
|
contents=contents,
|
|
config=config,
|
|
):
|
|
if chunk.text:
|
|
full_response_text += chunk.text
|
|
return full_response_text
|
|
except Exception as e:
|
|
error_msg = str(e).lower()
|
|
is_quota_error = "429" in error_msg or "quota" in error_msg or "exhausted" in error_msg
|
|
is_minute_limit = "minute" in error_msg or "rpm" in error_msg or "tpm" in error_msg
|
|
|
|
if is_minute_limit:
|
|
import re
|
|
# Extract wait time if present, else use default delay
|
|
retry_match = re.search(r"retry in ([\d.]+)s", error_msg)
|
|
wait_time = float(retry_match.group(1)) + 1.0 if retry_match else delays[attempt]
|
|
|
|
tprint(f"\tGemini Pro minute limit hit. Waiting {wait_time:.1f}s...")
|
|
time.sleep(wait_time)
|
|
continue # Retry same model
|
|
|
|
# Immediately fallback to Flash without waiting if it's a Pro quota error
|
|
if is_quota_error and model_id == MODEL_ID_pro and fallback_model_id:
|
|
tprint(f"\tGemini Pro quota hit ({e}). \n\n\tFalling back to Flash permanently...")
|
|
model_id = fallback_model_id
|
|
pro_quota_exhausted = True
|
|
continue # Retry immediately with Flash
|
|
|
|
if attempt < 2:
|
|
tprint(f"\tGemini API failure: {e}. Retrying in {delays[attempt]} seconds...")
|
|
time.sleep(delays[attempt])
|
|
else:
|
|
tprint(f"\tGemini API failure: {e}. Maximum retries reached.")
|
|
raise
|
|
|
|
def correct_boxes_with_gemini(pid, label, pdf_path, original_feedbacks,
|
|
yming, ymaxg, width_r, total_height):
|
|
"""Requests corrected bounding boxes from Gemini Flash on the single image."""
|
|
# pdf_path = COPIES_DIR / f"Copie{pid}" / f"{label}.pdf"
|
|
|
|
contents, config = prompting.request_for_box_correction(pdf_path, original_feedbacks)
|
|
response_text = call_gemini_with_retries(MODEL_ID_flash, contents, config)
|
|
corrected_feedbacks = json.loads(response_text)
|
|
|
|
global_feedbacks = [f for f in original_feedbacks if not f["box_2d"]]
|
|
|
|
# Map the coordinates back from the single image to the group canvas
|
|
for f in corrected_feedbacks:
|
|
b = f.get("box_2d")
|
|
if b:
|
|
if not valid_feedback_box(b) or any(value < 0 or value > 1000 for value in b):
|
|
f["box_2d"] = None
|
|
continue
|
|
ymin_s, xmin_s, ymax_s, xmax_s = b
|
|
|
|
# Y mapping: Add the group Y-offset (yming), then normalize to total_height
|
|
single_h = ymaxg - yming
|
|
new_ymin = int((yming + (ymin_s * single_h / 1000.0)) * 1000.0 / total_height)
|
|
new_ymax = int((yming + (ymax_s * single_h / 1000.0)) * 1000.0 / total_height)
|
|
|
|
# X mapping: Multiply by the width ratio of this sub-image vs the group image
|
|
new_xmin = int(xmin_s * width_r)
|
|
new_xmax = int(xmax_s * width_r)
|
|
|
|
f["box_2d"] = [new_ymin, new_xmin, new_ymax, new_xmax]
|
|
|
|
return global_feedbacks + corrected_feedbacks
|
|
|
|
def get_next_group_idx(label):
|
|
"""Finds the next available Group index for a given label."""
|
|
target_folder = GROUPS_DIR / label
|
|
target_folder.mkdir(exist_ok=True)
|
|
existing = list(target_folder.glob("Group_*.jpg"))
|
|
if not existing: return 0
|
|
return max([int(f.stem.split("_")[1]) for f in existing])
|
|
|
|
def handle_label_errors(pid, label, res, pdf_path):
|
|
"""Handles Gemini labeling errors, moves/copies files, and returns new tasks."""
|
|
new_tasks = []
|
|
error_type = res.get("error")
|
|
|
|
all_labels = read_all_labels(INPUT_DIR)
|
|
labels_txt = (INPUT_DIR / "labels").read_text(encoding="utf-8", errors="replace")
|
|
enonce = enonce_total(INPUT_DIR)
|
|
|
|
if error_type == "wrong-label":
|
|
tprint(f"\tHandling wrong-label for {pid} {label}")
|
|
contents, config = prompting.request_for_wrong_label(pdf_path, label, enonce, labels_txt)
|
|
new_label = call_gemini_with_retries(MODEL_ID_flash, contents, config).strip().strip('"\'')
|
|
if new_label not in all_labels:
|
|
tprint(f"\t\tCopie{pid} returned an incorrect label {new_label} from an initial wrong label {label}. Ignoring")
|
|
res["error"] = "wrg-lbl:cldtfix"
|
|
return []
|
|
if new_label == label:
|
|
res["error"] = ""
|
|
return []
|
|
|
|
base_new_pdf_path = COPIES_DIR / f"Copie{pid}" / f"{new_label}.pdf"
|
|
new_pdf_path = COPIES_DIR / f"Copie{pid}" / f"{new_label}_new.pdf"
|
|
|
|
if base_new_pdf_path.exists() or new_pdf_path.exists():
|
|
tprint(f"""\t\tCopie{pid} tried to move wrong {label} to {new_label},
|
|
but it already exists. Delaying.""")
|
|
# res["error"] = f"wrg-lbl:{new_label}?exists"
|
|
res["error"] = f"wrg-lbl:{new_label}?"
|
|
res.setdefault("delayed", []).append(["wrong-label", new_label])
|
|
else:
|
|
res["error"] = f"wrg-lbl-moved-to:{new_label}"
|
|
tprint(f"\t\tCopie{pid} : moving wrong {label} to {new_label}.")
|
|
|
|
# Copie vers _new, puis renommage de l'original vers _old
|
|
shutil.copy(str(pdf_path), str(new_pdf_path))
|
|
old_pdf_path = pdf_path.with_name(f"{label}_old.pdf")
|
|
if pdf_path != old_pdf_path:
|
|
shutil.move(str(pdf_path), str(old_pdf_path))
|
|
|
|
idx = get_next_group_idx(new_label)
|
|
height = grouping.get_pdf_height(str(new_pdf_path))
|
|
grouping.create_jpg(new_label, idx, [(pid, str(new_pdf_path), height)], GROUPS_DIR)
|
|
tprint(f"\t\tMaking {new_label} group {idx+1}")
|
|
new_tasks.append((str(GROUPS_DIR / new_label / f"Group_{idx+1}.jpg"),
|
|
new_label, False))
|
|
|
|
elif error_type == "additional-answer":
|
|
contents, config = prompting.request_for_additional_answer(pdf_path, label, enonce, labels_txt)
|
|
tprint(f"\tHandling additional-answer for {pid} {label}")
|
|
try:
|
|
add_labels = json.loads(call_gemini_with_retries(MODEL_ID_flash, contents, config))
|
|
except Exception: # noqa: BLE001 - invalid auxiliary model response
|
|
add_labels = []
|
|
|
|
keep_error = False
|
|
error = "al:"
|
|
for add_label in add_labels:
|
|
if add_label == label:
|
|
continue
|
|
if add_label not in all_labels:
|
|
tprint(f"\t\t Inexistent label ({add_label}) from additional-answer processing {pid} {label}. Ignoring")
|
|
error += f"{add_label}??"
|
|
keep_error = True
|
|
continue
|
|
|
|
base_add_pdf_path = COPIES_DIR / f"Copie{pid}" / f"{add_label}.pdf"
|
|
add_pdf_path = COPIES_DIR / f"Copie{pid}" / f"{add_label}_new.pdf"
|
|
|
|
if not base_add_pdf_path.exists() and not add_pdf_path.exists():
|
|
shutil.copy(str(pdf_path), str(add_pdf_path))
|
|
tprint(f"\t\tCopying Copie{pid} : {label} -> {add_label}")
|
|
idx = get_next_group_idx(add_label)
|
|
tprint(f"\t\tMaking {add_label} group {idx+1}")
|
|
height = grouping.get_pdf_height(str(add_pdf_path))
|
|
grouping.create_jpg(add_label, idx, [(pid, str(add_pdf_path), height)], GROUPS_DIR)
|
|
new_tasks.append((str(GROUPS_DIR / add_label / f"Group_{idx+1}.jpg"),
|
|
add_label, False))
|
|
error += f"(->){add_label}"
|
|
keep_error = True
|
|
else:
|
|
keep_error = True
|
|
error += f"(->){add_label}?"
|
|
res.setdefault("delayed", []).append(["add-label", add_label])
|
|
tprint(f"\t\tAlready present (not copied) Copie{pid} : {label} -> {add_label}. Delaying.")
|
|
if not keep_error:
|
|
res["error"] = ""
|
|
else:
|
|
res["error"] = error
|
|
|
|
return new_tasks
|
|
|
|
def process_single_task(task_tuple, precomputed_response=None):
|
|
try:
|
|
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
|
|
|
|
group_name = os.path.splitext(file_path)[0]
|
|
json_path = group_name + '.json'
|
|
new_tasks = []
|
|
|
|
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}
|
|
total_height = group_data[-1][2]
|
|
use_flash = n >= 4 or total_height <= 500
|
|
|
|
# Only apply limits and counts if we are making a live call
|
|
if precomputed_response is None:
|
|
if not use_flash:
|
|
with pro_lock:
|
|
if pro_quota_exhausted:
|
|
use_flash = True
|
|
elif limit is None or pro_count < limit:
|
|
pro_count += 1
|
|
else:
|
|
use_flash = True
|
|
|
|
if use_flash:
|
|
with pro_lock:
|
|
flash_count += 1
|
|
|
|
try:
|
|
contents, config = prompting.generate_request(INPUT_DIR, file_path, label)
|
|
model_to_use = MODEL_ID_flash if use_flash else MODEL_ID_pro
|
|
|
|
if precomputed_response:
|
|
tprint(f"Using batched response for: {label} {group_name}")
|
|
full_response_text = precomputed_response
|
|
else:
|
|
tprint(f"Asking Gemini {'Flash' if use_flash else 'Pro '}: {label} {group_name}")
|
|
full_response_text = call_gemini_with_retries(model_to_use, contents, config)
|
|
|
|
json_data = json.loads(full_response_text)
|
|
|
|
# Ensure consistency of answer placements
|
|
for p in json_data:
|
|
pid = p["id"]
|
|
res = p["result"]
|
|
yming, ymaxg, width_r = d_data[pid]
|
|
|
|
pdf_path = COPIES_DIR / f"Copie{pid}" / f"{label}.pdf"
|
|
current_suffix = ""
|
|
|
|
# Détection du vrai fichier s'il a un suffixe
|
|
if not pdf_path.exists():
|
|
if pdf_path.with_name(f"{label}_new.pdf").exists():
|
|
pdf_path = pdf_path.with_name(f"{label}_new.pdf")
|
|
current_suffix = "_new"
|
|
# Quand est-ce que ce chemin est utilisé ? Jamais ?
|
|
elif pdf_path.with_name(f"{label}_old.pdf").exists():
|
|
pdf_path = pdf_path.with_name(f"{label}_old.pdf")
|
|
current_suffix = "_old"
|
|
|
|
# 1. Gestion de empty-answer
|
|
if res.get("error") == "empty-answer":
|
|
old_path = pdf_path.with_name(f"{label}_old.pdf")
|
|
if pdf_path.exists() and pdf_path != old_path:
|
|
shutil.move(str(pdf_path), str(old_path))
|
|
pdf_path = old_path
|
|
current_suffix = "_old"
|
|
|
|
if (not can_spawn_tasks) and res["error"] == "additional-answer":
|
|
tprint("\tSwallowing an additional-answer from a subsequent task.")
|
|
res["error"]= ""
|
|
if res["error"] != "":
|
|
tprint("\tError :", res["error"], "for Copie", pid, group_name)
|
|
|
|
if can_spawn_tasks and res.get("error") in ["wrong-label", "additional-answer"]:
|
|
new_tasks.extend(handle_label_errors(pid, label, res, pdf_path))
|
|
# Si "wrong-label" a déplacé le fichier courant vers _old
|
|
if res.get("error", "").startswith("wrg-lbl-moved-to:"):
|
|
current_suffix = "_old"
|
|
|
|
# 5. Enregistrer l'information dans correction.json
|
|
if current_suffix:
|
|
res["suffix"] = current_suffix
|
|
|
|
needs_correction = []
|
|
for (i,f) in enumerate(res["feedback"]):
|
|
b = f.get("box_2d")
|
|
if b:
|
|
if not valid_feedback_box(b):
|
|
needs_correction.append(i)
|
|
continue
|
|
ymin, _xmin, ymax, xmax = b
|
|
ymin = ymin * total_height // 1000
|
|
ymax = ymax * total_height // 1000
|
|
|
|
if pid not in d_data:
|
|
tprint("Error : Gemini answered a copie id not present",
|
|
pid, label, group_name)
|
|
continue
|
|
|
|
if (ymin < yming - 50 or ymax > ymaxg + 50
|
|
or ymin > ymaxg + 50 or ymax < yming - 50
|
|
or _xmin < 0 or xmax / 1000 > width_r):
|
|
needs_correction.append(i)
|
|
continue
|
|
if ymin < yming - 5:
|
|
ymin = yming - 5
|
|
b[0] = ymin * 1000 // total_height
|
|
if ymax > ymaxg + 5:
|
|
ymax = ymaxg + 5
|
|
b[2] = ymax * 1000 // total_height
|
|
|
|
|
|
if needs_correction:
|
|
tprint(f"\tBox anomalies detected for Copie {pid} {group_name}. \n\tRequesting isolated correction from Gemini Flash...")
|
|
try:
|
|
# Pensez à passer pdf_path à la fonction modifiée !
|
|
res["feedback"] = correct_boxes_with_gemini(
|
|
pid, label, pdf_path, res["feedback"],
|
|
yming, ymaxg, width_r, total_height)
|
|
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"]):
|
|
if i in needs_correction:
|
|
f["box_2d"] = None
|
|
|
|
# --- Use Lock for writing shared data ---
|
|
with io_lock:
|
|
if label not in results:
|
|
results[label] = []
|
|
results[label].append(json_data)
|
|
|
|
atomic_write_json(output_path, results)
|
|
|
|
# To track progress
|
|
completed_tasks.append((file_path, label))
|
|
atomic_write_json(progress_path, completed_tasks)
|
|
|
|
except json.JSONDecodeError:
|
|
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: # noqa: BLE001 - per-task processing boundary
|
|
error_msg = f"Exception processing {file_path}: {e}"
|
|
print(error_msg, file=sys.stderr)
|
|
with io_lock:
|
|
errors_summary.append((error_msg, file_path))
|
|
return new_tasks
|
|
finally:
|
|
flush_thread_log()
|
|
|
|
def resolve_delayed_moves():
|
|
"""Scans the current results to find delayed moves and executes them if space was freed."""
|
|
new_tasks = []
|
|
with io_lock:
|
|
for label, batches in results.items():
|
|
for batch in batches:
|
|
for p in batch:
|
|
res = p.get("result", {})
|
|
delayed_list = res.get("delayed", [])
|
|
if not delayed_list:
|
|
continue
|
|
|
|
pid = p["id"]
|
|
pdf_path = COPIES_DIR / f"Copie{pid}" / f"{label}.pdf"
|
|
|
|
if not pdf_path.exists():
|
|
if pdf_path.with_name(f"{label}_new.pdf").exists():
|
|
pdf_path = pdf_path.with_name(f"{label}_new.pdf")
|
|
elif pdf_path.with_name(f"{label}_old.pdf").exists():
|
|
pdf_path = pdf_path.with_name(f"{label}_old.pdf")
|
|
|
|
remaining_delayed = []
|
|
resolved_any = False
|
|
|
|
for delay_type, target_label in delayed_list:
|
|
# 1. Résolution de wrong-label
|
|
if delay_type == "wrong-label":
|
|
base_new_pdf_path = COPIES_DIR / f"Copie{pid}" / f"{target_label}.pdf"
|
|
new_pdf_path = COPIES_DIR / f"Copie{pid}" / f"{target_label}_new.pdf"
|
|
|
|
if not base_new_pdf_path.exists() and not new_pdf_path.exists():
|
|
tprint(f"Resolving delayed move: Copie{pid} {label} -> {target_label}")
|
|
res["error"] = f"wrg-lbl-moved-to:{target_label}"
|
|
res["suffix"] = "_old" # Fixed typo: was suffixe
|
|
resolved_any = True
|
|
|
|
shutil.copy(str(pdf_path), str(new_pdf_path))
|
|
old_pdf_path = pdf_path.with_name(f"{label}_old.pdf")
|
|
if pdf_path != old_pdf_path:
|
|
shutil.move(str(pdf_path), str(old_pdf_path))
|
|
|
|
idx = get_next_group_idx(target_label)
|
|
height = grouping.get_pdf_height(str(new_pdf_path))
|
|
grouping.create_jpg(target_label, idx, [(pid, str(new_pdf_path), height)], GROUPS_DIR)
|
|
new_tasks.append((str(GROUPS_DIR / target_label / f"Group_{idx+1}.jpg"), target_label, False))
|
|
else:
|
|
remaining_delayed.append([delay_type, target_label])
|
|
|
|
# 2. Résolution de additional-answer
|
|
elif delay_type == "add-label":
|
|
base_add_pdf_path = COPIES_DIR / f"Copie{pid}" / f"{target_label}.pdf"
|
|
add_pdf_path = COPIES_DIR / f"Copie{pid}" / f"{target_label}_new.pdf"
|
|
|
|
if not base_add_pdf_path.exists() and not add_pdf_path.exists():
|
|
tprint(f"Resolving delayed additional-answer: Copie{pid} {label} -> {target_label}")
|
|
res["error"] = res["error"].replace(f"(xx){target_label}", f"(->){target_label}")
|
|
resolved_any = True
|
|
|
|
shutil.copy(str(pdf_path), str(add_pdf_path))
|
|
idx = get_next_group_idx(target_label)
|
|
height = grouping.get_pdf_height(str(add_pdf_path))
|
|
grouping.create_jpg(target_label, idx, [(pid, str(add_pdf_path), height)], GROUPS_DIR)
|
|
new_tasks.append((str(GROUPS_DIR / target_label / f"Group_{idx+1}.jpg"), target_label, False))
|
|
else:
|
|
remaining_delayed.append([delay_type, target_label])
|
|
|
|
if resolved_any:
|
|
if remaining_delayed:
|
|
res["delayed"] = remaining_delayed
|
|
else:
|
|
del res["delayed"]
|
|
|
|
if new_tasks:
|
|
atomic_write_json(output_path, results)
|
|
|
|
return new_tasks
|
|
|
|
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():
|
|
refaire_list = read_json(refaire_path)
|
|
|
|
overwritten_data = []
|
|
if overwritten_path.exists():
|
|
overwritten_data = read_json(overwritten_path)
|
|
|
|
dirty_results = False
|
|
|
|
for copie_name, labels in refaire_list:
|
|
pid = copie_name.replace("Copie", "")
|
|
copie_dir = COPIES_DIR / copie_name
|
|
|
|
# If list is empty, redo all labels available for this Copie
|
|
if not labels:
|
|
labels_set = set()
|
|
for p in copie_dir.glob("*.pdf"):
|
|
if p.name.endswith("_old.pdf"):
|
|
continue # Strictly ignore old files
|
|
elif p.name.endswith("_new.pdf"):
|
|
labels_set.add(p.stem[:-4]) # Strip '_new' to get base label
|
|
else:
|
|
labels_set.add(p.stem)
|
|
labels = list(labels_set)
|
|
|
|
for label in labels:
|
|
# 1. Extract and backup old corrections
|
|
if label in results:
|
|
for batch in results[label]:
|
|
to_remove = None
|
|
for item in batch:
|
|
if item.get("id") == pid:
|
|
to_remove = item
|
|
break
|
|
if to_remove:
|
|
batch.remove(to_remove)
|
|
overwritten_data.append({
|
|
"pid": pid,
|
|
"label": label,
|
|
"data": to_remove,
|
|
"timestamp": time.time()
|
|
})
|
|
dirty_results = True
|
|
# Clean up empty batches
|
|
results[label] = [b for b in results[label] if b]
|
|
|
|
# 2. Make new group and add to tasks
|
|
pdf_path = copie_dir / f"{label}.pdf"
|
|
is_new = False
|
|
|
|
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"
|
|
|
|
if pdf_path.exists():
|
|
idx = get_next_group_idx(label)
|
|
height = grouping.get_pdf_height(str(pdf_path))
|
|
grouping.create_jpg(label, idx, [(pid, str(pdf_path), height)], GROUPS_DIR)
|
|
new_group_path = str(GROUPS_DIR / label / f"Group_{idx+1}.jpg")
|
|
tasks_to_process.append((new_group_path, label, not is_new))
|
|
|
|
if dirty_results:
|
|
atomic_write_json(output_path, results)
|
|
atomic_write_json(overwritten_path, overwritten_data)
|
|
else:
|
|
print(f"Warning: --refaire flag used, but {refaire_path} not found.", file=sys.stderr)
|
|
|
|
|
|
if args.batch or args.batch_from:
|
|
all_labels = read_all_labels(INPUT_DIR)
|
|
batch_tasks = []
|
|
if args.batch_from:
|
|
for label in all_labels:
|
|
if label.startswith(args.batch_from):
|
|
args.batch_from = label
|
|
input(f"About to batch from: {args.batch_from}. Press Enter to confirm...")
|
|
break
|
|
if args.batch_from not in 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 = []
|
|
|
|
for task in tasks_to_process:
|
|
lbl = task[1]
|
|
# Any label found sequentially equal or after `args.batch_from` gets batched
|
|
if lbl in all_labels and all_labels.index(lbl) >= target_idx:
|
|
batch_tasks.append(task)
|
|
else:
|
|
live_tasks.append(task)
|
|
|
|
tasks_to_process = live_tasks # Keep live tasks to be run right after
|
|
else:
|
|
batch_tasks = tasks_to_process
|
|
tasks_to_process = [] # Run nothing live if just `--batch`
|
|
|
|
if batch_tasks:
|
|
batch_flash_file = INPUT_DIR / "batch_requests_flash.jsonl"
|
|
batch_pro_file = INPUT_DIR / "batch_requests_pro.jsonl"
|
|
|
|
count_flash = 0
|
|
count_pro = 0
|
|
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,
|
|
},
|
|
},
|
|
}
|
|
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 ""),
|
|
)
|
|
|
|
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:
|
|
return ExitCode.SUCCESS
|
|
|
|
batched_responses = {}
|
|
if args.deal_with_batched:
|
|
batch_results_path = INPUT_DIR / "batched_correction_result.jsonl"
|
|
if batch_results_path.exists():
|
|
print(f"Loading batch results from {batch_results_path}...")
|
|
with open(batch_results_path, "r", encoding="utf-8") as f:
|
|
for line in f:
|
|
if not line.strip(): continue
|
|
data = json.loads(line)
|
|
task_id = data.get("key") # Corresponds to the key sent in the request
|
|
|
|
if "response" in data:
|
|
try:
|
|
# Extract the JSON response text per standard Batch API schema
|
|
resp_text = data["response"]["candidates"][0]["content"]["parts"][0]["text"]
|
|
batched_responses[task_id] = resp_text
|
|
except (KeyError, IndexError) as e:
|
|
print(f"Warning: Could not parse response for {task_id}: {e}", file=sys.stderr)
|
|
elif "error" in data:
|
|
print(f"Batch API Error for {task_id}: {data['error']}", file=sys.stderr)
|
|
else:
|
|
print(f"Warning: Batch results file {batch_results_path} not found.", file=sys.stderr)
|
|
|
|
made_progress = True
|
|
while tasks_to_process or made_progress:
|
|
if tasks_to_process:
|
|
print(f"Starting processing on {len(tasks_to_process)} tasks with {NB_THREADS} threads...")
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=NB_THREADS) as executor:
|
|
futures = {}
|
|
for task in tasks_to_process:
|
|
file_path = task[0]
|
|
precomp = batched_responses.get(file_path)
|
|
futures[executor.submit(process_single_task, task, precomp)] = task
|
|
|
|
for future in concurrent.futures.as_completed(futures):
|
|
try:
|
|
new_generated_tasks = future.result()
|
|
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: # 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
|
|
|
|
# Après avoir traité toutes les tâches actuelles (live ou batched),
|
|
# on tente de débloquer les mouvements qui étaient en attente
|
|
delayed_tasks = resolve_delayed_moves()
|
|
if delayed_tasks:
|
|
print(f"Resolved {len(delayed_tasks)} delayed moves! Running executor for new tasks...")
|
|
tasks_to_process.extend(delayed_tasks)
|
|
made_progress = True
|
|
else:
|
|
made_progress = False
|
|
|
|
# Check for remaining unresolved delayed tasks
|
|
unresolved_delayed = []
|
|
with io_lock:
|
|
for label, batches in sorted(results.items()):
|
|
for batch in batches:
|
|
for p in batch:
|
|
res = p.get("result", {})
|
|
delayed = res.get("delayed", [])
|
|
pid = p["id"]
|
|
|
|
for delay_type, target_label in delayed:
|
|
if delay_type == "wrong-label":
|
|
unresolved_delayed.append(f"Copie{pid} {label} x> {target_label}|")
|
|
elif delay_type == "add-label":
|
|
unresolved_delayed.append(f"Copie{pid} {label} -> {target_label}|")
|
|
|
|
if unresolved_delayed:
|
|
manual_path = INPUT_DIR / "manual_resolutions.txt"
|
|
atomic_write_text(
|
|
manual_path,
|
|
"### Use -> x>, -x, ss, sx, xx, xs, c{43}1>, c{43}2x\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 -m copienator resolve-manual <InputDir>`")
|
|
|
|
end_time = time.time()
|
|
print("Time elapsed : ", end_time - start_time)
|
|
print("Requests to pro / flash : ", pro_count, flash_count)
|
|
if errors_summary:
|
|
print("\n--- Summary of Exceptions (You can use several images on one instance) ---", file=sys.stderr)
|
|
for (err, file) in errors_summary:
|
|
print(err, file=sys.stderr)
|
|
escaped_path = shlex.quote(str(file))
|
|
print(f"Run : python -m copienator correct {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())
|