lamino-reco/lamino.c

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


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

typedef void (*CommandFunction)(Params *, gchar **argv);

typedef struct {
    const gchar *command;
    CommandFunction func;
} CommandMap;


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 gboolean
params_okay (Params *params)
{
    return 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;
}

static void
run_reconstruction (Params *params, gchar **argv)
{
    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;

    if (!params_okay (params)) {
        err ("Parameters missing.\n");
        return;
    }

    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 (ffc, "dark-scale", params->dark_scale, NULL);
        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
print (Params *params, gchar **argv)
{
    for (gint i = 0; params->entries[i].long_name != NULL; i++) {
        GOptionEntry *entry;

        entry = &params->entries[i];

        if (argv[0] == NULL || g_strcmp0 (entry->long_name, argv[0]) == 0) {
            switch (entry->arg) {
                case G_OPTION_ARG_STRING:
                    info ("%-15s%s\n", entry->long_name, *((gchar **) entry->arg_data));
                    break;
                case G_OPTION_ARG_INT:
                    info ("%-15s%i\n", entry->long_name, *((gint *) entry->arg_data));
                    break;
                case G_OPTION_ARG_DOUBLE:
                    info ("%-15s%f\n", entry->long_name, *((gdouble *) entry->arg_data));
                    break;
                default:
                    break;
            }
        }
    }
}

static void
set (Params *params, gchar **argv)
{
    if (argv[0] == NULL) {
        warn ("No parameter given\n");
        return;
    }

    if (argv[1] == NULL) {
        warn ("No value given\n");
        return;
    }

    for (gint i = 0; params->entries[i].long_name != NULL; i++) {
        GOptionEntry *entry;

        entry = &params->entries[i];

        if (g_strcmp0 (entry->long_name, argv[0]) == 0) {
            switch (entry->arg) {
                case G_OPTION_ARG_STRING:
                    {
                        gchar *s;

                        s = *((gchar **) entry->arg_data);
                        g_free (s);
                        *(gchar **) entry->arg_data = g_strjoinv (" ", &argv[1]);
                    }
                    break;
                case G_OPTION_ARG_INT:
                    *(gint *) entry->arg_data = atoi (argv[1]);
                    break;
                case G_OPTION_ARG_DOUBLE:
                    *(gdouble *) entry->arg_data = atof (argv[1]);
                    break;
                default:
                    break;
            }

            /* print (params, arg); */
            break;
        }
    }
}

static void
quit (Params *params, gchar **argv)
{
    exit (0);
}

static void
start_shell (Params *params)
{
    static CommandMap map[] = {
        { "run", run_reconstruction },
        { "print", print },
        { "set", set },
        { "quit", quit },
        { NULL },
    };

    while (1) {
        gchar *line;
        gint i;
        gchar **split = NULL;

        line = readline ("> ");

        if (line && *line && (g_strchomp (line) == g_strchug (line))) {
            add_history (line);
        }
        else {
            goto cleanup_line;
        }

        split = g_strsplit (line, " ", 0);

        for (i = 0; map[i].command != NULL; i++) {
            /* Run command even if only partially written out */
            if (g_str_has_prefix (map[i].command, split[0])) {
                map[i].func (params, &split[1]);
                break;
            }
        }

        if (map[i].command == NULL)
            warn ("Unknown command `%s'\n", split[0]);

cleanup_line:
        g_free (line);
        g_strfreev (split);
    }
}

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", "[int]" },
        { "height", 0, 0, G_OPTION_ARG_INT, &params->height, "Height", "[int]" },
        { "num-radios", 0, 0, G_OPTION_ARG_INT, &params->num_radios, "Number of radios", "[int]" },
        { "num-darks", 0, 0, G_OPTION_ARG_INT, &params->num_darks, "Number of darks", "[int]" },
        { "radios", 0, 0, G_OPTION_ARG_STRING, &params->radios, "Radios", "[path|glob]" },
        { "darks", 0, 0, G_OPTION_ARG_STRING, &params->darks, "Darks", "[path|glob]" },
        { "flats", 0, 0, G_OPTION_ARG_STRING, &params->flats, "Flats", "[path|glob]" },
        { "dark-scale", 0, 0, G_OPTION_ARG_DOUBLE, &params->dark_scale, "Dark scale", "[float]" },
        { "output", 0, 0, G_OPTION_ARG_STRING, &params->output, "Output", "[path]" },
        { "theta", 0, 0, G_OPTION_ARG_DOUBLE, &params->theta, "Tilt (theta)", "[float]" },
        { "tau", 0, 0, G_OPTION_ARG_DOUBLE, &params->tau, "Pixel size (theta)", "[float]" },
        { "psi", 0, 0, G_OPTION_ARG_DOUBLE, &params->psi, "Misalignment (psi)", "[float]" },
        { "px", 0, 0, G_OPTION_ARG_DOUBLE, &params->px, "X coordinate of axis", "[float]" },
        { "py", 0, 0, G_OPTION_ARG_DOUBLE, &params->py, "Y coordinate of axis", "[float]" },
        { "px-variation", 0, 0, G_OPTION_ARG_DOUBLE, &params->px_variation, "X axis variation", "[float]" },
        { "vx", 0, 0, G_OPTION_ARG_DOUBLE, &params->v_origin[0], "X coordinate of box origin", "[float]" },
        { "vy", 0, 0, G_OPTION_ARG_DOUBLE, &params->v_origin[1], "Y coordinate of box origin", "[float]" },
        { "vz", 0, 0, G_OPTION_ARG_DOUBLE, &params->v_origin[2], "Z coordinate of box origin", "[float]" },
        { "vw", 0, 0, G_OPTION_ARG_INT, &params->v_size[0], "Width of box", "[int]" },
        { "vh", 0, 0, G_OPTION_ARG_INT, &params->v_size[1], "Height of box", "[int]" },
        { "vd", 0, 0, G_OPTION_ARG_INT, &params->v_size[2], "Depth of box", "[int]" },
        { "interactive", 0, 0, G_OPTION_ARG_NONE, &params->interactive, "Start interactive mode", "" },
        { NULL }
    };

    context = g_option_context_new ("- laminographic reconstruction");
    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);
    }

    /* Copy entry information for later reference */
    params->entries = g_memdup (entries, sizeof (entries));

    if (params->interactive)
        goto parse_params_cleanup;

    if (!params_okay (params)) {
        err ("Parameters missing.\n\n");
        g_print ("%s\n", g_option_context_get_help (context, TRUE, NULL));
        exit (1);
    }

parse_params_cleanup:
    g_option_context_free (context);
}

int
main (int argc, char **argv)
{
    Params params = {
        .interactive = FALSE,
        .radios = NULL,
        .darks = NULL,
        .flats = NULL,
        .output = "volume-%i.tif",
        .width = 0,
        .height = 0,
        .num_radios = 0,
        .num_darks = 1,
        .theta = G_MAXDOUBLE,
        .dark_scale = 1.0,
        .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);

    if (params.interactive)
        start_shell (&params);
    else
        run_reconstruction (&params, NULL);

    return 0;
}