// Segment: what is where in an image, from a panoptic segmentation model (Mask2Former): every pixel
// given a class ("person", "dog", "sky", "wall"), and countable things split into instances.
//
//   engine, ok, error := segment_create(models_folder, "mask2former-swin-base-coco-panoptic");
//   results := segment_images(engine, pictures);   // one Segment_Result per picture
//   for results[0].segments  print("% %\n", engine.classes[it.class].name, it.score);
//
// The model is listed in models_folder/models.config, made by tools/export_mask2former.py, and run
// with onnxruntime (modules/onnxruntime): on an NVIDIA GPU through CUDA when it can, else on the CPU.
// Its two outputs, a class and a mask for each of its 100 queries, are made into segments here, as
// transformers' post_process_panoptic_segmentation does (tools/export_mask2former.py checks that
// they match).

Segment_Provider :: enum { CPU; CUDA; }

// Tightly packed or not: stride is the bytes from one row to the next.
Segment_Image :: struct {
  pixels: *u8;   // RGBA
  width, height, stride: s32;
}

Segment_Model_Info :: struct {
  name: string;
  title: string;
  model: string;                       // the .onnx, in the models folder
  classes: string;                     // the class list, in the models folder
  image_size: s32 = 384;               // the input's width and height: images are resized to it, not keeping their shape
  mean := float.[0.485, 0.456, 0.406]; // ImageNet's, per channel, after scaling to 0..1
  std  := float.[0.229, 0.224, 0.225];
  batch_gpu: s32 = 4;
  batch_cpu: s32 = 2;
  score_threshold: float = 0.5;        // a query's class score must be above this (transformers' default)
  mask_threshold: float = 0.5;         // a pixel is in a query's own mask at this (score times mask)
  overlap_threshold: float = 0.8;      // a segment keeps more than this of its own mask, or is dropped
  map_max_side: s32 = 1024;            // the segment map: the image scaled so its longest side is at most this
}

Segment_Class :: struct {
  name: string;      // as people see and search it: "wall"
  model_name: string;// the model's own: "wall-other-merged"
  thing: bool;       // countable (person, dog), not stuff (sky, wall)
}

Segment_Engine :: struct {
  model: Segment_Model_Info;
  classes: [] Segment_Class;
  classes_hash: u64;                   // of the class list, so results from another list are out of date
  provider: Segment_Provider;
  gpu_explanation: string;             // why the GPU is not in use; "" when it is
  env: *OrtEnv;
  memory_info: *OrtMemoryInfo;
  session: Ort_Session;
  cancel: *bool;                       // when set, and true, segment_images stops between batches: the rest fail as "cancelled"
  stats: Segment_Stats;                // since it was made
}

// Time spent in each step of segment_images, summed.
Segment_Stats :: struct {
  images: s64;
  prepare_us: s64;   // resizing and normalizing
  model_us: s64;     // the model's run, on the GPU or the CPU
  post_us: s64;      // making segments of its outputs
}

// One segment: a thing, or a stretch of stuff. Its pixels are those of map with its id.
Segment :: struct {
  id: u8;                // 1 up, in the order made
  class: s32;            // into engine.classes
  score: float;          // the model's confidence in its class, 0 to 1
  pixels: s64;           // its area in the map
  x0, y0, x1, y1: s32;   // its bounds in the map; x1 and y1 are past its last pixel
}

Segment_Result :: struct {
  failed: bool;
  error: string;
  width, height: s32;    // the map's: the image's shape, longest side at most map_max_side
  map: [] u8;            // a segment id per pixel, row by row; 0: none
  segments: [] Segment;
}

SEGMENT_MODELS_CONFIG :: "models.config";
SEGMENT_CANCELLED     :: "cancelled";

