lamino-reco/lamino.c

#include <stdlib.h>
#include <ufo/ufo.h>

typedef struct {
    gchar *radios;
    gchar *darks;
    gchar *flats;
    gchar *output;
    gint width;
    gint height;
    guint num_radios;
    guint num_darks;
    gdouble theta;
    gdouble tau;
    gdouble psi;
    gdouble px;
    gdouble py;
    gdouble px_variation;
    gdouble v_origin[3];
    guint v_size[3];
} Params;

const int COLOR_RED = 31;
const int COLOR_GREEN = 32;
const int COLOR_YELLOW = 33;

static void
check (GError *error)
{
    if (error != NULL) {
        if (error->message != NULL)
            g_printerr ("%s", error->message);

        exit (1);
    }
}

static UfoTaskNode *
make_task (UfoPluginManager *pm, const gchar *name)
{
    GError *error = NULL;
    UfoTaskNode *task;

    task = ufo_plugin_manager_get_task (pm, name, &error);
    check (error);

    return task;
}

static guint
next_power_of_two (guint32 x)
{
    --x;
    x |= x >> 1;
    x |= x >> 2;
    x |= x >> 4;
    x |= x >> 8;
    x |= x >> 16;
    return x + 1;
}

static void
colored_print (gint color, const gchar *fmt, va_list args)
{
    gchar *s;

    s = g_strdup_vprintf (fmt, args);
    g_print ("%c[%d;%dm%s%c[%dm", 0x1B, 1, color, s, 0x1b, 0);
    g_free (s);
}

static void
warn (const gchar *fmt, ...)
{
    va_list args;
    va_start (args,  fmt);
    colored_print (COLOR_YELLOW, fmt, args);
    va_end (args);
}

static void
err (const gchar *fmt, ...)
{
    va_list args;
    va_start (args,  fmt);
    colored_print (COLOR_RED, fmt, args);
    va_end (args);
}

static void
info (const gchar *fmt, ...)
{
    va_list args;
    gchar *s;

    va_start (args,  fmt);
    g_print ("%c[%d;%dm>%c[%dm ", 0x1B, 1, COLOR_GREEN, 0x1b, 0);

    s = g_strdup_vprintf (fmt, args);
    g_print ("%s", s);
    va_end (args);
    g_free (s);
}

static void
print_dots (gpointer user)
{
    static int n = 0;

    if ((n++ % 10) == 0)
        g_print (".");
}

