Files
Copies/copienator/commands/plotting.py
T

476 lines
18 KiB
Python

from __future__ import annotations
import argparse
import queue
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 copienator.platform import open_path
from copienator.utils import natural_key, read_all_labels
# --- Configuration & Globals ---
padding = 60
try:
font = ImageFont.truetype("DejaVuSans.ttf", size=30)
except OSError:
font = ImageFont.load_default()
# --- Helper Functions (Shared) ---
def page_number(b, nb_pages):
column_width = 1000 // nb_pages
center_x = (b[1] + b[3]) // 2
return center_x // column_width
def convert_box2d(b, pn_ori, npn, tot_ori, tot_dest):
l = b.copy()
l[1] = (l[1] - (1000 // tot_ori) * (pn_ori-1)) * tot_ori // tot_dest\
+ (1000 // tot_dest) * (npn - 1)
l[3] = (l[3] - (1000 // tot_ori) * (pn_ori-1)) * tot_ori // tot_dest\
+ (1000 // tot_dest) * (npn - 1)
return l
def convert_list(l, group_id, json_schema):
ll = []
nb_pages = json_schema["columns_per_file"][group_id-1]
nb_previous_pages = sum([json_schema["columns_per_file"][i] for i in range(group_id-1)])
nb_tot_pages = sum([e for e in json_schema["columns_per_file"]])
for e in l:
ee = e.copy()
pn = page_number(e["box_2d"], nb_pages)
npn = pn + nb_previous_pages
ee["box_2d"] = convert_box2d(ee["box_2d"], pn, npn, nb_pages, nb_tot_pages)
ee["part"] = group_id
ee["pn"] = npn
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()
width, height = im.size
new_im = Image.new(im.mode, (width + padding, height), "white")
new_im.paste(im, (0, 0))
draw = ImageDraw.Draw(new_im)
bounding_boxes.sort(key=lambda b: (page_number(b["box_2d"], nb_pages), b["box_2d"][0]))
for bbox in bounding_boxes:
raw_y_min = int(bbox["box_2d"][0] * height / 1000)
raw_x_min = int(bbox["box_2d"][1] * width / 1000)
raw_y_max = int(bbox["box_2d"][2] * height / 1000)
raw_x_max = int(bbox["box_2d"][3] * width / 1000)
abs_y_min = max(0, raw_y_min - 10)
abs_x_min = max(0, raw_x_min - 10)
abs_y_max = min(height, raw_y_max + 10)
abs_x_max = min(width, raw_x_max + 10)
color = "black"
label = bbox.get("label")
if label and label in all_labels:
current_index = all_labels.index(label)
if current_index < last_label_index or (last_label_index == -1 and current_index != 0):
color = "red"
elif current_index > last_label_index + 1:
color = "orange"
last_label_index = current_index
draw.rectangle(((abs_x_min, abs_y_min), (abs_x_max, abs_y_max)), outline=color, width=4)
if label:
if abs_y_min > 80:
draw.text((abs_x_min + 8, abs_y_min - 30), label, fill=color, font=font)
else:
draw.text((abs_x_min + 8, abs_y_max + 6), label, fill=color, font=font)
return (new_im, last_label_index)
# --- Processing Logic (Worker Thread) ---
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.
"""
previous_copie = None
last_label_index = None
for img_path in files_to_process:
json_path = base_dir / "Copies" / f"{img_path.stem}.json"
copie_part = int(img_path.stem[-2:])
copie = img_path.stem[:-3]
if copie != previous_copie:
last_label_index = -1
previous_copie = copie
json_schema_path = base_dir / 'Cutleft' / f"{copie}_schema.json"
try:
json_schema = read_json(json_schema_path)
except (OSError, TypeError, ValueError):
print("No json_schema : ", json_schema_path)
continue
nb_pages = json_schema["columns_per_file"][copie_part-1]
if json_path.exists():
# Read strictly for visualization purposes
bb_list = []
json_name = ""
try:
json_result = read_json(json_path)
bb_list = json_result.get("list", [])
json_name = json_result.get("name", "")
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
try:
print(f"Buffering {img_path.name}...")
(pil_image, last_label_index) = \
prepare_image(str(img_path), bb_list, all_labels, nb_pages, last_label_index)
error_msg = None
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)
metadata = {
"copie": copie,
"part": copie_part,
"schema": json_schema,
"name": json_name,
"error": error_msg
}
output_queue.put((pil_image, json_path, metadata))
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, workspace, valid_labels, input_queue):
self.root = root
self.root.resizable(False, False) # If you resize, coordinates will be wrong
screen_w = root.winfo_screenwidth()
screen_h = root.winfo_screenheight()
x = int(screen_w * 0.1)
y = int(screen_h * 0.05)
root.geometry(f"+{x}+{y}")
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")
# Display State
self.current_image = None
self.current_json_path = None
self.current_meta = None # Stores schema/copie info
self.is_viewing = False
self.scale_factor = 1.0
self.orig_size = (1, 1)
# Data Aggregation State
self.active_copie_name = None
self.accumulated_results = None # Dict with "name" and "list"
# To go back
self.history = []
self.forward_stack = []
self.current_pil_image = None
self.failed = False
from config import PLOTTING_KB
# Bindings
self.root.bind(PLOTTING_KB["OK"], self.on_enter)
self.root.bind(PLOTTING_KB["previous"], self.on_previous)
self.root.bind(PLOTTING_KB["edit"], self.on_edit)
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 _event: self.close())
self.root.protocol("WM_DELETE_WINDOW", self.close)
self.label.bind('<Button-1>', self.on_click)
self.poll_queue()
def poll_queue(self):
if not self.is_viewing:
try:
# pil_image, json_path, metadata = image_queue.get_nowait()
if self.forward_stack:
pil_image, json_path, metadata = self.forward_stack.pop()
else:
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()
return
# Check if we switched to a new "Copie" group
if self.active_copie_name != metadata["copie"]:
self.save_current_batch() # Write previous group to disk
# Start new batch
self.active_copie_name = metadata["copie"]
self.accumulated_results = {"name": metadata["name"], "list": []}
self.history.clear()
self.display_image(pil_image, json_path, metadata)
except queue.Empty:
pass
self.root.after(100, self.poll_queue)
def save_current_batch(self):
"""Writes the accumulated data to the main JSON file."""
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}")
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:
print("Going back to previous image...")
prev_pil, prev_json, prev_meta, num_added = self.history.pop()
# Undo the accumulation to prevent duplicates when we hit Enter again
if self.accumulated_results and num_added > 0:
self.accumulated_results["list"] = self.accumulated_results["list"][:-num_added]
# Push current image to the forward stack so we don't lose it
self.forward_stack.append((self.current_pil_image,
self.current_json_path, self.current_meta))
# Display the previous image immediately
self.display_image(prev_pil, prev_json, prev_meta)
def display_image(self, pil_image, json_path, metadata):
self.current_pil_image = pil_image # ADD THIS LINE
self.orig_size = pil_image.size
self.scale_factor = 1.0
screen_h = self.root.winfo_screenheight() - 100
if pil_image.height > screen_h:
self.scale_factor = screen_h / pil_image.height
pil_image = pil_image.resize((int(pil_image.width * self.scale_factor),
int(pil_image.height * self.scale_factor)))
self.tk_image = ImageTk.PhotoImage(pil_image)
self.label.config(image=self.tk_image, text=f"Processing: {json_path.name}")
self.current_json_path = json_path
self.current_meta = metadata
self.is_viewing = True
self.root.lift()
if metadata.get("error"):
msg = f"Error generating boxes for {json_path.name}:\n\n{metadata['error']}\n\nPlease press 'e' to fix the JSON file, then press Enter to retry."
messagebox.showerror("Processing Error", msg)
def on_enter(self, event):
if self.is_viewing:
print(f"Committing data for {self.current_json_path.name}...")
num_added = 0 # ADD THIS LINE
try:
current_data = read_json(self.current_json_path)
# Perform the conversion now, post-edit
converted_items = convert_list(
current_data["list"],
self.current_meta["part"],
self.current_meta["schema"]
)
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."
print(msg)
messagebox.showerror("Label Error", msg)
return
num_added = len(converted_items)
# Add to accumulator
if self.accumulated_results:
self.accumulated_results["list"].extend(converted_items)
# Update name just in case (though usually consistent per group)
if "name" in current_data and current_data["name"] != "Continued":
self.accumulated_results["name"] = current_data["name"]
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)
messagebox.showerror("JSON Error", msg)
return # Abort advancement
self.history.append((self.current_pil_image, self.current_json_path,
self.current_meta, num_added))
# Advance UI
self.is_viewing = False
self.label.config(image="", text="Loading next...")
def on_open_pdf(self, event):
if self.is_viewing and self.current_json_path:
a = self.current_json_path.stem.split('_')[0] + ".pdf"
pdf_path = self.current_json_path.with_name(a)
print(f"Opening {pdf_path}")
open_path(pdf_path)
def on_open_interro(self, event):
if self.is_viewing and self.current_json_path:
# Check local directory first
local_accent = self.base_dir / "énoncé.pdf"
local_plain = self.base_dir / "enonce.pdf"
if local_accent.exists():
pdf_path = str(local_accent)
elif local_plain.exists():
pdf_path = local_plain
else:
messagebox.showerror(
"PDF not found",
f"Neither {local_accent.name} nor {local_plain.name} exists in {self.base_dir}.",
)
return
print(f"Opening {pdf_path}")
open_path(pdf_path)
def on_open_ori_pdf(self, event):
if self.is_viewing and self.current_json_path:
new_filename = self.current_json_path.stem.split('_')[0] + ".pdf"
pdf_path = self.base_dir / "Copies Originales" / new_filename
print(f"Opening {pdf_path}")
open_path(pdf_path)
def on_edit(self, event):
if self.is_viewing and self.current_json_path:
print(f"Opening {self.current_json_path}")
open_path(self.current_json_path)
def on_click(self, event):
if not self.is_viewing: return
x = int(event.x / self.scale_factor)
y = int(event.y / self.scale_factor)
w, h = self.orig_size
box = [
int(max(0, y - 5) / h * 1000),
int(max(0, x - 5) / (w- padding) * 1000),
int(min(h, y + 5) / h * 1000),
int(min(w, x + 5) / (w - padding) * 1000),
]
box_str = "{ \"box_2d\": " + str(box) + ", \"label\": \"\" },"
print(f"Copied box at ({x},{y}): {box_str}")
self.root.clipboard_clear()
self.root.clipboard_append(box_str)
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__":
raise SystemExit(main())