// The models models.config lists, in its order; from the context allocator.
segment_list_models :: (models_folder: string) -> [] Segment_Model_Info, error: string {
  models: [..] Segment_Model_Info;
  path := tprint("%/%", models_folder, SEGMENT_MODELS_CONFIG);
  text, ok := read_entire_file(path, log_errors = false,, temp);
  if !ok  return models, tprint("Cannot read %.", path);

  model: *Segment_Model_Info;
  version_seen := false;
  line_number := 0;
  for raw: split(text, "\n",, temp) {
    line_number += 1;
    line := trim(segment_strip_comment(raw));
    if !line.count  continue;
    if begins_with(line, "@version") { version_seen = true; continue; }
    if !version_seen  return models, tprint("% does not start with @version 1.", path);
    if begins_with(line, "[model ") && ends_with(line, "]") {
      model = array_add(*models);
      model.* = .{};
      model.name = copy_string(trim(slice(line, 7, line.count - 8)));
      continue;
    }
    colon := find_index_from_left(line, #char ":");
    if colon < 0 || !model { log_error("[Segment] %:%: expected `key: value` under a [model NAME]. Ignored.", path, line_number); continue; }
    key   := trim(slice(line, 0, colon));
    value := trim(advance(line, colon + 1));
    if value.count >= 2 && value[0] == #char "\"" && value[value.count - 1] == #char "\""  value = slice(value, 1, value.count - 2);
    number, number_ok := string_to_float(value);
    if key == {
      case "title";             model.title   = copy_string(value);
      case "model";             model.model   = copy_string(value);
      case "classes";           model.classes = copy_string(value);
      case "image_size";        if number_ok  model.image_size = cast(s32) number;
      case "mean";              segment_parse_three(value, *model.mean, path, line_number);
      case "std";               segment_parse_three(value, *model.std, path, line_number);
      case "batch_gpu";         if number_ok  model.batch_gpu = cast(s32) number;
      case "batch_cpu";         if number_ok  model.batch_cpu = cast(s32) number;
      case "score_threshold";   if number_ok  model.score_threshold = number;
      case "mask_threshold";    if number_ok  model.mask_threshold = number;
      case "overlap_threshold"; if number_ok  model.overlap_threshold = number;
      case "map_max_side";      if number_ok  model.map_max_side = cast(s32) number;
      case;  log_error("[Segment] %:%: unknown key %. Ignored.", path, line_number, key);
    }
  }
  return models, "";
}

// An engine for the model named in models_folder/models.config: on CUDA with prefer_gpu when every
// library that needs is present, else on the CPU.
segment_create :: (models_folder: string, model_name: string, prefer_gpu := true) -> *Segment_Engine, ok: bool, error: string {
  push_allocator(context.default_allocator);
  loaded, load_error := ort_load();
  if !loaded  return null, false, load_error;

  models, models_error := segment_list_models(models_folder);
  if models_error  return null, false, models_error;
  model: *Segment_Model_Info;
  for * models  if it.name == model_name { model = it; break; }
  if !model  return null, false, tprint("% has no [model %].", SEGMENT_MODELS_CONFIG, model_name);

  engine := New(Segment_Engine);
  engine.model = model.*;
  ok, error := segment_load_classes(engine, tprint("%/%", models_folder, model.classes));
  if !ok { segment_destroy(engine); return null, false, error; }

  api := ort.api;
  ok, error = ort_check(api.CreateEnv(.ORT_LOGGING_LEVEL_WARNING, "segment", *engine.env));
  if !ok { segment_destroy(engine); return null, false, error; }
  ort_check(api.CreateCpuMemoryInfo(.OrtArenaAllocator, .OrtMemTypeDefault, *engine.memory_info));

  HINT :: "make it with tools/export_mask2former.py (see modules/Segment/README.md) and rebuild";
  path := tprint("%/%", models_folder, model.model);
  use_gpu := false;
  if prefer_gpu {
    available, explanation := ort_cuda_check("Segmentation");
    use_gpu = available;
    if !available  engine.gpu_explanation = copy_string(explanation);
  } else {
    engine.gpu_explanation = copy_string("Segmentation runs on the CPU, as it was asked to.");
  }
  if use_gpu {
    ort_gpu_lock();   // one CUDA engine made or released at a time (modules/onnxruntime/session.jai)
    engine.session, ok, error = ort_create_session(engine.env, path, true, 2, missing_hint = HINT);
    if ok  ok, error = segment_warm_up(engine);
    if !ok  ort_release_session(*engine.session);
    ort_gpu_unlock();
    if ok {
      engine.provider = .CUDA;
    } else {
      engine.gpu_explanation = sprint("Segmentation runs on the CPU: the CUDA libraries were found, but onnxruntime could not start on the GPU:\n%", error);
      use_gpu = false;
    }
  }
  if !use_gpu {
    cores := cast(s32) max(get_number_of_processors(.ALL_LOGICAL), 1);
    engine.session, ok, error = ort_create_session(engine.env, path, false, cores, missing_hint = HINT);
    if !ok { segment_destroy(engine); return null, false, error; }
    engine.provider = .CPU;
  }
  return engine, true, "";
}

segment_destroy :: (engine: *Segment_Engine) {
  if !engine  return;
  push_allocator(context.default_allocator);
  if engine.provider == .CUDA  ort_gpu_lock();
  ort_release_session(*engine.session);
  if engine.provider == .CUDA  ort_gpu_unlock();
  if engine.memory_info  ort.api.ReleaseMemoryInfo(engine.memory_info);
  if engine.env          ort.api.ReleaseEnv(engine.env);
  for engine.classes { free(it.name); free(it.model_name); }
  array_free(engine.classes);
  free(engine.gpu_explanation);
  free(engine);
}

// The map size for an image of width x height: its shape, the longest side at most max_side.
segment_map_size :: (width: s32, height: s32, max_side: s32) -> s32, s32 {
  scale := min(cast(float64) 1, cast(float64) max_side / max(width, height));
  return max(cast(s32) 1, cast(s32) (width * scale + 0.5)), max(cast(s32) 1, cast(s32) (height * scale + 0.5));
}

// A segmentation of each image, in batches; from allocator. An image that fails says why in its
// result; with engine.cancel set and true between batches, every image left fails as SEGMENT_CANCELLED.
segment_images :: (engine: *Segment_Engine, images: [] Segment_Image, allocator := context.allocator) -> [] Segment_Result {
  results := NewArray(images.count, Segment_Result,, allocator);
  size := engine.model.image_size;
  plane := cast(s64) size * size;
  batch := cast(s64) max(ifx engine.provider == .CUDA then engine.model.batch_gpu else engine.model.batch_cpu, 1);
  input := NewArray(plane * 3 * batch, float,, context.default_allocator);
  defer array_free(input,, context.default_allocator);

  first := 0;
  while first < images.count {
    count := min(batch, images.count - first);
    if engine.cancel && engine.cancel.* {
      for first..images.count - 1  results[it] = .{failed = true, error = copy_string(SEGMENT_CANCELLED,, allocator)};
      break;
    }
    start := current_time_monotonic();
    for 0..count - 1  segment_prepare_image(engine.model, images[first + it], input.data + plane * 3 * it);
    shape := s64.[count, 3, size, size];
    outs: [2] Ort_Output;
    ran := current_time_monotonic();
    engine.stats.prepare_us += to_microseconds(ran - start);
    ok, error := ort_run_float_outputs(*engine.session, engine.memory_info, array_view(input, 0, plane * 3 * count), shape, .["class_queries_logits", "masks_queries_logits"], outs);
    posted := current_time_monotonic();
    engine.stats.model_us += to_microseconds(posted - ran);
    engine.stats.images += count;
    defer engine.stats.post_us += to_microseconds(current_time_monotonic() - posted);
    if !ok {
      for first..first + count - 1  results[it] = .{failed = true, error = copy_string(error,, allocator)};
    } else {
      classes, masks := outs[0], outs[1];
      queries := classes.dims[1];
      labels  := classes.dims[2];   // the classes and "no object", last
      mask_h, mask_w := masks.dims[2], masks.dims[3];
      if classes.dim_count != 3 || masks.dim_count != 4 || classes.dims[0] != count || masks.dims[0] != count || masks.dims[1] != queries || labels != engine.classes.count + 1 {
        message := tprint("The model gave outputs of shape % and %, not [%, Q, %] and [%, Q, H, W].", array_view(classes.dims, 0, classes.dim_count), array_view(masks.dims, 0, masks.dim_count), count, engine.classes.count + 1, count);
        for first..first + count - 1  results[it] = .{failed = true, error = copy_string(message,, allocator)};
      } else {
        for n: 0..count - 1 {
          image := images[first + n];
          results[first + n] = segment_panoptic(engine, classes.data + n * queries * labels, masks.data + n * queries * mask_h * mask_w,
                                                cast(s32) queries, cast(s32) labels, cast(s32) mask_w, cast(s32) mask_h, image.width, image.height, allocator);
        }
      }
      ort_release_output(*outs[0]);
      ort_release_output(*outs[1]);
    }
    first += count;
  }
  return results;
}

segment_result_free :: (result: *Segment_Result, allocator := context.allocator) {
  free(result.error,, allocator);
  array_free(result.map,, allocator);
  array_free(result.segments,, allocator);
  result.* = .{};
}

// A result as the chunks an .mcap stores, keyed by the model's name (modules/Capture_File/segment_chunks.jai):
// segs, with each segment's box in the pixels of the image read (image_width x image_height; which
// is "rendered" or "original"), largest first; and the mask. Temp.
segment_mcap_chunks :: (engine: *Segment_Engine, result: Segment_Result, which: string, image_width: s32, image_height: s32) -> segs: Mcap_Chunk, mask: Mcap_Chunk {
  s: Mcap_Segments;
  s.engine = tprint("onnxruntime %, %", ort.version, ifx engine.provider == .CUDA then "CUDA" else "CPU");
  s.date = mcap_ocr_now();
  s.image = which;
  s.image_width, s.image_height = image_width, image_height;
  s.classes = segment_classes_text(engine);
  s.segments.allocator = temp;
  sx := cast(float64) image_width / result.width;
  sy := cast(float64) image_height / result.height;
  map_area := cast(float64) result.width * result.height;
  for result.segments {
    x := cast(s32) (it.x0 * sx);
    y := cast(s32) (it.y0 * sy);
    array_add(*s.segments, .{id = it.id, score = it.score, area = cast(float) (it.pixels / map_area), x = x, y = y,
                             width = cast(s32) (it.x1 * sx + 0.5) - x, height = cast(s32) (it.y1 * sy + 0.5) - y, class = engine.classes[it.class].name});
  }
  quick_sort(s.segments, (a: Mcap_Segment, b: Mcap_Segment) -> s64 { return ifx a.area > b.area then -1 else ifx a.area < b.area then 1 else 0; });
  name := engine.model.name;
  return mcap_segments_chunk(name, s), mcap_mask_chunk(name, result.width, result.height, result.map);
}

// The class list's hash as segs stores it: 16 hex digits. Temp.
segment_classes_text :: (engine: *Segment_Engine) -> string {
  return tprint("%", formatInt(engine.classes_hash, base = 16, minimum_digits = 16));
}

// A class's colour, the same everywhere: golden-angle hues, so neighbouring class numbers differ.
segment_colour :: (class: s32) -> [3] u8 {
  degrees := cast(float64) class * 137.508;
  hue := cast(float) (degrees - floor(degrees / 360) * 360) / 60.0;
  sector := cast(s32) hue;
  f := hue - sector;
  v, p, q, t := 230.0, 60.0, 230.0 - 170.0 * f, 60.0 + 170.0 * f;
  rgb: [3] float;
  if sector == {
    case 0; rgb = .[v, t, p];
    case 1; rgb = .[q, v, p];
    case 2; rgb = .[p, v, t];
    case 3; rgb = .[p, q, v];
    case 4; rgb = .[t, p, v];
    case;   rgb = .[v, p, q];
  }
  return .[cast(u8) rgb[0], cast(u8) rgb[1], cast(u8) rgb[2]];
}

// The model's input for an image (3 x image_size x image_size floats at dest), as segment_images makes
// it: for checking against the checkpoint's own processor (segment-probe -dump-input).
segment_prepare :: (engine: *Segment_Engine, image: Segment_Image, dest: *float) {
  segment_prepare_image(engine.model, image, dest);
}

#scope_module

// Transformers' post_process_panoptic_segmentation, for one image (no stuff fused): queries whose
// best class is not "no object" and scores above score_threshold are kept; their masks are scaled
// up to image_size, then to the map, through their sigmoid, and weighted by their scores; each
// pixel goes to the query weighing most there; and a query keeps its segment only when that is more
// than overlap_threshold of its own mask (its pixels weighing at least mask_threshold).
segment_panoptic :: (engine: *Segment_Engine, class_logits: *float, mask_logits: *float, queries: s32, labels: s32, mask_w: s32, mask_h: s32,
                     image_w: s32, image_h: s32, allocator: Allocator) -> Segment_Result {
  model := engine.model;
  result: Segment_Result;
  result.width, result.height = segment_map_size(image_w, image_h, model.map_max_side);
  map_w, map_h := result.width, result.height;
  result.map = NewArray(cast(s64) map_w * map_h, u8, initialized = true,, allocator);

  // The queries kept: softmax over the classes, then the best.
  Kept :: struct { query: s32; class: s32; score: float; }
  kept: [..] Kept;
  kept.allocator = temp;
  for q: 0..queries - 1 {
    row := class_logits + cast(s64) q * labels;
    top := row[0];
    for 1..labels - 1  top = max(top, row[it]);
    total: float64;
    for 0..labels - 1  total += exp(cast(float64) (row[it] - top));
    best: s64;
    best_score: float64 = -1;
    for 0..labels - 1 {
      p := exp(cast(float64) (row[it] - top)) / total;
      if p > best_score { best, best_score = it, p; }   // the first of equals, as torch's max
    }
    if best == labels - 1 || best_score <= model.score_threshold  continue;
    array_add(*kept, .{q, cast(s32) best, cast(float) best_score});
  }
  if !kept.count || kept.count > 255  return result;

  // Each kept query's mask at image_size: its logits scaled up (bilinear, corners not aligned), then
  // through the sigmoid.
  size := model.image_size;
  plane := cast(s64) size * size;
  probs := NewArray(kept.count * plane, float, initialized = false,, temp);
  across_in := segment_linear_taps(mask_w, size,, temp);
  down_in   := segment_linear_taps(mask_h, size,, temp);
  for k, n: kept {
    source := mask_logits + cast(s64) k.query * mask_w * mask_h;
    dest := probs.data + n * plane;
    for y: 0..size - 1 {
      ty := down_in[y];
      row0 := source + cast(s64) ty.i0 * mask_w;
      row1 := source + cast(s64) ty.i1 * mask_w;
      for x: 0..size - 1 {
        tx := across_in[x];
        top    := row0[tx.i0] * (1 - tx.t) + row0[tx.i1] * tx.t;
        bottom := row1[tx.i0] * (1 - tx.t) + row1[tx.i1] * tx.t;
        logit := top * (1 - ty.t) + bottom * ty.t;
        dest[cast(s64) y * size + x] = 1.0 / (1.0 + exp(-logit));
      }
    }
  }

  // Then at the map's size, weighted by the scores: each pixel to the query weighing most, and each
  // query's own mask counted.
  across := segment_linear_taps(size, map_w,, temp);
  down   := segment_linear_taps(size, map_h,, temp);
  best_of := NewArray(cast(s64) map_w * map_h, u8, initialized = false,, temp);   // index into kept
  won  := NewArray(kept.count, s64, initialized = true,, temp);
  own  := NewArray(kept.count, s64, initialized = true,, temp);
  for y: 0..map_h - 1 {
    ty := down[y];
    for x: 0..map_w - 1 {
      tx := across[x];
      best, best_value := 0, -1.0;
      for k, n: kept {
        p := probs.data + n * plane;
        row0 := p + cast(s64) ty.i0 * size;
        row1 := p + cast(s64) ty.i1 * size;
        top    := row0[tx.i0] * (1 - tx.t) + row0[tx.i1] * tx.t;
        bottom := row1[tx.i0] * (1 - tx.t) + row1[tx.i1] * tx.t;
        value := (top * (1 - ty.t) + bottom * ty.t) * k.score;
        if value >= model.mask_threshold  own[n] += 1;
        if value > best_value { best, best_value = n, value; }   // the first of equals, as torch's argmax
      }
      best_of[cast(s64) y * map_w + x] = cast(u8) best;
      won[best] += 1;
    }
  }

  // The segments: the queries, in order, that keep enough of their own mask.
  id_of := NewArray(kept.count, u8, initialized = true,, temp);   // 0: dropped
  segments: [..] Segment;
  segments.allocator = allocator;
  for k, n: kept {
    if !won[n] || !own[n] || !(cast(float64) won[n] / own[n] > model.overlap_threshold)  continue;
    id := cast(u8) (segments.count + 1);
    id_of[n] = id;
    array_add(*segments, .{id = id, class = k.class, score = k.score, x0 = map_w, y0 = map_h});
  }
  for y: 0..map_h - 1  for x: 0..map_w - 1 {
    at := cast(s64) y * map_w + x;
    id := id_of[best_of[at]];
    result.map[at] = id;
    if !id  continue;
    s := *segments[id - 1];
    s.pixels += 1;
    s.x0 = min(s.x0, cast(s32) x);  s.x1 = max(s.x1, cast(s32) x + 1);
    s.y0 = min(s.y0, cast(s32) y);  s.y1 = max(s.y1, cast(s32) y + 1);
  }
  result.segments = segments;
  return result;
}

// For each of out positions, the two in positions and the weight of the second, for torch's bilinear
// interpolate with align_corners = False: no averaging when shrinking. Temp.
Linear_Tap :: struct { i0, i1: s32; t: float; }

segment_linear_taps :: (in: s32, out: s32) -> [] Linear_Tap {
  taps := NewArray(out, Linear_Tap);
  scale := cast(float) in / out;
  for i: 0..out - 1 {
    source := max((i + 0.5) * scale - 0.5, 0.0);
    i0 := min(cast(s32) source, in - 1);
    i1 := min(i0 + 1, in - 1);
    taps[i] = .{i0, i1, source - i0};
  }
  return taps;
}

// The image resized to image_size x image_size (not keeping its shape, as the checkpoint's processor
// does) with PIL's bilinear filter, each pass rounded to 8 bits as PIL's is; then as planar RGB,
// scaled to 0..1 and normalized by the model's mean and std, at dest.
segment_prepare_image :: (model: Segment_Model_Info, image: Segment_Image, dest: *float) {
  size := model.image_size;
  across := segment_resample_weights(image.width, size);
  down   := segment_resample_weights(image.height, size);
  rows := NewArray(cast(s64) image.height * size * 3, u8, initialized = false,, temp);
  for y: 0..image.height - 1 {
    row := image.pixels + cast(s64) y * image.stride;
    for x: 0..size - 1 {
      w := *across[x];
      r, g, b: float;
      for k: 0..w.weights.count - 1 {
        p := row + (w.first + k) * 4;
        r += p[0] * w.weights[k];  g += p[1] * w.weights[k];  b += p[2] * w.weights[k];
      }
      at := (cast(s64) y * size + x) * 3;
      rows[at] = segment_byte(r);  rows[at + 1] = segment_byte(g);  rows[at + 2] = segment_byte(b);
    }
  }
  plane := cast(s64) size * size;
  scale: [3] float;
  offset: [3] float;
  for c: 0..2 { scale[c] = 1 / (255 * model.std[c]); offset[c] = model.mean[c] / model.std[c]; }
  for y: 0..size - 1 {
    w := *down[y];
    for x: 0..size - 1 {
      r, g, b: float;
      for k: 0..w.weights.count - 1 {
        at := (cast(s64) (w.first + k) * size + x) * 3;
        r += rows[at] * w.weights[k];  g += rows[at + 1] * w.weights[k];  b += rows[at + 2] * w.weights[k];
      }
      i := cast(s64) y * size + x;
      dest[i]             = segment_byte(r) * scale[0] - offset[0];
      dest[i + plane]     = segment_byte(g) * scale[1] - offset[1];
      dest[i + plane * 2] = segment_byte(b) * scale[2] - offset[2];
    }
  }
}

segment_byte :: (v: float) -> u8 {
  return cast(u8) clamp(cast(s32) (v + 0.5), 0, 255);
}

Resample_Weights :: struct { first: s32; weights: [] float; }

// For each of out positions, the input positions it takes and their weights: a triangle filter
// widened by the scale when shrinking, as PIL's bilinear resize (and modules/Labels). Temp.
segment_resample_weights :: (in: s32, out: s32) -> [] Resample_Weights {
  result := NewArray(out, Resample_Weights,, temp);
  scale := cast(float) in / out;
  filter_scale := max(scale, 1.0);
  support := filter_scale;
  for i: 0..out - 1 {
    centre := (i + 0.5) * scale;
    first := max(cast(s32) (centre - support + 0.5), 0);
    last  := min(cast(s32) (centre + support + 0.5), in);
    weights := NewArray(max(last - first, 1), float,, temp);
    total: float;
    for k: 0..last - first - 1 {
      x := (first + k - centre + 0.5) / filter_scale;
      w := max(0.0, 1 - abs(x));
      weights[k] = w;
      total += w;
    }
    if total > 0  for * weights  it.* /= total;
    else { first = clamp(cast(s32) centre, 0, in - 1); weights[0] = 1; }
    result[i] = .{first, weights};
  }
  return result;
}

// Runs the model once, on a blank image, when it is made on the GPU: see labels_warm_up.
segment_warm_up :: (engine: *Segment_Engine) -> ok: bool, error: string {
  size := engine.model.image_size;
  input := NewArray(3 * size * size, float, initialized = true,, context.default_allocator);
  defer array_free(input,, context.default_allocator);
  outs: [2] Ort_Output;
  ok, error := ort_run_float_outputs(*engine.session, engine.memory_info, input, s64.[1, 3, size, size], .["class_queries_logits", "masks_queries_logits"], outs);
  if !ok  return false, tprint("The first run on the GPU failed: %", error);
  ort_release_output(*outs[0]);
  ort_release_output(*outs[1]);
  return true, "";
}

// The class list: `class: INDEX thing|stuff "model's name" "name"`, a line each, in index order.
segment_load_classes :: (engine: *Segment_Engine, path: string) -> ok: bool, error: string {
  text, ok := read_entire_file(path, log_errors = false,, temp);
  if !ok  return false, tprint("Cannot read %; make it with tools/export_mask2former.py.", path);
  classes: [..] Segment_Class;
  for raw, line_index: split(text, "\n",, temp) {
    line := trim(raw);
    if !begins_with(line, "class:")  continue;
    rest := trim(advance(line, 6));
    space := find_index_from_left(rest, #char " ");
    index, index_ok := string_to_int(slice(rest, 0, max(space, 0)));
    if space < 0 || !index_ok || index != classes.count  return false, tprint("%:%: expected class % next.", path, line_index + 1, classes.count);
    rest = trim(advance(rest, space + 1));
    thing := begins_with(rest, "thing");
    if !thing && !begins_with(rest, "stuff")  return false, tprint("%:%: expected thing or stuff.", path, line_index + 1);
    quoted: [2] string;
    for 0..1 {
      open := find_index_from_left(rest, #char "\"");
      if open < 0  return false, tprint("%:%: expected two quoted names.", path, line_index + 1);
      close := find_index_from_left(rest, #char "\"", open + 1);
      if close < 0  return false, tprint("%:%: a name is not closed.", path, line_index + 1);
      quoted[it] = slice(rest, open + 1, close - open - 1);
      rest = advance(rest, close + 1);
    }
    array_add(*classes, .{name = copy_string(quoted[1]), model_name = copy_string(quoted[0]), thing = thing});
  }
  if !classes.count  return false, tprint("% lists no classes.", path);
  engine.classes = classes;

  // FNV-1a over each class's index, kind and name, as a line.
  hash := FNV_64_OFFSET_BIAS;
  for engine.classes {
    line := tprint("%\t%\t%\n", it_index, ifx it.thing then "thing" else "stuff", it.model_name);
    hash = fnv1a_hash(line.data, line.count, hash);
  }
  engine.classes_hash = hash;
  return true, "";
}

segment_parse_three :: (value: string, dest: *[3] float, path: string, line_number: s64) {
  parts := split(value, " ",, temp);
  got := 0;
  for parts {
    if !it.count  continue;
    number, ok := string_to_float(it);
    if !ok || got >= 3 { got = -1; break; }
    dest.*[got] = number;
    got += 1;
  }
  if got != 3  log_error("[Segment] %:%: expected three numbers.", path, line_number);
}

segment_strip_comment :: (line: string) -> string {
  in_quotes := false;
  for i: 0..line.count - 1 {
    c := line[i];
    if c == #char "\""  in_quotes = !in_quotes;
    else if c == #char "#" && !in_quotes && (i == 0 || line[i - 1] == #char " " || line[i - 1] == #char "\t")  return slice(line, 0, i);
  }
  return line;
}

#import "Basic";
#import "String";
#import "File";
#import "Math";
#import "System";
#import "Hash";
#import "Sort";
#import "Capture_File";   // segment_mcap_chunks
#import "onnxruntime";