Initial commit
This commit is contained in:
+272
@@ -0,0 +1,272 @@
|
||||
import sys
|
||||
import json
|
||||
import threading
|
||||
import queue
|
||||
import subprocess
|
||||
import tkinter as tk
|
||||
from pathlib import Path
|
||||
from PIL import Image, ImageDraw, ImageFont, ImageTk
|
||||
from pypdf import PdfReader
|
||||
|
||||
# --- Configuration & Globals ---
|
||||
padding = 60 # White margin to the right
|
||||
|
||||
|
||||
image_queue = queue.Queue(maxsize=5) # Buffer a few images ahead
|
||||
try:
|
||||
font = ImageFont.truetype("DejaVuSans.ttf", size=30)
|
||||
except OSError:
|
||||
font = ImageFont.load_default()
|
||||
|
||||
# --- Processing Logic (Worker Thread) ---
|
||||
|
||||
def page_number(b, nb_pages):
|
||||
column_width = 1000 // nb_pages
|
||||
center_x = (b[1] + b[3]) // 2
|
||||
return center_x // column_width
|
||||
|
||||
def prepare_image(image_path: str, bounding_boxes, all_labels, nb_pages):
|
||||
"""
|
||||
Draws boxes on the image and returns the PIL Image object.
|
||||
Does NOT display it.
|
||||
"""
|
||||
im = Image.open(image_path)
|
||||
# Ensure image is loaded so we can pass it between threads safely
|
||||
im.load()
|
||||
|
||||
width, height = im.size
|
||||
|
||||
# Add white padding to the right
|
||||
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]))
|
||||
|
||||
last_label_index = -1
|
||||
|
||||
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:
|
||||
color = "red"
|
||||
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:
|
||||
# draw.text((abs_x_min + 8, abs_y_min + 6), label, fill=color, font=font)
|
||||
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
|
||||
|
||||
def worker_thread(base_dir, files_to_process, all_labels):
|
||||
"""
|
||||
Iterates through files, processes them, and puts them in the queue.
|
||||
"""
|
||||
for img_path in files_to_process:
|
||||
json_path = base_dir / f"{img_path.stem}.json"
|
||||
pdf_path = base_dir / f"{img_path.stem}.pdf"
|
||||
|
||||
nb_pages = 1
|
||||
if pdf_path.exists():
|
||||
try:
|
||||
nb_pages = len(PdfReader(pdf_path).pages)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if json_path.exists():
|
||||
try:
|
||||
with open(json_path, 'r') as f:
|
||||
json_result = json.load(f)
|
||||
|
||||
bb_list = json_result.get("list", [])
|
||||
print(f"Processing {img_path.name}...")
|
||||
|
||||
# Draw boxes
|
||||
pil_image = prepare_image(str(img_path), bb_list, all_labels, nb_pages)
|
||||
|
||||
# Block if queue is full (waiting for user to view)
|
||||
image_queue.put((pil_image, json_path))
|
||||
except Exception as e:
|
||||
print(f"Error processing {img_path.name}: {e}")
|
||||
|
||||
# Sentinel to indicate finished
|
||||
image_queue.put((None, None))
|
||||
|
||||
# --- GUI Logic (Main Thread) ---
|
||||
|
||||
class ImageViewer:
|
||||
def __init__(self, root, base_dir):
|
||||
self.root = root
|
||||
self.base_dir = base_dir
|
||||
self.root.title("Bounding Box Viewer")
|
||||
|
||||
# UI Elements
|
||||
self.label = tk.Label(root, text="Waiting for images...")
|
||||
self.label.pack(expand=True, fill="both")
|
||||
|
||||
# State
|
||||
self.current_image = None
|
||||
self.current_json_path = None
|
||||
self.is_viewing = False
|
||||
self.scale_factor = 1.0 # To track resizing
|
||||
self.orig_size = (1, 1) # To track original dimensions
|
||||
|
||||
# Input Bindings
|
||||
self.root.bind('<Return>', self.on_enter)
|
||||
self.root.bind('e', self.on_edit)
|
||||
self.root.bind('o', self.on_open_pdf) # <--- 3. Add Key Binding
|
||||
self.root.bind('<Escape>', lambda e: self.root.quit())
|
||||
self.label.bind('<Button-1>', self.on_click) # Bind left mouse click
|
||||
|
||||
# Start polling queue
|
||||
self.poll_queue()
|
||||
|
||||
def poll_queue(self):
|
||||
if not self.is_viewing:
|
||||
try:
|
||||
pil_image, json_path = image_queue.get_nowait()
|
||||
|
||||
if pil_image is None:
|
||||
print("All images processed.")
|
||||
self.root.quit() # Stop the program
|
||||
return
|
||||
|
||||
self.display_image(pil_image, json_path)
|
||||
except queue.Empty:
|
||||
pass
|
||||
self.root.after(100, self.poll_queue)
|
||||
|
||||
def on_open_pdf(self, event):
|
||||
if self.is_viewing and self.current_json_path:
|
||||
# Replace .json extension with .pdf
|
||||
pdf_path = self.current_json_path.with_suffix(".pdf")
|
||||
|
||||
print(f"Opening {pdf_path}")
|
||||
# Use subprocess to run xdg-open without blocking
|
||||
subprocess.Popen(['xdg-open', str(pdf_path)])
|
||||
|
||||
def display_image(self, pil_image, json_path):
|
||||
self.orig_size = pil_image.size
|
||||
self.scale_factor = 1.0
|
||||
|
||||
# Resize if too large for screen
|
||||
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="")
|
||||
self.current_json_path = json_path
|
||||
self.is_viewing = True
|
||||
self.root.lift()
|
||||
|
||||
def on_enter(self, event):
|
||||
if self.is_viewing:
|
||||
print("Next...")
|
||||
self.is_viewing = False
|
||||
self.label.config(image="", text="Loading next...")
|
||||
|
||||
def on_edit(self, event):
|
||||
if self.is_viewing and self.current_json_path:
|
||||
print(f"Opening {self.current_json_path}")
|
||||
subprocess.Popen(['xdg-open', str(self.current_json_path)])
|
||||
|
||||
def on_click(self, event):
|
||||
if not self.is_viewing: return
|
||||
|
||||
# Map click to original image coordinates
|
||||
x = int(event.x / self.scale_factor)
|
||||
y = int(event.y / self.scale_factor)
|
||||
w, h = self.orig_size
|
||||
|
||||
# Create 10px box (5px radius)
|
||||
# Coordinate format: [y_min, x_min, y_max, x_max] (0-1000 scale)
|
||||
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)
|
||||
|
||||
if __name__ == "__main__":
|
||||
if len(sys.argv) < 2:
|
||||
print("Usage: python plotting_gui.py <directory_or_file>")
|
||||
sys.exit(1)
|
||||
|
||||
input_path = Path(sys.argv[1])
|
||||
files_to_process = []
|
||||
|
||||
if input_path.is_file():
|
||||
# File mode
|
||||
base_dir = input_path.parent
|
||||
stem = input_path.stem
|
||||
|
||||
# Try to locate the image in Cutleft directory
|
||||
img_path = base_dir / "Cutleft" / f"{stem}.jpg"
|
||||
|
||||
# Fallback: Check if user provided the jpg inside Cutleft directly
|
||||
if not img_path.exists() and input_path.parent.name == "Cutleft" and input_path.suffix.lower() == ".jpg":
|
||||
base_dir = input_path.parent.parent
|
||||
img_path = input_path
|
||||
|
||||
if not img_path.exists():
|
||||
print(f"Error: Could not find image at {img_path}")
|
||||
sys.exit(1)
|
||||
|
||||
files_to_process = [img_path]
|
||||
|
||||
else:
|
||||
# Directory mode
|
||||
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 = list(filter(None, (base_dir / "labels").read_text().splitlines()))
|
||||
except FileNotFoundError:
|
||||
all_labels = []
|
||||
|
||||
# Start Processing Thread
|
||||
t = threading.Thread(target=worker_thread, args=(base_dir, files_to_process, all_labels))
|
||||
t.daemon = True # Kill thread if main app closes
|
||||
t.start()
|
||||
|
||||
# Start GUI
|
||||
root = tk.Tk()
|
||||
app = ImageViewer(root, base_dir)
|
||||
root.mainloop()
|
||||
Reference in New Issue
Block a user