940 lines
34 KiB
Python
Executable File
940 lines
34 KiB
Python
Executable File
#!/usr/bin/env -S uv run
|
|
# /// script
|
|
# requires-python = ">=3.14"
|
|
# dependencies = [
|
|
# "bottle>=0.13.0",
|
|
# ]
|
|
# ///
|
|
|
|
import argparse
|
|
import functools
|
|
import os
|
|
import re
|
|
import socket
|
|
import sys
|
|
import threading
|
|
import time
|
|
import webbrowser
|
|
from pathlib import Path
|
|
from datetime import datetime
|
|
from typing import TypedDict
|
|
import bottle
|
|
|
|
IMAGE_EXTENSIONS: set[str] = {".jpg", ".jpeg", ".png", ".webp", ".webm"}
|
|
|
|
|
|
class ImageGroup(TypedDict):
|
|
"""Type definition for grouped images."""
|
|
|
|
base: str
|
|
original: Path
|
|
edits: list[Path]
|
|
|
|
|
|
def find_available_port() -> int:
|
|
"""Find an available random port."""
|
|
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
|
|
s.bind(("127.0.0.1", 0))
|
|
s.listen(1)
|
|
port = s.getsockname()[1]
|
|
return port
|
|
|
|
|
|
def is_image(path: str | Path) -> bool:
|
|
"""Check if file is a supported image."""
|
|
return Path(path).suffix.lower() in IMAGE_EXTENSIONS
|
|
|
|
|
|
def discover_images(paths: list[str]) -> list[Path]:
|
|
"""Discover all images from mixed list of files and folders."""
|
|
images = []
|
|
for path_str in paths:
|
|
path = Path(path_str)
|
|
if path.is_file() and is_image(path):
|
|
images.append(path.resolve())
|
|
elif path.is_dir():
|
|
for item in path.iterdir():
|
|
if item.is_file() and is_image(item):
|
|
images.append(item.resolve())
|
|
return sorted(set(images))
|
|
|
|
|
|
def group_images(images: list[Path]) -> list[ImageGroup]:
|
|
"""
|
|
Group images by original + edits.
|
|
Original: image.jpg
|
|
Edits: image-edit1.jpg, image-edit2.jpg, etc.
|
|
Returns list of dicts: {original: Path, edits: [Path, ...]}
|
|
"""
|
|
groups = {}
|
|
|
|
for img in images:
|
|
stem = img.stem
|
|
|
|
# Check if this is an edit
|
|
match = re.match(r"^(.+?)-edit", stem)
|
|
if match:
|
|
base = match.group(1)
|
|
if base not in groups:
|
|
groups[base] = {"original": None, "edits": []}
|
|
groups[base]["edits"].append(img)
|
|
else:
|
|
# This is an original
|
|
if stem not in groups:
|
|
groups[stem] = {"original": None, "edits": []}
|
|
groups[stem]["original"] = img
|
|
|
|
# Filter out groups without originals, convert to list
|
|
result = []
|
|
for base, group in groups.items():
|
|
if group["original"]:
|
|
result.append(ImageGroup(base=base, original=group["original"], edits=sorted(group["edits"])))
|
|
|
|
return sorted(result, key=lambda g: g["original"].name)
|
|
|
|
|
|
def safe_image_path(filepath: str) -> Path:
|
|
"""Verify image path is safe (prevent directory traversal)."""
|
|
# Search for the image in all_images by filename
|
|
for img in all_images:
|
|
if img.name == filepath:
|
|
return img
|
|
raise ValueError(f"Unauthorized path: {filepath}")
|
|
|
|
|
|
def hardlink_images(src_paths: list[Path | str], dest_dir: str | Path) -> None:
|
|
"""Hard link multiple images to destination directory."""
|
|
dest_path = Path(dest_dir)
|
|
make_dirs(dest_path)
|
|
|
|
for src in src_paths:
|
|
src = Path(src)
|
|
dest = dest_path / src.name
|
|
try:
|
|
# Remove existing file if present
|
|
if dest.exists():
|
|
dest.unlink()
|
|
os.link(src, dest)
|
|
except Exception as e:
|
|
print(f"Error hard linking {src} to {dest}: {e}", file=sys.stderr)
|
|
raise
|
|
|
|
|
|
# Global state
|
|
app = bottle.Bottle()
|
|
all_images = set()
|
|
image_groups = []
|
|
current_group_idx = 0
|
|
last_heartbeat = datetime.now()
|
|
should_exit = False
|
|
picks_dir = Path()
|
|
flags = []
|
|
flag_dirs = {}
|
|
|
|
HTML_TEMPLATE = """
|
|
<!DOCTYPE html>
|
|
<html lang="en">
|
|
<head>
|
|
<meta charset="UTF-8">
|
|
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
|
<title>Image Culler</title>
|
|
<script defer src="https://cdn.jsdelivr.net/npm/alpinejs@3.x.x/dist/cdn.min.js"></script>
|
|
<style>
|
|
* {
|
|
margin: 0;
|
|
padding: 0;
|
|
box-sizing: border-box;
|
|
}
|
|
body {
|
|
font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', sans-serif;
|
|
background: #1e1e1e;
|
|
color: #e0e0e0;
|
|
overflow: hidden;
|
|
height: 100vh;
|
|
}
|
|
.container {
|
|
display: flex;
|
|
height: calc(100vh - 46px);
|
|
gap: 20px;
|
|
padding: 20px;
|
|
}
|
|
.panel {
|
|
display: flex;
|
|
flex-direction: column;
|
|
gap: 15px;
|
|
}
|
|
.left-panel {
|
|
flex: 0 0 40%;
|
|
}
|
|
.right-panel {
|
|
flex: 0 0 60%;
|
|
overflow-y: auto;
|
|
}
|
|
.image-display {
|
|
width: 100%;
|
|
height: 100%;
|
|
background: #2a2a2a;
|
|
border-radius: 8px;
|
|
display: flex;
|
|
align-items: center;
|
|
justify-content: center;
|
|
overflow: hidden;
|
|
border: 2px solid transparent;
|
|
transition: border-color 0.2s;
|
|
}
|
|
.image-display:hover {
|
|
border-color: #4a9eff;
|
|
}
|
|
.image-display.picked {
|
|
border-color: #4ade80;
|
|
background: rgba(74, 222, 128, 0.1);
|
|
}
|
|
.image-display img {
|
|
max-width: 100%;
|
|
max-height: 100%;
|
|
object-fit: contain;
|
|
}
|
|
.image-info {
|
|
background: #2a2a2a;
|
|
padding: 12px;
|
|
border-radius: 8px;
|
|
font-size: 14px;
|
|
}
|
|
dialog {
|
|
background: #2a2a2a;
|
|
color: #e0e0e0;
|
|
border: 1px solid #404040;
|
|
border-radius: 8px;
|
|
padding: 20px;
|
|
max-width: 400px;
|
|
}
|
|
dialog::backdrop {
|
|
background: rgba(0, 0, 0, 0.5);
|
|
}
|
|
.grid {
|
|
display: grid;
|
|
grid-template-columns: repeat(2, 1fr);
|
|
gap: 10px;
|
|
}
|
|
.grid[data-image-count="1"] {
|
|
grid-template-columns: 1fr;
|
|
}
|
|
.grid-item {
|
|
position: relative;
|
|
background: #2a2a2a;
|
|
border-radius: 8px;
|
|
overflow: hidden;
|
|
cursor: pointer;
|
|
border: 2px solid transparent;
|
|
transition: all 0.2s;
|
|
}
|
|
.grid-item:hover {
|
|
border-color: #4a9eff;
|
|
}
|
|
.grid-item img {
|
|
width: 100%;
|
|
height: auto;
|
|
display: block;
|
|
object-fit: contain;
|
|
}
|
|
.grid[data-image-count="1"] .grid-item img {
|
|
max-height: 90vh;
|
|
}
|
|
.grid-item.picked {
|
|
border-color: #4ade80;
|
|
background: rgba(74, 222, 128, 0.1);
|
|
}
|
|
.flag-pill {
|
|
background: #404040;
|
|
color: #e0e0e0;
|
|
border: 1px solid #555;
|
|
}
|
|
.flag-pill:hover {
|
|
background: #505050;
|
|
border-color: #666;
|
|
}
|
|
.flag-pill-active {
|
|
background: #4a9eff;
|
|
color: #000;
|
|
border: 1px solid #4a9eff;
|
|
font-weight: 500;
|
|
}
|
|
.flag-pill-active:hover {
|
|
background: #6ab0ff;
|
|
border-color: #6ab0ff;
|
|
}
|
|
</style>
|
|
</head>
|
|
<body>
|
|
<div x-data="app()" x-init="init()" @keydown.window="handleKeydown($event)">
|
|
<div style="background: #2a2a2a; padding: 12px 20px; border-bottom: 1px solid #404040;">
|
|
<div style="font-size: 16px; font-weight: 500; color: #e0e0e0; word-break: break-all;">
|
|
<span x-text="currentGroup ? currentGroup.original.name : 'No images'"></span>
|
|
</div>
|
|
</div>
|
|
<div class="container">
|
|
<!-- Left Panel: Original Image -->
|
|
<div class="panel left-panel">
|
|
<div class="image-display" :class="{ picked: originalPicked }" style="position: relative;">
|
|
<template x-if="currentGroup">
|
|
<img :src="`/image/${currentGroup.original.name}`" :alt="currentGroup.original.name">
|
|
<template x-if="imageFlagAssignments[currentGroup.original.name]">
|
|
<div style="position: absolute; bottom: 12px; left: 12px; background: #4a9eff; color: #000; padding: 4px 8px; border-radius: 4px; font-size: 12px; font-weight: bold;">
|
|
<span x-text="imageFlagAssignments[currentGroup.original.name]"></span>
|
|
</div>
|
|
</template>
|
|
</template>
|
|
<template x-if="!currentGroup">
|
|
<div style="text-align: center; color: #888;">No images</div>
|
|
</template>
|
|
</div>
|
|
<div class="image-info">
|
|
<template x-if="currentGroup">
|
|
<div>
|
|
<div style="display: flex; align-items: center; gap: 8px;">
|
|
<strong>Position:</strong>
|
|
<input
|
|
type="number"
|
|
:value="currentGroupIdx + 1"
|
|
@keyup.enter="jumpToIndex($event)"
|
|
@change="jumpToIndex($event)"
|
|
min="1"
|
|
:max="imageGroups.length"
|
|
style="width: 90px; background: #404040; color: #e0e0e0; border: 1px solid #555; padding: 4px 8px; border-radius: 4px; font-size: 14px;"
|
|
>
|
|
<span style="color: #888;">/ <span x-text="imageGroups.length"></span></span>
|
|
</div>
|
|
<template x-if="availableFlags.length > 0">
|
|
<div style="margin-top: 10px;">
|
|
<strong style="font-size: 13px;">Flags:</strong>
|
|
<div style="display: flex; flex-wrap: wrap; gap: 6px; margin-top: 6px;">
|
|
<template x-for="(flag, idx) in availableFlags" :key="idx">
|
|
<div
|
|
@click="handleFlagAssignment(flag)"
|
|
:class="imageFlagAssignments[currentGroup.original.name] === flag ? 'flag-pill-active' : 'flag-pill'"
|
|
style="cursor: pointer; padding: 4px 8px; border-radius: 4px; font-size: 12px; white-space: nowrap; transition: all 0.2s;"
|
|
>
|
|
<span x-text="`${flag} (${idx + 1})`"></span>
|
|
</div>
|
|
</template>
|
|
</div>
|
|
</div>
|
|
</template>
|
|
</div>
|
|
</template>
|
|
</div>
|
|
</div>
|
|
|
|
<!-- Right Panel: Edit Variations Grid -->
|
|
<div class="panel right-panel">
|
|
<template x-if="currentGroup && currentGroup.edits.length > 0">
|
|
<div class="grid" :data-image-count="currentGroup.edits.length">
|
|
<template x-for="(edit, idx) in currentGroup.edits" :key="idx">
|
|
<div
|
|
class="grid-item"
|
|
:class="{ picked: pickedEdits.includes(idx) }"
|
|
@click="toggleEditPick(idx)"
|
|
>
|
|
<img :src="`/image/${edit.name}`" :alt="`Edit ${idx + 1}`">
|
|
</div>
|
|
</template>
|
|
</div>
|
|
</template>
|
|
<template x-if="!currentGroup || currentGroup.edits.length === 0">
|
|
<div style="color: #888; text-align: center; padding: 40px;">No variations</div>
|
|
</template>
|
|
</div>
|
|
</div>
|
|
|
|
<!-- Help Dialog -->
|
|
<dialog id="helpDialog">
|
|
<h2>Keyboard Shortcuts</h2>
|
|
<div style="margin-top: 15px; font-size: 14px; line-height: 1.8;">
|
|
<div style="margin-bottom: 10px;">
|
|
<span style="background: #404040; padding: 2px 6px; border-radius: 3px; font-weight: bold;">P</span>
|
|
Toggle pick (all edits or hover one)
|
|
</div>
|
|
<div style="margin-bottom: 10px;">
|
|
<span style="background: #404040; padding: 2px 6px; border-radius: 3px; font-weight: bold;">←→</span>
|
|
Navigate any image
|
|
</div>
|
|
<div style="margin-bottom: 10px;">
|
|
<span style="background: #404040; padding: 2px 6px; border-radius: 3px; font-weight: bold;">K</span>
|
|
Skip to next pending
|
|
</div>
|
|
<div style="margin-top: 15px; font-size: 13px; color: #aaa;">
|
|
<strong>Flags:</strong>
|
|
</div>
|
|
<div>
|
|
<span style="background: #404040; padding: 2px 6px; border-radius: 3px; font-weight: bold;">1-9</span>
|
|
Assign/unassign to flag
|
|
</div>
|
|
<div>
|
|
<span style="background: #404040; padding: 2px 6px; border-radius: 3px; font-weight: bold;">/</span>
|
|
Toggle reject flag
|
|
</div>
|
|
</div>
|
|
<div style="margin-top: 20px; text-align: right;">
|
|
<button @click="closeHelp()" style="background: #404040; color: #e0e0e0; border: none; padding: 8px 16px; border-radius: 4px; cursor: pointer;">Close</button>
|
|
</div>
|
|
</dialog>
|
|
</div>
|
|
|
|
<script>
|
|
function app() {
|
|
return {
|
|
imageGroups: [],
|
|
currentGroupIdx: 0,
|
|
currentGroup: null,
|
|
pickedEdits: [],
|
|
originalPicked: false,
|
|
groupStatus: {},
|
|
availableFlags: [],
|
|
imageFlagAssignments: {},
|
|
|
|
async init() {
|
|
// Start heartbeat
|
|
this.startHeartbeat();
|
|
|
|
// Track mouse position globally for hover detection
|
|
document.addEventListener('mousemove', (e) => {
|
|
window.lastMouseEvent = e;
|
|
});
|
|
|
|
// Notify server when tab/window is closed
|
|
window.addEventListener('beforeunload', () => {
|
|
fetch('/api/shutdown', { method: 'POST' }).catch(() => {});
|
|
});
|
|
|
|
// Load images
|
|
await this.loadImages();
|
|
|
|
// Load first group
|
|
this.loadGroup(0);
|
|
},
|
|
|
|
async loadImages() {
|
|
try {
|
|
const resp = await fetch('/api/images');
|
|
const data = await resp.json();
|
|
this.imageGroups = data.groups;
|
|
this.availableFlags = data.flags_list || [];
|
|
this.groupStatus = {};
|
|
this.imageFlagAssignments = data.flag_assignments || {};
|
|
this.imageGroups.forEach((_, idx) => {
|
|
this.groupStatus[idx] = 'pending';
|
|
});
|
|
} catch (err) {
|
|
console.error('Error loading images:', err.message);
|
|
}
|
|
},
|
|
|
|
loadGroup(idx) {
|
|
if (idx < 0 || idx >= this.imageGroups.length) return;
|
|
this.currentGroupIdx = idx;
|
|
this.currentGroup = this.imageGroups[idx];
|
|
this.pickedEdits = [];
|
|
this.originalPicked = false;
|
|
},
|
|
|
|
getGroupStatus(idx) {
|
|
return this.groupStatus[idx] || 'pending';
|
|
},
|
|
|
|
jumpToIndex(event) {
|
|
const value = parseInt(event.target.value);
|
|
if (!isNaN(value) && value >= 1 && value <= this.imageGroups.length) {
|
|
this.loadGroup(value - 1);
|
|
}
|
|
},
|
|
|
|
toggleEditPick(idx) {
|
|
if (this.pickedEdits.includes(idx)) {
|
|
this.pickedEdits = this.pickedEdits.filter(i => i !== idx);
|
|
} else {
|
|
this.pickedEdits.push(idx);
|
|
}
|
|
},
|
|
|
|
getHoveredElement() {
|
|
// Find the element currently under the mouse pointer
|
|
const mouseEvent = window.lastMouseEvent;
|
|
if (!mouseEvent) return null;
|
|
return document.elementFromPoint(mouseEvent.clientX, mouseEvent.clientY);
|
|
},
|
|
|
|
async handleKeydown(e) {
|
|
// Skip keybindings if focused on input/textarea
|
|
const target = e.target;
|
|
const isInput = target.tagName === 'INPUT' || target.tagName === 'TEXTAREA' || target.contentEditable === 'true';
|
|
if (isInput) {
|
|
return;
|
|
}
|
|
|
|
const key = e.key.toLowerCase();
|
|
|
|
if (key === 'p') {
|
|
e.preventDefault();
|
|
await this.handlePick();
|
|
} else if (key === '?') {
|
|
e.preventDefault();
|
|
this.toggleHelp();
|
|
} else if (key === '/') {
|
|
e.preventDefault();
|
|
await this.handleFlagAssignment('reject');
|
|
} else if (key === 'arrowleft') {
|
|
e.preventDefault();
|
|
let nextIdx = this.currentGroupIdx - 1;
|
|
if (nextIdx < 0) nextIdx = this.imageGroups.length - 1;
|
|
this.loadGroup(nextIdx);
|
|
} else if (key === 'arrowright') {
|
|
e.preventDefault();
|
|
let nextIdx = this.currentGroupIdx + 1;
|
|
if (nextIdx >= this.imageGroups.length) nextIdx = 0;
|
|
this.loadGroup(nextIdx);
|
|
} else if (key === 'k') {
|
|
e.preventDefault();
|
|
for (let i = this.currentGroupIdx + 1; i < this.imageGroups.length; i++) {
|
|
if (this.getGroupStatus(i) === 'pending') {
|
|
this.loadGroup(i);
|
|
return;
|
|
}
|
|
}
|
|
} else if (key >= '1' && key <= '9') {
|
|
e.preventDefault();
|
|
const flagIdx = parseInt(key) - 1;
|
|
if (flagIdx < this.availableFlags.length) {
|
|
const flagName = this.availableFlags[flagIdx];
|
|
await this.handleFlagAssignment(flagName);
|
|
}
|
|
}
|
|
},
|
|
|
|
async handlePick() {
|
|
if (!this.currentGroup) return;
|
|
|
|
const currentStatus = this.getGroupStatus(this.currentGroupIdx);
|
|
|
|
// If already picked, unpick it
|
|
if (currentStatus === 'picked') {
|
|
try {
|
|
const resp = await fetch('/api/unpick', {
|
|
method: 'POST',
|
|
headers: { 'Content-Type': 'application/json' },
|
|
body: JSON.stringify({
|
|
original: this.currentGroup.original.name
|
|
})
|
|
});
|
|
if (resp.ok) {
|
|
this.groupStatus[this.currentGroupIdx] = 'pending';
|
|
this.originalPicked = false;
|
|
this.pickedEdits = [];
|
|
}
|
|
} catch (err) {
|
|
console.error('Error unpicking:', err);
|
|
}
|
|
return;
|
|
}
|
|
|
|
const hoveredEl = this.getHoveredElement();
|
|
const isHoveringOriginal = hoveredEl?.closest('.image-display:not(.right-panel *)');
|
|
|
|
// If hovering over original, pick just the original
|
|
if (isHoveringOriginal) {
|
|
try {
|
|
const resp = await fetch('/api/pick', {
|
|
method: 'POST',
|
|
headers: { 'Content-Type': 'application/json' },
|
|
body: JSON.stringify({
|
|
original: this.currentGroup.original.name,
|
|
edits: [],
|
|
includeOriginal: true
|
|
})
|
|
});
|
|
if (resp.ok) {
|
|
this.originalPicked = true;
|
|
}
|
|
} catch (err) {
|
|
console.error('Error picking original:', err);
|
|
}
|
|
return;
|
|
}
|
|
|
|
// If hovering on edit, pick just that edit
|
|
const hoveredGridItem = hoveredEl?.closest('.grid-item');
|
|
if (hoveredGridItem && this.currentGroup.edits.length > 0) {
|
|
const idx = Array.from(hoveredGridItem.parentElement.querySelectorAll('.grid-item')).indexOf(hoveredGridItem);
|
|
if (idx !== -1) {
|
|
const editName = this.currentGroup.edits[idx].name;
|
|
try {
|
|
const resp = await fetch('/api/pick', {
|
|
method: 'POST',
|
|
headers: { 'Content-Type': 'application/json' },
|
|
body: JSON.stringify({
|
|
original: this.currentGroup.original.name,
|
|
edits: [editName],
|
|
includeOriginal: false
|
|
})
|
|
});
|
|
if (resp.ok) {
|
|
this.pickedEdits.push(idx);
|
|
}
|
|
} catch (err) {
|
|
console.error('Error picking edit:', err);
|
|
}
|
|
return;
|
|
}
|
|
}
|
|
|
|
// If not hovering, pick all selected edits and advance
|
|
let editsToPick = this.pickedEdits.map(i => this.currentGroup.edits[i].name);
|
|
|
|
// If no edits selected but there's exactly 1 edit, pick it automatically
|
|
if (editsToPick.length === 0 && this.currentGroup.edits.length === 1) {
|
|
editsToPick = [this.currentGroup.edits[0].name];
|
|
}
|
|
|
|
// Only proceed if there are edits to pick
|
|
if (editsToPick.length === 0) return;
|
|
|
|
try {
|
|
const resp = await fetch('/api/pick', {
|
|
method: 'POST',
|
|
headers: { 'Content-Type': 'application/json' },
|
|
body: JSON.stringify({
|
|
original: this.currentGroup.original.name,
|
|
edits: editsToPick,
|
|
includeOriginal: false
|
|
})
|
|
});
|
|
if (resp.ok) {
|
|
this.groupStatus[this.currentGroupIdx] = 'picked';
|
|
this.loadGroup(this.currentGroupIdx + 1);
|
|
}
|
|
} catch (err) {
|
|
console.error('Error picking edits:', err);
|
|
}
|
|
},
|
|
|
|
startHeartbeat() {
|
|
setInterval(async () => {
|
|
try {
|
|
await fetch('/api/heartbeat', { method: 'POST' });
|
|
} catch (err) {
|
|
console.error('Heartbeat failed:', err);
|
|
}
|
|
}, 1000);
|
|
},
|
|
|
|
toggleHelp() {
|
|
const dialog = document.getElementById('helpDialog');
|
|
if (dialog.open) {
|
|
dialog.close();
|
|
} else {
|
|
dialog.showModal();
|
|
}
|
|
},
|
|
|
|
closeHelp() {
|
|
document.getElementById('helpDialog').close();
|
|
},
|
|
|
|
async handleFlagAssignment(flagName) {
|
|
if (!this.currentGroup) return;
|
|
|
|
const imageKey = this.currentGroup.original.name;
|
|
|
|
// Toggle: if already assigned to this flag, unassign
|
|
if (this.imageFlagAssignments[imageKey] === flagName) {
|
|
try {
|
|
await fetch('/api/unassign-flag', {
|
|
method: 'POST',
|
|
headers: { 'Content-Type': 'application/json' },
|
|
body: JSON.stringify({
|
|
original: imageKey,
|
|
flag: flagName
|
|
})
|
|
});
|
|
delete this.imageFlagAssignments[imageKey];
|
|
} catch (err) {
|
|
console.error('Error unassigning flag:', err);
|
|
}
|
|
} else {
|
|
// Assign to new flag
|
|
try {
|
|
const resp = await fetch('/api/assign-flag', {
|
|
method: 'POST',
|
|
headers: { 'Content-Type': 'application/json' },
|
|
body: JSON.stringify({
|
|
original: imageKey,
|
|
flag: flagName
|
|
})
|
|
});
|
|
if (resp.ok) {
|
|
this.imageFlagAssignments[imageKey] = flagName;
|
|
}
|
|
} catch (err) {
|
|
console.error('Error assigning flag:', err);
|
|
}
|
|
}
|
|
}
|
|
};
|
|
}
|
|
</script>
|
|
</body>
|
|
</html>
|
|
"""
|
|
|
|
|
|
@app.get("/")
|
|
def index() -> str:
|
|
return HTML_TEMPLATE
|
|
|
|
|
|
@app.get("/api/images")
|
|
def api_images():
|
|
"""Return grouped images and current flag assignments."""
|
|
global image_groups, flags, flag_dirs
|
|
result = []
|
|
for group in image_groups:
|
|
result.append(
|
|
{
|
|
"base": group["base"],
|
|
"original": {"name": group["original"].name, "path": str(group["original"])},
|
|
"edits": [{"name": e.name, "path": str(e)} for e in group["edits"]],
|
|
}
|
|
)
|
|
|
|
# Build flag assignments by checking which flag directories contain which original images
|
|
flag_assignments = {}
|
|
for flag, flag_dir in flag_dirs.items():
|
|
if flag_dir.exists():
|
|
for item in flag_dir.iterdir():
|
|
if item.is_file():
|
|
# Map the file name to its flag
|
|
flag_assignments[item.name] = flag
|
|
|
|
return {"groups": result, "flags_list": flags, "flag_assignments": flag_assignments}
|
|
|
|
|
|
@app.get("/image/<filename>")
|
|
def serve_image(filename: str):
|
|
"""Serve an image file safely."""
|
|
try:
|
|
target = safe_image_path(filename)
|
|
return bottle.static_file(target.name, root=str(target.parent), mimetype="image/jpeg")
|
|
except ValueError:
|
|
bottle.response.status = 403
|
|
return {"error": "Unauthorized path"}
|
|
|
|
|
|
@app.post("/api/pick")
|
|
def api_pick():
|
|
"""Hard link picked images to _picks directory."""
|
|
global image_groups, current_group_idx, picks_dir
|
|
|
|
# Create directory before writing
|
|
make_dirs(picks_dir)
|
|
|
|
data = bottle.request.json
|
|
original_name = data.get("original")
|
|
edits_names = data.get("edits", [])
|
|
include_original = data.get("includeOriginal", False)
|
|
|
|
files_to_pick = []
|
|
|
|
# Add original if requested
|
|
if include_original:
|
|
for img in all_images:
|
|
if img.name == original_name:
|
|
files_to_pick.append(img)
|
|
break
|
|
|
|
# Add edits by finding them in all_images
|
|
for edit_name in edits_names:
|
|
for img in all_images:
|
|
if img.name == edit_name:
|
|
files_to_pick.append(img)
|
|
break
|
|
|
|
try:
|
|
hardlink_images(files_to_pick, picks_dir)
|
|
return {"status": "ok"}
|
|
except Exception as e:
|
|
bottle.response.status = 500
|
|
return {"error": str(e)}
|
|
|
|
|
|
@app.post("/api/unpick")
|
|
def api_unpick():
|
|
"""Remove picked images from _picks directory."""
|
|
global picks_dir
|
|
|
|
data = bottle.request.json
|
|
original_name = data.get("original")
|
|
|
|
try:
|
|
# Find and remove the original and any edits from picks directory
|
|
if picks_dir.exists():
|
|
for item in picks_dir.iterdir():
|
|
if item.is_file() and (item.name == original_name or item.stem.startswith(original_name.rsplit(".", 1)[0] + "-edit")):
|
|
item.unlink()
|
|
return {"status": "ok"}
|
|
except Exception as e:
|
|
bottle.response.status = 500
|
|
return {"error": str(e)}
|
|
|
|
|
|
@app.post("/api/assign-flag")
|
|
def api_assign_flag():
|
|
"""Assign an image to a flag."""
|
|
global flag_dirs, all_images
|
|
|
|
data = bottle.request.json
|
|
original_name = data.get("original")
|
|
flag_name = data.get("flag")
|
|
|
|
# Find original file
|
|
original_path = None
|
|
for img in all_images:
|
|
if img.name == original_name:
|
|
original_path = img
|
|
break
|
|
|
|
if not original_path or flag_name not in flag_dirs:
|
|
bottle.response.status = 400
|
|
return {"error": "Invalid image or flag"}
|
|
|
|
# Create directory before writing
|
|
make_dirs(flag_dirs[flag_name])
|
|
|
|
try:
|
|
hardlink_images([original_path], flag_dirs[flag_name])
|
|
return {"status": "ok"}
|
|
except Exception as e:
|
|
bottle.response.status = 500
|
|
return {"error": str(e)}
|
|
|
|
|
|
@app.post("/api/unassign-flag")
|
|
def api_unassign_flag():
|
|
"""Remove an image from a flag."""
|
|
global flag_dirs
|
|
|
|
data = bottle.request.json
|
|
original_name = data.get("original")
|
|
flag_name = data.get("flag")
|
|
|
|
if flag_name not in flag_dirs:
|
|
bottle.response.status = 400
|
|
return {"error": "Invalid flag"}
|
|
|
|
try:
|
|
target_file = flag_dirs[flag_name] / original_name
|
|
if target_file.exists():
|
|
target_file.unlink()
|
|
return {"status": "ok"}
|
|
except Exception as e:
|
|
bottle.response.status = 500
|
|
return {"error": str(e)}
|
|
|
|
|
|
@app.post("/api/heartbeat")
|
|
def api_heartbeat():
|
|
"""Update last heartbeat timestamp."""
|
|
global last_heartbeat
|
|
last_heartbeat = datetime.now()
|
|
return {"status": "ok"}
|
|
|
|
|
|
@app.post("/api/shutdown")
|
|
def api_shutdown():
|
|
"""Shutdown the server."""
|
|
global should_exit
|
|
should_exit = True
|
|
return {"status": "ok"}
|
|
|
|
|
|
@functools.cache
|
|
def make_dirs(path: Path) -> None:
|
|
"""Create directory if it doesn't exist."""
|
|
path.mkdir(parents=True, exist_ok=True)
|
|
|
|
|
|
def heartbeat_monitor() -> None:
|
|
"""Monitor for shutdown signal."""
|
|
global should_exit
|
|
|
|
while not should_exit:
|
|
time.sleep(0.5)
|
|
|
|
print("\nShutting down server.", file=sys.stderr)
|
|
os._exit(0)
|
|
|
|
|
|
def main() -> None:
|
|
global all_images, image_groups, last_heartbeat, picks_dir, flags, flag_dirs
|
|
|
|
parser = argparse.ArgumentParser(description="Image culling app for AI image2image transforms")
|
|
parser.add_argument("paths", nargs="+", help="Image files or folders to cull")
|
|
parser.add_argument("--picks-dir", default=None, help="Directory for picked images (default: $cwd/_picks)")
|
|
parser.add_argument("--rejects-dir", default=None, help="Directory for rejected images (default: $cwd/_rejects)")
|
|
parser.add_argument("--flags", default=None, help="Comma-separated flag names (e.g., 'a,b,c,d')")
|
|
parser.add_argument("--flags-dir", default=None, help="Directory for flag subdirectories (default: $cwd)")
|
|
|
|
args = parser.parse_args()
|
|
|
|
# Set directories
|
|
cwd = Path.cwd()
|
|
picks_dir = Path(args.picks_dir or (cwd / "_picks"))
|
|
rejects_dir = Path(args.rejects_dir or (cwd / "_rejects"))
|
|
flags_dir = Path(args.flags_dir or cwd)
|
|
|
|
# Set up flags
|
|
if args.flags:
|
|
flags = [f.strip() for f in args.flags.split(",")]
|
|
for flag in flags:
|
|
flag_dir = flags_dir / f"_picks_flag_{flag}"
|
|
flag_dirs[flag] = flag_dir
|
|
|
|
# Always add reject as a flag at the end
|
|
flags.append("reject")
|
|
flag_dirs["reject"] = rejects_dir
|
|
|
|
# Discover images
|
|
images = discover_images(args.paths)
|
|
all_images = set(images)
|
|
|
|
if not images:
|
|
print("No images found in provided paths.", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
# Group images
|
|
image_groups = group_images(images)
|
|
|
|
if not image_groups:
|
|
print("No original images found (looking for files matching naming convention).", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
print(f"Found {len(image_groups)} original image(s) with variations", file=sys.stderr)
|
|
|
|
# Find available port
|
|
port = find_available_port()
|
|
url = f"http://127.0.0.1:{port}"
|
|
|
|
# Start heartbeat monitor
|
|
monitor_thread = threading.Thread(target=heartbeat_monitor, daemon=True)
|
|
monitor_thread.start()
|
|
|
|
# Open browser
|
|
print(f"Opening browser at {url}", file=sys.stderr)
|
|
webbrowser.open(url)
|
|
|
|
# Start Flask server
|
|
last_heartbeat = datetime.now()
|
|
bottle.run(app, host="127.0.0.1", port=port, quiet=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|