I just had one of those wonderful moments where one of my free-time projects actually turns out to be useful in a real-world task.
So in a project for one of my university courses we have to train a neural network. Not exactly on-topic for my studies, but oh well. Unfortunately, the project involved getting the training data ourselves. Even more unfortunately, the way we took the pictures made it impossible to pre-sort them.
So here I was sitting with a directory of a few thousand pictures that needed labeling. We are sorting chess pieces by the way. I don't think there is any common tool out there to sort images like this and going through a normal file manager would have been insanely painful.
Luckily two things saved me:
My image viewer of choice, imv, can be remote-controlled.
I have written a TUI library for zig, allowing me to create user interfaces in basically no time.
So I created a little program to help with this task. It goes through all images in a directory, commanding the image viewer to display them and subsequently allowing me to quickly select the right label with keyboard shortcuts. For example if the image shows a white pawn, I press "wp" and the image gets automatically renamed. I implemented this such that I could also press "pw", because pressing keys in the right order gets kinda hard when you're tired. If the image is unusable (we picked up a decent amount of shadows, specs of dust and yes even rain drops with our detection method; don't ask), I press "r" for "reject".
Still super annoying; I went through four Stargate SG1 episodes and three cups of coffee for the first set alone. Yes, I basically CAPTCHA'd myself. But a considerable improvement over doing this manually. And all I had to do for this was read the manpage of my image viewer and slightly adapt one of zig-spoons example programs, which all-in-all took me about half an hour. That's basically nothing compared to the time sorting the pictures without this tools would have taken me.
Turns out that with the right libraries, zig is a great fit even for super-specific single-use programs like this one.
const std = @import("std");
const heap = std.heap;
const math = std.math;
const mem = std.mem;
const os = std.os;
const fs = std.fs;
const fmt = std.fmt;
const spoon = @import("spoon");
const gpa = heap.page_allocator;
var arena: heap.ArenaAllocator = undefined;
var term: spoon.Term = undefined;
var loop: bool = true;
var files: std.ArrayList([]const u8) = undefined;
var file_index: usize = 0;
var running_index: usize = 0;
const Piece = enum { none, king, queen, pawn, bishop, horse, tower };
const Colour = enum { none, black, white, reject };
var current_piece: Piece = .none;
var current_colour: Colour = .none;
var imv_pid: [*:0]const u8 = undefined;
pub fn main() !void {
imv_pid = os.argv[1];
arena = heap.ArenaAllocator.init(gpa);
defer arena.deinit();
const arloc = arena.allocator();
try term.init(.{});
defer term.deinit();
// Get all files in cwd.
{
files = try std.ArrayList([]const u8).initCapacity(arloc, 2000);
var dir = try fs.cwd().openIterableDir(".", .{});
defer dir.close();
var it = dir.iterate();
while (try it.next()) |entry| {
if (!mem.endsWith(u8, entry.name, ".jpg")) continue;
try files.append(try arloc.dupe(u8, entry.name));
}
if (files.items.len == 0) return;
}
try os.sigaction(os.SIG.WINCH, &os.Sigaction{
.handler = .{ .handler = handleSigWinch },
.mask = os.empty_sigset,
.flags = 0,
}, null);
var fds: [1]os.pollfd = undefined;
fds[0] = .{
.fd = term.tty.handle,
.events = os.POLL.IN,
.revents = undefined,
};
try term.uncook(.{});
defer term.cook() catch {};
try term.fetchSize();
try term.setWindowTitle("sort-o-matic", .{});
try render();
try imvOpen();
var buf: [16]u8 = undefined;
while (loop) {
_ = try os.poll(&fds, -1);
const read = try term.readInput(&buf);
var it = spoon.inputParser(buf[0..read]);
while (it.next()) |in| {
if (in.eqlDescription("C-c")) {
loop = false;
break;
} else if (in.eqlDescription("w")) {
current_colour = .white;
} else if (in.eqlDescription("b")) {
current_colour = .black;
} else if (in.eqlDescription("r")) {
current_colour = .reject;
} else if (in.eqlDescription("p")) {
current_piece = .pawn;
} else if (in.eqlDescription("h")) {
current_piece = .horse;
} else if (in.eqlDescription("B")) {
current_piece = .bishop;
} else if (in.eqlDescription("q")) {
current_piece = .queen;
} else if (in.eqlDescription("k")) {
current_piece = .king;
} else if (in.eqlDescription("t")) {
current_piece = .tower;
}
}
if (current_colour == .reject or (current_colour != .none and current_piece != .none)) {
try imvClose();
try mvFile();
file_index += 1;
if (file_index >= files.items.len) break;
try imvOpen();
}
try render();
}
}
fn render() !void {
var rc = try term.getRenderContext();
defer rc.done() catch {};
try rc.clear();
if (term.width < 6) {
try rc.setAttribute(.{ .fg = .red, .bold = true });
try rc.writeAllWrapping("Terminal too small!");
return;
}
try rc.moveCursorTo(0, 0);
try rc.setAttribute(.{ .fg = .green, .reverse = true });
var rpw = rc.restrictedPaddingWriter(term.width);
try rpw.writer().writeAll(" sort-o-matic");
try rpw.pad();
try rc.moveCursorTo(2, 0);
try rc.setAttribute(.{ .bold = true });
rpw = rc.restrictedPaddingWriter(term.width);
try rpw.writer().print(" [{}/{}] {s}", .{ file_index, files.items.len, files.items[file_index] });
try rpw.finish();
try rc.moveCursorTo(4, 0);
rpw = rc.restrictedPaddingWriter(term.width);
try rpw.writer().writeByte(' ');
try rc.setAttribute(.{ .reverse = current_colour == .black });
try rpw.writer().writeAll("[b]lack");
try rc.setAttribute(.{});
try rpw.writer().writeByte(' ');
try rc.setAttribute(.{ .reverse = current_colour == .white });
try rpw.writer().writeAll("[w]hite");
try rpw.finish();
try rc.moveCursorTo(6, 0);
rpw = rc.restrictedPaddingWriter(term.width);
try rc.setAttribute(.{});
try rpw.writer().writeByte(' ');
try rc.setAttribute(.{ .reverse = current_piece == .queen });
try rpw.writer().writeAll("[q]ueen");
try rc.setAttribute(.{});
try rpw.writer().writeByte(' ');
try rc.setAttribute(.{ .reverse = current_piece == .king });
try rpw.writer().writeAll("[k]ing");
try rc.setAttribute(.{});
try rpw.writer().writeByte(' ');
try rc.setAttribute(.{ .reverse = current_piece == .horse });
try rpw.writer().writeAll("[h]orse");
try rc.setAttribute(.{});
try rpw.writer().writeByte(' ');
try rc.setAttribute(.{ .reverse = current_piece == .pawn });
try rpw.writer().writeAll("[p]awn");
try rc.setAttribute(.{});
try rpw.writer().writeByte(' ');
try rc.setAttribute(.{ .reverse = current_piece == .tower });
try rpw.writer().writeAll("[t]ower");
try rc.setAttribute(.{});
try rpw.writer().writeByte(' ');
try rc.setAttribute(.{ .reverse = current_piece == .bishop });
try rpw.writer().writeAll("[B]bishop");
try rc.setAttribute(.{});
try rpw.finish();
}
fn handleSigWinch(_: c_int) callconv(.C) void {
term.fetchSize() catch {};
render() catch {};
}
/// Custom panic handler, so that we can try to cook the terminal on a crash,
/// as otherwise all messages will be mangled.
pub fn panic(msg: []const u8, trace: ?*std.builtin.StackTrace, ret_addr: ?usize) noreturn {
@setCold(true);
term.cook() catch {};
std.builtin.default_panic(msg, trace, ret_addr);
}
fn imvOpen() !void {
const cmd = try fmt.allocPrintZ(gpa, "imv-msg {s} open {s}", .{ imv_pid, files.items[file_index] });
try runBackground(cmd);
}
fn imvClose() !void {
const cmd = try fmt.allocPrintZ(gpa, "imv-msg {s} close", .{imv_pid});
try runBackground(cmd);
}
fn mvFile() !void {
defer {
current_colour = .none;
current_piece = .none;
running_index += 1;
}
const cmd = if (current_colour == .reject)
try fmt.allocPrintZ(gpa, "mv \"{s}\" \"reject-{}.jpg\"", .{ files.items[file_index], running_index })
else
try fmt.allocPrintZ(gpa, "mv \"{s}\" \"{s}-{s}-{}.jpg\"", .{ files.items[file_index], @tagName(current_colour), @tagName(current_piece), running_index });
defer gpa.free(cmd);
try runBackground(cmd);
}
fn runBackground(cmd: [:0]const u8) !void {
const args = [_:null]?[*:0]const u8{ "/bin/sh", "-c", cmd, null };
const pid = try os.fork();
if (pid == 0) {
const pid2 = os.fork() catch os.fork() catch os.exit(1);
if (pid2 == 0) os.execveZ("/bin/sh", &args, @ptrCast([*:null]?[*:0]u8, os.environ.ptr)) catch os.exit(1);
os.exit(0);
} else {
_ = os.waitpid(pid, 0);
}
}
Top comments (2)
Haha I went through a very similar experience in uni when I wanted to make an app to rate pizzas algorithmically. In my case the classifier was supposed to give a star rating 1-5 so I had slightly different key bindings, but that was it.
Very nice! I hacked a much worse thing for naming pictures a while back but this approach is much better. Thanks!