Seed image label review tool (PyQt6) — blind review with keyboard controls, back/undo, CSV+JSON output

review_exported.py

# review_exported.py
"""
Simplified review tool that loads pre-exported composite images from
output/review_images/. No data dependencies — just the PNGs.

Arrow keys: Left=EMERGED, Down=AMBIGUOUS, Right=NOT EMERGED
Backspace/Up: Go back to previous image (re-label)
Press 'q' or Escape to quit. Progress saves automatically.

Usage:
    pip install pyqt6
    python review_exported.py
"""

import os
import sys
import csv
import json

from PIL import Image
from PyQt6.QtWidgets import (
    QApplication, QLabel, QMainWindow, QVBoxLayout, QWidget,
    QInputDialog,
)
from PyQt6.QtGui import QPixmap, QImage, QKeyEvent
from PyQt6.QtCore import Qt

IMG_DIR = "output/review_images"


def save_results(reviewed, json_path):
    """Save as both JSON and CSV."""
    with open(json_path, "w") as f:
        json.dump(reviewed, f, indent=2)

    csv_path = json_path.replace(".json", ".csv")
    with open(csv_path, "w", newline="") as f:
        writer = csv.writer(f)
        writer.writerow(["image", "label"])
        for fname in sorted(reviewed.keys()):
            writer.writerow([fname, reviewed[fname]])


def log_progress(reviewed, log_path, msg=""):
    emerged = sum(1 for v in reviewed.values() if v == "emerged")
    ambiguous = sum(1 for v in reviewed.values() if v == "ambiguous")
    not_emerged = sum(1 for v in reviewed.values() if v == "not_emerged")
    total = len(reviewed)

    line = f"Total: {total} | Emerged: {emerged} | Ambiguous: {ambiguous} | Not Emerged: {not_emerged}"
    if msg:
        line = f"{msg} | {line}"

    print(line)
    with open(log_path, "a") as f:
        f.write(line + "\n")