static void
run_reconstruction (Params *params)
{
    guint xl, xr, yt, yb;
    guint padded_width;
    guint padded_height;
    gdouble theta_rad;
    gdouble angle_step;
    guint fwidth;
    gdouble time = 0.0;
    UfoPluginManager *pm = NULL;
    UfoTaskGraph *graph = NULL;
    UfoBaseScheduler *sched = NULL;
    UfoTaskNode *radio_reader = NULL;
    UfoTaskNode *dark_reader = NULL;
    UfoTaskNode *flat_reader = NULL;
    UfoTaskNode *ffc = NULL;
    UfoTaskNode *pad = NULL;
    UfoTaskNode *ramp = NULL;
    UfoTaskNode *fft1 = NULL;
    UfoTaskNode *fft2 = NULL;
    UfoTaskNode *ifft = NULL;
    UfoTaskNode *conv = NULL;
    UfoTaskNode *reco = NULL;
    UfoTaskNode *writer = NULL;
    GError *error = NULL;

    pm = ufo_plugin_manager_new (NULL);
    graph = UFO_TASK_GRAPH (ufo_task_graph_new ());

    radio_reader = make_task (pm, "reader");
    pad = make_task (pm, "padding-2d");
    ramp = make_task (pm, "lamino-ramp");
    fft1 = make_task (pm, "fft");
    fft2 = make_task (pm, "fft");
    ifft = make_task (pm, "ifft");
    conv = make_task (pm, "lamino-conv");
    reco = make_task (pm, "lamino-bp");
    writer = make_task (pm, "writer");

    g_object_set (radio_reader,
                  "path", params->radios,
                  "end", params->num_radios - 1,
                  NULL);

    g_object_set (fft1, "dimensions", 2, NULL);
    g_object_set (fft2, "dimensions", 2, NULL);
    g_object_set (ifft, "dimensions", 2, NULL);

    fwidth = ((guint) params->px) * 2;
    padded_width = next_power_of_two (fwidth) * 2;
    padded_height = next_power_of_two ((guint32) params->height + 1);

    xl = params->px - params->width / 2;
    xr = padded_width - params->width - xl;
    yt = params->py - params->height / 2;
    yb = padded_height - params->height - yt;

    angle_step = (G_PI * 2.0) / params->num_radios;
    theta_rad = params->theta / 360. * G_PI * 2;

    info ("Axis: x=%.1f  y=%.1f  variation=%.1f\n",
          params->px, params->py, params->px_variation);

    info ("Lamino: theta=%.3f  tau=%.3f  step=%.5f  fwidth=%d\n",
          theta_rad, params->tau, angle_step, fwidth);

    info ("Padding: size=[%d %d]  (xl=%d xr=%d yt=%d yb=%d)\n",
          padded_width, padded_height, xl, xr, yt, yb);

    info ("Volume: origin=[%.1f %.1f %.1f]  size=[%d %d %d]\n",
          params->v_origin[0], params->v_origin[1], params->v_origin[2],
          params->v_size[0], params->v_size[1], params->v_size[2]);

    g_object_set (pad,
                  "xl", xl, "xr", xr,
                  "yt", yt, "yb", yb,
                  "mode", "brep",
                  NULL);

    g_object_set (ramp,
                  "width", padded_width,
                  "height", padded_height,
                  "theta", theta_rad,
                  "tau", params->tau,
                  "fwidth", fwidth,
                  NULL);

    g_object_set (reco,
                  "angle-step", angle_step,
                  "theta", theta_rad,
                  "psi", params->psi,
                  "proj-ox", params->px,
                  "proj-oy", params->py,
                  "proj-ox-variation", params->px_variation,
                  "vol-ox", params->v_size[0] / 2 + params->v_origin[0],
                  "vol-oy", params->v_size[1] / 2 + params->v_origin[1],
                  "vol-oz", params->v_size[2] / 2 + params->v_origin[2],
                  "vol-sx", params->v_size[0],
                  "vol-sy", params->v_size[1],
                  "vol-sz", params->v_size[2],
                  NULL);

    g_object_set (writer,
                  "filename", params->output,
                  NULL);

    if (params->darks != NULL && params->flats != NULL) {
        flat_reader = make_task (pm, "reader");
        dark_reader = make_task (pm, "reader");
        ffc = make_task (pm, "flat-field-correction");

        g_object_set (flat_reader, "path", params->flats, NULL);
        g_object_set (dark_reader, "path", params->darks, NULL);

        ufo_task_graph_connect_nodes_full (graph, radio_reader, ffc, 0);
        ufo_task_graph_connect_nodes_full (graph, dark_reader, ffc, 1);
        ufo_task_graph_connect_nodes_full (graph, flat_reader, ffc, 2);
        ufo_task_graph_connect_nodes (graph, ffc, pad);
    }
    else {
        warn ("> No flat/dark field correction\n");
        ufo_task_graph_connect_nodes (graph, radio_reader, pad);
    }

    ufo_task_graph_connect_nodes (graph, pad, fft1);
    ufo_task_graph_connect_nodes (graph, ramp, fft2);
    ufo_task_graph_connect_nodes_full (graph, fft1, conv, 0);
    ufo_task_graph_connect_nodes_full (graph, fft2, conv, 1);
    ufo_task_graph_connect_nodes (graph, conv, ifft);
    ufo_task_graph_connect_nodes (graph, ifft, reco);
    ufo_task_graph_connect_nodes (graph, reco, writer);

    g_signal_connect (reco, "processed", G_CALLBACK (print_dots), NULL);

    sched = UFO_BASE_SCHEDULER (ufo_scheduler_new (NULL, NULL));
    ufo_base_scheduler_run (sched, graph, &error);
    check (error);

    g_object_get (sched, "time", &time, NULL);
    g_print ("\n");
    info("Finished in %3.3fs\n", time);

    g_object_unref (pm);
    g_object_unref (graph);
    g_object_unref (sched);

    g_object_unref (radio_reader);
    g_object_unref (pad);
    g_object_unref (fft1);
    g_object_unref (fft2);
    /* g_object_unref (conv); */
    g_object_unref (reco);
    g_object_unref (writer);

    if (params->darks != NULL && params->flats != NULL) {
        g_object_unref (ffc);
        g_object_unref (flat_reader);
        g_object_unref (dark_reader);
    }
}

