267 lines
13 KiB
Plaintext
267 lines
13 KiB
Plaintext
// segment-probe: what modules/Segment makes of image files. For each FILE, its segments (class, score,
|
|
// area); with -out FOLDER, the segments drawn over the image as a PNG, to look at; with -check FOLDER,
|
|
// how they compare with PyTorch's own segmentation of the same image (NAME.ref and NAME.ref.txt, from
|
|
// tools/export_mask2former.py --reference). A FILE may be an .mcap: the image it is shown as is read
|
|
// (a video is skipped), and with -store its segments are stored in it (segs and mask chunks).
|
|
//
|
|
// Build: jai first.jai - segment_probe Run: bin/segment-probe FILE... [-out FOLDER] [-check FOLDER] [-store] [-cpu] [-dump-input FOLDER]
|
|
|
|
main :: () {
|
|
args := get_command_line_arguments();
|
|
files: [..] string;
|
|
out_folder, check_folder, dump_folder: string;
|
|
models_folder := tprint("%segment-models", path_strip_filename(get_path_of_running_executable()));
|
|
model_name := "mask2former-swin-base-coco-panoptic";
|
|
prefer_gpu := true;
|
|
store := false;
|
|
i := 1;
|
|
while i < args.count {
|
|
if args[i] == "-out" && i + 1 < args.count { out_folder = args[i + 1]; i += 2; continue; }
|
|
if args[i] == "-check" && i + 1 < args.count { check_folder = args[i + 1]; i += 2; continue; }
|
|
if args[i] == "-models" && i + 1 < args.count { models_folder = args[i + 1]; i += 2; continue; }
|
|
if args[i] == "-model" && i + 1 < args.count { model_name = args[i + 1]; i += 2; continue; }
|
|
if args[i] == "-cpu" { prefer_gpu = false; i += 1; continue; }
|
|
if args[i] == "-store" { store = true; i += 1; continue; }
|
|
if args[i] == "-dump-input" && i + 1 < args.count { dump_folder = args[i + 1]; i += 2; continue; }
|
|
array_add(*files, args[i]);
|
|
i += 1;
|
|
}
|
|
if !files.count { print("Usage: segment-probe FILE... [-out FOLDER] [-check FOLDER] [-cpu] [-model NAME] [-models FOLDER]\n"); exit(2); }
|
|
if out_folder make_directory_if_it_does_not_exist(out_folder, recursive = true);
|
|
|
|
start := current_time_monotonic();
|
|
engine, ok, error := segment_create(models_folder, model_name, prefer_gpu);
|
|
if !ok { print("Cannot start %: %\n", model_name, error); exit(1); }
|
|
defer segment_destroy(engine);
|
|
print("% on %, % classes (made in % ms)%\n", model_name, segment_provider_name(engine), engine.classes.count, ms(start),
|
|
ifx engine.gpu_explanation then tprint(": %", engine.gpu_explanation) else "");
|
|
|
|
failures := 0;
|
|
checked, matched, extra: s64;
|
|
worst_iou: float64 = 1;
|
|
agree, total_pixels, missed_confident: s64;
|
|
model_ms, decode_ms: float64;
|
|
done := 0;
|
|
GROUP :: 8; // images decoded and segmented together
|
|
position := 0;
|
|
while position < files.count {
|
|
group := array_view(files, position, min(GROUP, files.count - position));
|
|
position += group.count;
|
|
|
|
images: [..] Image;
|
|
names: [..] string;
|
|
whiches: [..] string; // for an .mcap, which of its images: "rendered" or "original"
|
|
t := current_time_monotonic();
|
|
for path: group {
|
|
bytes: [] u8;
|
|
which := "original";
|
|
data: string;
|
|
if ends_with_nocase(path, ".mcap") {
|
|
// The image the capture is shown as; never a video's.
|
|
m, shown, shown_which, read_ok, read_error := mcap_read_shown_image(path,, temp);
|
|
if !read_ok { print("%: %\n", path, read_error); failures += 1; continue; }
|
|
if contains(m.text, "VideoCodec:") { print("%: a video, skipped (segmentation is for images)\n", path_filename(path)); continue; }
|
|
stored := mcap_find_chunk(m.front_chunks, mcap_tag(MCAP_SEGMENTS_TAG), model_name);
|
|
if stored {
|
|
s, s_ok := mcap_segments_parse(cast(string) stored.data,, temp);
|
|
if s_ok print("%: already holds % segments by % (%, %)\n", path_filename(path), s.segments.count, model_name, s.date, s.engine);
|
|
}
|
|
bytes, which = shown, shown_which;
|
|
} else {
|
|
read_ok: bool;
|
|
data, read_ok = read_entire_file(path, log_errors = false);
|
|
if !read_ok { print("%: cannot be read\n", path); failures += 1; continue; }
|
|
bytes = cast([] u8) data;
|
|
}
|
|
image, decoded, decode_error := image_decode(bytes); // turned by its EXIF orientation, as mcap shows it
|
|
free(data);
|
|
if !decoded { print("%: cannot decode: %\n", path, decode_error); failures += 1; continue; }
|
|
array_add(*images, image);
|
|
array_add(*names, path);
|
|
array_add(*whiches, which);
|
|
}
|
|
decode_ms += ms(t);
|
|
if !images.count continue;
|
|
|
|
inputs := NewArray(images.count, Segment_Image,, temp);
|
|
for images inputs[it_index] = .{it.pixels, it.width, it.height, it.width * 4};
|
|
if dump_folder { // the model's input, as raw floats, to compare with the checkpoint's processor
|
|
make_directory_if_it_does_not_exist(dump_folder, recursive = true);
|
|
size := engine.model.image_size;
|
|
floats := NewArray(3 * size * size, float,, temp);
|
|
for inputs {
|
|
segment_prepare(engine, it, floats.data);
|
|
write_entire_file(tprint("%/%.input", dump_folder, path_filename(names[it_index])), floats.data, floats.count * size_of(float));
|
|
}
|
|
}
|
|
t = current_time_monotonic();
|
|
results := segment_images(engine, inputs,, temp);
|
|
model_ms += ms(t);
|
|
|
|
for result, n: results {
|
|
path := names[n];
|
|
image := images[n];
|
|
done += 1;
|
|
if result.failed { print("%: failed: %\n", path, result.error); failures += 1; continue; }
|
|
print("% (% x %, map % x %): % segments\n", path_filename(path), image.width, image.height, result.width, result.height, result.segments.count);
|
|
map_area := cast(float64) result.width * result.height;
|
|
for result.segments print(" % % %\n", formatFloat(it.score, trailing_width = 3), pad(engine.classes[it.class].name, 22), percent(it.pixels / map_area));
|
|
|
|
if out_folder write_overlay(engine, image, result, tprint("%/%.segments.png", out_folder, path_strip_extension(path_filename(path))));
|
|
if store && ends_with_nocase(path, ".mcap") {
|
|
segs, mask := segment_mcap_chunks(engine, result, whiches[n], image.width, image.height);
|
|
stored, store_error := mcap_store(path, front = .[segs], back = .[mask]);
|
|
if stored print(" stored: % bytes of segs, % bytes of mask\n", segs.data.count, mask.data.count);
|
|
else { print(" cannot store: %\n", store_error); failures += 1; }
|
|
}
|
|
if check_folder {
|
|
c := compare(engine, result, tprint("%/%.ref", check_folder, path_filename(path)), tprint("%/%.ref.txt", check_folder, path_filename(path)));
|
|
if !c.ok { print(" check: %\n", c.error); failures += 1; continue; }
|
|
print(" check: pixels agreeing on their class %; segments: % of % matched, worst IoU %%\n", percent(c.agree / cast(float64) c.pixels),
|
|
c.matched, c.checked, formatFloat(c.worst, trailing_width = 4), ifx c.extra then tprint(", % extra", c.extra) else "");
|
|
checked += c.checked; matched += c.matched; extra += c.extra; worst_iou = min(worst_iou, c.worst);
|
|
agree += c.agree; total_pixels += c.pixels; missed_confident += c.missed_confident;
|
|
}
|
|
}
|
|
for * images image_free(it);
|
|
reset_temporary_storage();
|
|
}
|
|
|
|
print("\n% images: decoding % ms each, segmenting % ms each (model and post-processing, batch %)\n", done,
|
|
formatFloat(decode_ms / max(done, 1), trailing_width = 0), formatFloat(model_ms / max(done, 1), trailing_width = 0),
|
|
segment_batch(engine));
|
|
per := max(engine.stats.images, 1) * 1000.0;
|
|
print(" of which: preparing % ms, the model % ms, making segments % ms, per image\n", formatFloat(engine.stats.prepare_us / per, trailing_width = 1),
|
|
formatFloat(engine.stats.model_us / per, trailing_width = 1), formatFloat(engine.stats.post_us / per, trailing_width = 1));
|
|
if check_folder && total_pixels {
|
|
print("check: pixels agreeing on their class %; segments: % of % matched, % extra, worst IoU %\n",
|
|
percent(agree / cast(float64) total_pixels), matched, checked, extra, formatFloat(worst_iou, trailing_width = 4));
|
|
// Pass: nearly every pixel the same class, and every segment PyTorch is fairly sure of found. A
|
|
// decoder or a rounding a step apart can tip a segment near a threshold either way.
|
|
print("check: % confident segments (score at least %) missed\n", missed_confident, CONFIDENT);
|
|
if cast(float64) agree < total_pixels * 0.995 || missed_confident failures += 1;
|
|
}
|
|
exit(xx ifx failures then 1 else 0);
|
|
}
|
|
|
|
CONFIDENT :: 0.6; // a segment PyTorch scores this or more must be matched
|
|
|
|
Comparison :: struct {
|
|
ok: bool;
|
|
error: string;
|
|
pixels, agree: s64; // pixels, and those given the same class (or none) by both
|
|
checked, matched, extra: s64; // PyTorch's segments, those with one of ours of the same class at IoU >= 0.5, and ours left over
|
|
missed_confident: s64; // PyTorch's segments scoring CONFIDENT or more that ours do not match
|
|
worst: float64 = 1; // the lowest IoU of a PyTorch segment with its best match
|
|
}
|
|
|
|
// Ours against PyTorch's: classes pixel by pixel, and each of its segments against our best of the same class.
|
|
compare :: (engine: *Segment_Engine, result: Segment_Result, ref_path: string, info_path: string) -> Comparison {
|
|
c: Comparison;
|
|
data, ok := read_entire_file(ref_path, log_errors = false,, temp);
|
|
if !ok { c.error = tprint("no reference % (tools/export_mask2former.py --reference)", ref_path); return c; }
|
|
if data.count < 8 || slice(data, 0, 4) != "M2FR" { c.error = tprint("% is not a reference", ref_path); return c; }
|
|
w := (cast(s32) data[4]) + ((cast(s32) data[5]) << 8);
|
|
h := (cast(s32) data[6]) + ((cast(s32) data[7]) << 8);
|
|
if w != result.width || h != result.height { c.error = tprint("the reference is % x %, ours % x %", w, h, result.width, result.height); return c; }
|
|
ref_map := cast([] u8) advance(data, 8);
|
|
if ref_map.count != cast(s64) w * h { c.error = "the reference is cut short"; return c; }
|
|
|
|
// Its segments' classes, by id.
|
|
info, info_ok := read_entire_file(info_path, log_errors = false,, temp);
|
|
if !info_ok { c.error = tprint("no %", info_path); return c; }
|
|
ref_class: [256] s32;
|
|
for * ref_class it.* = -1;
|
|
ref_score: [256] float64;
|
|
ref_ids: [..] u8;
|
|
ref_ids.allocator = temp;
|
|
for line: split(info, "\n",, temp) {
|
|
parts := split(trim(line), " ",, temp);
|
|
if parts.count < 2 continue;
|
|
id := string_to_int(parts[0]);
|
|
ref_class[id] = cast(s32) string_to_int(parts[1]);
|
|
if parts.count >= 3 ref_score[id] = string_to_float64(parts[2]);
|
|
array_add(*ref_ids, cast(u8) id);
|
|
}
|
|
our_class: [256] s32;
|
|
for * our_class it.* = -1;
|
|
for result.segments our_class[it.id] = it.class;
|
|
|
|
c.pixels = ref_map.count;
|
|
for 0..ref_map.count - 1 if ref_class[ref_map[it]] == our_class[result.map[it]] c.agree += 1;
|
|
|
|
used: [256] bool;
|
|
for rid: ref_ids {
|
|
c.checked += 1;
|
|
best: float64;
|
|
best_id: u8;
|
|
for result.segments {
|
|
if it.class != ref_class[rid] continue;
|
|
both, either: s64;
|
|
for p: 0..ref_map.count - 1 {
|
|
a := ref_map[p] == rid;
|
|
b := result.map[p] == it.id;
|
|
if a && b both += 1;
|
|
if a || b either += 1;
|
|
}
|
|
iou := ifx either then cast(float64) both / either else 1.0;
|
|
if iou > best { best, best_id = iou, it.id; }
|
|
}
|
|
if best >= 0.5 { c.matched += 1; used[best_id] = true; }
|
|
else if ref_score[rid] >= CONFIDENT c.missed_confident += 1;
|
|
if best < 0.5 print(" check: PyTorch's % (score %) has no match here (best IoU %)\n", engine.classes[ref_class[rid]].name, formatFloat(ref_score[rid], trailing_width = 3), formatFloat(best, trailing_width = 3));
|
|
c.worst = min(c.worst, best);
|
|
}
|
|
for result.segments if !used[it.id] c.extra += 1;
|
|
c.ok = true;
|
|
return c;
|
|
}
|
|
|
|
// The segments in their classes' colours over the image, at the map's size, as a PNG.
|
|
write_overlay :: (engine: *Segment_Engine, image: Image, result: Segment_Result, path: string) {
|
|
w, h := result.width, result.height;
|
|
pixels := NewArray(cast(s64) w * h * 4, u8,, temp);
|
|
for y: 0..h - 1 for x: 0..w - 1 {
|
|
sx := min(cast(s32) ((x + 0.5) * image.width / w), image.width - 1);
|
|
sy := min(cast(s32) ((y + 0.5) * image.height / h), image.height - 1);
|
|
source := image.pixels + (cast(s64) sy * image.width + sx) * 4;
|
|
dest := pixels.data + (cast(s64) y * w + x) * 4;
|
|
id := result.map[cast(s64) y * w + x];
|
|
if id {
|
|
colour := segment_colour(result.segments[id - 1].class);
|
|
for c: 0..2 dest[c] = cast(u8) ((cast(s32) source[c] + colour[c]) / 2);
|
|
} else {
|
|
for c: 0..2 dest[c] = source[c] / 3; // no segment: darkened
|
|
}
|
|
dest[3] = 255;
|
|
}
|
|
stbi_write_png(temp_c_string(path), w, h, 4, pixels.data, w * 4);
|
|
}
|
|
|
|
ms :: (since: Apollo_Time) -> float64 {
|
|
return to_float64_seconds(current_time_monotonic() - since) * 1000;
|
|
}
|
|
|
|
percent :: (fraction: float64) -> string {
|
|
return join(tprint("%", formatFloat(fraction * 100, trailing_width = 1)), "%",, temp);
|
|
}
|
|
|
|
pad :: (text: string, width: s64) -> string {
|
|
if text.count >= width return text;
|
|
builder: String_Builder;
|
|
builder.allocator = temp;
|
|
append(*builder, text);
|
|
for 1..width - text.count append(*builder, " ");
|
|
return builder_to_string(*builder,, temp);
|
|
}
|
|
|
|
#import "Basic";
|
|
#import "String";
|
|
#import "File";
|
|
#import "File_Utilities";
|
|
#import "System";
|
|
#import "Image_Decode";
|
|
#import "Segment";
|
|
#import "Capture_File";
|
|
#import "stb_image_write";
|