class ReviewWindow(QMainWindow):
    def __init__(self, images, reviewed, output_path, log_path):
        super().__init__()
        self.images = images
        self.reviewed = reviewed
        self.output_path = output_path
        self.log_path = log_path
        self.history = []  # stack of visited indices for back navigation

        # Find starting index (first unreviewed)
        self.current_idx = 0
        for i, fname in enumerate(images):
            if fname not in reviewed:
                self.current_idx = i
                break
        else:
            # All reviewed — start at end
            self.current_idx = len(images)

        central = QWidget(self)
        layout = QVBoxLayout(central)
        layout.setContentsMargins(0, 0, 0, 0)
        layout.setSpacing(0)

        self.progress_label = QLabel(self)
        self.progress_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.progress_label.setStyleSheet(
            "background-color: black; color: white; font-size: 20px; padding: 10px;"
        )
        self.progress_label.setFixedHeight(50)

        self.panels_label = QLabel("Left: RGB  |  Centre: 405nm  |  Right: 630nm", self)
        self.panels_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.panels_label.setStyleSheet(
            "background-color: black; color: #888; font-size: 14px; padding: 4px;"
        )
        self.panels_label.setFixedHeight(25)

        self.label = QLabel(self)
        self.label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.label.setStyleSheet("background-color: black;")

        self.status_label = QLabel(self)
        self.status_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        self.status_label.setFixedHeight(30)

        layout.addWidget(self.progress_label)
        layout.addWidget(self.panels_label)
        layout.addWidget(self.label, 1)
        layout.addWidget(self.status_label)

        self.setCentralWidget(central)
        self.showFullScreen()

        self.update_progress()
        self.show_image()

    def update_progress(self):
        done = len(self.reviewed)
        total = len(self.images)
        pct = 100 * done / total if total > 0 else 0
        remaining = total - done

        # Show current label if image already reviewed
        label_info = ""
        if self.current_idx < len(self.images):
            fname = self.images[self.current_idx]
            existing = self.reviewed.get(fname)
            if existing:
                label_info = f"   [CURRENT: {existing.upper()}]"

        self.progress_label.setText(
            f"{done}/{total} ({pct:.1f}%) — {remaining} left{label_info}   |   "
            f"← EMERGED   ↓ AMBIGUOUS   → NOT EMERGED   "
            f"Backspace/↑ BACK   Q quit"
        )

    def update_status(self):
        if self.current_idx >= len(self.images):
            self.status_label.setText("")
            self.status_label.setFixedHeight(0)
            return
        fname = self.images[self.current_idx]
        existing = self.reviewed.get(fname)
        if existing:
            self.status_label.setFixedHeight(30)
            self.status_label.setText(f"Current label: {existing.upper()}")
            self.status_label.setStyleSheet(
                "background-color: #333; color: #ffcc00; font-size: 16px; "
                "font-weight: bold; padding: 4px;"
            )
        else:
            self.status_label.setFixedHeight(0)
            self.status_label.setText("")

    def show_image(self):
        if self.current_idx >= len(self.images):
            self.label.setText("All done! Press Q to quit.")
            self.label.setStyleSheet(
                "background-color: black; color: white; font-size: 48px;"
            )
            self.status_label.setText("")
            return

        fname = self.images[self.current_idx]
        img = Image.open(os.path.join(IMG_DIR, fname)).convert("RGB")

        data = img.tobytes("raw", "RGB")
        qimg = QImage(data, img.width, img.height, 3 * img.width,
                       QImage.Format.Format_RGB888)
        pixmap = QPixmap.fromImage(qimg)

        screen = self.screen().size()
        scaled = pixmap.scaled(screen, Qt.AspectRatioMode.KeepAspectRatio,
                                Qt.TransformationMode.SmoothTransformation)
        self.label.setPixmap(scaled)
        self.update_status()

    def keyPressEvent(self, event: QKeyEvent):
        key = event.key()

        if key in (Qt.Key.Key_Q, Qt.Key.Key_Escape):
            self.save_and_quit()
            return

        # Back navigation
        if key in (Qt.Key.Key_Backspace, Qt.Key.Key_Up):
            if self.history:
                self.current_idx = self.history.pop()
                self.label.clear()
                self.update_progress()
                self.show_image()
            return

        if self.current_idx >= len(self.images):
            return

        human_label = None
        if key == Qt.Key.Key_Left:
            human_label = "emerged"
        elif key == Qt.Key.Key_Down:
            human_label = "ambiguous"
        elif key == Qt.Key.Key_Right:
            human_label = "not_emerged"

        if human_label:
            fname = self.images[self.current_idx]
            self.reviewed[fname] = human_label

            save_results(self.reviewed, self.output_path)

            log_progress(self.reviewed, self.log_path,
                         f"{fname}: {human_label.upper()}")

            self.history.append(self.current_idx)
            self.update_progress()
            self.current_idx += 1
            # Advance to next unreviewed (but don't skip if going forward
            # normally — user might want to re-review sequentially)
            self.show_image()

    def save_and_quit(self):
        save_results(self.reviewed, self.output_path)

        if self.reviewed:
            emerged = sum(1 for v in self.reviewed.values() if v == "emerged")
            ambiguous = sum(1 for v in self.reviewed.values() if v == "ambiguous")
            not_emerged = sum(1 for v in self.reviewed.values() if v == "not_emerged")

            print(f"\nFinal summary ({len(self.reviewed)} reviewed):")
            print(f"  Emerged: {emerged}, Ambiguous: {ambiguous}, Not Emerged: {not_emerged}")
            print(f"  Saved to: {self.output_path}")

        self.close()


def main():
    images = sorted(f for f in os.listdir(IMG_DIR) if f.endswith(".png"))
    print(f"Found {len(images)} images in {IMG_DIR}")

    # Ask for reviewer name / output file
    app = QApplication(sys.argv)
    name, ok = QInputDialog.getText(
        None, "Reviewer name",
        "Enter your name (used for output filename):",
    )
    if not ok or not name.strip():
        print("Cancelled.")
        sys.exit(0)

    name = name.strip().replace(" ", "_").lower()
    output_path = f"output/review_{name}.json"
    log_path = f"output/review_{name}.log"

    # Load existing progress for this reviewer
    reviewed = {}
    if os.path.exists(output_path):
        with open(output_path) as f:
            reviewed = json.load(f)
        print(f"Resuming: {len(reviewed)} already reviewed")
        log_progress(reviewed, log_path, "RESUMED")
    else:
        with open(log_path, "w") as f:
            f.write(f"=== Review Log: {name} ===\n")

    print(f"Output: {output_path}")
    print("\nControls: ← EMERGED | ↓ AMBIGUOUS | → NOT EMERGED | Backspace/↑ BACK | Q/Esc quit\n")

    window = ReviewWindow(images, reviewed, output_path, log_path)
    sys.exit(app.exec())


if __name__ == "__main__":
    main()