static void
parse_params (Params *params, int argc, char **argv)
{
    GOptionContext *context;
    GError *error = NULL;

    GOptionEntry entries[] = {
        { "width", 0, 0, G_OPTION_ARG_INT, &params->width, "Width", "" },
        { "height", 0, 0, G_OPTION_ARG_INT, &params->height, "Height", "" },
        { "num-radios", 0, 0, G_OPTION_ARG_INT, &params->num_radios, "Number of radios", "" },
        { "num-darks", 0, 0, G_OPTION_ARG_INT, &params->num_darks, "Number of darks", "" },
        { "radios", 0, 0, G_OPTION_ARG_STRING, &params->radios, "Radios", "" },
        { "darks", 0, 0, G_OPTION_ARG_STRING, &params->darks, "Darks", "" },
        { "flats", 0, 0, G_OPTION_ARG_STRING, &params->flats, "Flats", "" },
        { "output", 0, 0, G_OPTION_ARG_STRING, &params->output, "Output", "" },
        { "theta", 0, 0, G_OPTION_ARG_DOUBLE, &params->theta, "Tilt (theta)", "" },
        { "tau", 0, 0, G_OPTION_ARG_DOUBLE, &params->tau, "Pixel size (theta)", "" },
        { "psi", 0, 0, G_OPTION_ARG_DOUBLE, &params->psi, "Misalignment (psi)", "" },
        { "px", 0, 0, G_OPTION_ARG_DOUBLE, &params->px, "X coordinate of axis", "" },
        { "py", 0, 0, G_OPTION_ARG_DOUBLE, &params->py, "Y coordinate of axis", "" },
        { "px-variation", 0, 0, G_OPTION_ARG_DOUBLE, &params->px_variation, "X axis variation", "" },
        { "vx", 0, 0, G_OPTION_ARG_DOUBLE, &params->v_origin[0], "X coordinate of box origin", "" },
        { "vy", 0, 0, G_OPTION_ARG_DOUBLE, &params->v_origin[1], "Y coordinate of box origin", "" },
        { "vz", 0, 0, G_OPTION_ARG_DOUBLE, &params->v_origin[2], "Z coordinate of box origin", "" },
        { "vw", 0, 0, G_OPTION_ARG_INT, &params->v_size[0], "Width of box", "" },
        { "vh", 0, 0, G_OPTION_ARG_INT, &params->v_size[1], "Height of box", "" },
        { "vd", 0, 0, G_OPTION_ARG_INT, &params->v_size[2], "Depth of box", "" },
        { NULL }
    };

    context = g_option_context_new ("- test tree model performance");
    g_option_context_add_main_entries (context, entries, NULL);

    if (!g_option_context_parse (context, &argc, &argv, &error)) {
        g_print ("option parsing failed: %s\n", error->message);
        exit (1);
    }

    if (params->width == 0 || params->height == 0 || params->num_radios == 0 ||
        params->radios == NULL ||
        params->theta == G_MAXDOUBLE ||
        params->px == G_MAXDOUBLE || params->py == G_MAXDOUBLE) {
            err ("Parameters missing.\n\n");
            g_print ("%s\n", g_option_context_get_help (context, TRUE, NULL));
            exit (1);
    }

    g_option_context_free (context);
}

int
main (int argc, char **argv)
{
    Params params = {
        .radios = NULL,
        .darks = NULL,
        .flats = NULL,
        .output = "volume-%i.tif",
        .width = 0,
        .height = 0,
        .num_radios = 0,
        .num_darks = 1,
        .theta = G_MAXDOUBLE,
        .tau = 1.0,
        .psi = 0.0,
        .px = G_MAXDOUBLE,
        .py = G_MAXDOUBLE,
        .px_variation = 0.0,
        .v_size = {256, 256, 256},
        .v_origin = {0.0, 0.0, 0.0},
    };

#if !(GLIB_CHECK_VERSION (2, 36, 0))
    g_type_init ();
#endif

    parse_params (&params, argc, argv);
    run_reconstruction (&params);

    return 0;
}