Info2-Pr-NeuroNetz/Startcode/mnistVisualization.c

#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <math.h>
#include "mnistVisualization.h"
#include "raylib.h"

#define MAX_TEXT_LEN 100
#define BRUSH_RADIUS 2.0f
#define BRUSH_STEP 0.25f
#define BROWSE_TEXT "Use left and right key to navigate ... (D: draw mode)"
#define DRAW_TEXT "Left mouse: draw  Right mouse: erase  C: clear  D: browse"

typedef enum
{
    BROWSE_MODE,
    DRAW_MODE
} VisualizationMode;

typedef struct
{
    Vector2 position;
    const GrayScaleImageSeries *series;
    const unsigned char *predictions;
    unsigned int currentIdx;
    Vector2 pixelSize;
} MnistVisualization;

typedef struct
{
    char text[MAX_TEXT_LEN];
    Vector2 position;
    unsigned int fontSize;
    Color color;
} TextLabel;

typedef struct
{
    GrayScaleImage image;
    Vector2 lastPosition;
    int isStroking;
    int isEmpty;
    int prediction;
} DrawingCanvas;

static TextLabel *createTextLabel(const char *text, unsigned int fontSize, Color color)
{
    TextLabel *label = NULL;

    if(text != NULL)
    {
        label = (TextLabel *)malloc(sizeof(TextLabel));

        if(label != NULL)
        {
            strncpy(label->text, text, MAX_TEXT_LEN);
            label->text[MAX_TEXT_LEN-1] = '\0';
            label->position.x = 0;
            label->position.y = 0;
            label->fontSize = fontSize;
            label->color = color;
        }
    }

    return label;
}

static MnistVisualization *createVisualizationContainer(const GrayScaleImageSeries *series, const unsigned char predictions[], Vector2 size)
{
    MnistVisualization *container = NULL;
    
    if(size.x > 0 && size.y > 0 && series != NULL && series->images != NULL && series->count > 0 && predictions != NULL)
    {
        container = (MnistVisualization *)calloc(1, sizeof(MnistVisualization));

        if(container != NULL)
        {
            Vector2 pixelSize = {(int)(size.x / series->images[0].width), (int)(size.y / series->images[0].height)};
            container->pixelSize = pixelSize;
            container->series = series;
            container->predictions = predictions;
        }
    }

    return container;
}

static void setLabelText(TextLabel *label, const char *text)
{
    strncpy(label->text, text, MAX_TEXT_LEN);
    label->text[MAX_TEXT_LEN-1] = '\0';
}

static DrawingCanvas *createDrawingCanvas(unsigned int width, unsigned int height)
{
    DrawingCanvas *canvas = (DrawingCanvas *)calloc(1, sizeof(DrawingCanvas));

    if(canvas != NULL)
    {
        canvas->image.buffer = (GrayScalePixelType *)calloc(width * height, sizeof(GrayScalePixelType));

        if(canvas->image.buffer != NULL)
        {
            canvas->image.width = width;
            canvas->image.height = height;
            canvas->isEmpty = 1;
        }
        else
        {
            free(canvas);
            canvas = NULL;
        }
    }

    return canvas;
}

static void freeDrawingCanvas(DrawingCanvas *canvas)
{
    if(canvas != NULL)
    {
        free(canvas->image.buffer);
        free(canvas);
    }
}

static void clearDrawingCanvas(DrawingCanvas *canvas)
{
    memset(canvas->image.buffer, 0, canvas->image.width * canvas->image.height * sizeof(GrayScalePixelType));
    canvas->isEmpty = 1;
}

// Weicher Pinsel: innerer Bereich voll weiss, nach aussen linear auslaufend (aehnlich den MNIST-Strichen)
static void paintAt(DrawingCanvas *canvas, Vector2 position, int erase)
{
    int minX = (int)floorf(position.x - BRUSH_RADIUS);
    int maxX = (int)ceilf(position.x + BRUSH_RADIUS);
    int minY = (int)floorf(position.y - BRUSH_RADIUS);
    int maxY = (int)ceilf(position.y + BRUSH_RADIUS);

    for(int y = minY; y <= maxY; y++)
    {
        for(int x = minX; x <= maxX; x++)
        {
            if(x >= 0 && y >= 0 && x < (int)canvas->image.width && y < (int)canvas->image.height)
            {
                float dx = x + 0.5f - position.x;
                float dy = y + 0.5f - position.y;
                float strength = 2.0f * (1.0f - sqrtf(dx * dx + dy * dy) / BRUSH_RADIUS);
                GrayScalePixelType *pixel = &canvas->image.buffer[y * canvas->image.width + x];

                if(strength > 0)
                {
                    if(erase)
                        *pixel = 0;
                    else
                    {
                        float value = strength >= 1.0f ? 255.0f : 255.0f * strength;

                        if(value > *pixel)
                            *pixel = (GrayScalePixelType)value;
                    }
                }
            }
        }
    }
}

// Zwischenpunkte interpolieren, damit bei schnellen Mausbewegungen keine Luecken entstehen
static void paintLine(DrawingCanvas *canvas, Vector2 from, Vector2 to, int erase)
{
    float dx = to.x - from.x;
    float dy = to.y - from.y;
    int steps = (int)ceilf(sqrtf(dx * dx + dy * dy) / BRUSH_STEP);

    for(int i = 0; i <= steps; i++)
    {
        float t = steps > 0 ? (float)i / steps : 0.0f;
        Vector2 point = {from.x + t * dx, from.y + t * dy};
        paintAt(canvas, point, erase);
    }
}

// MNIST-Ziffern sind ueber ihren Schwerpunkt zentriert; die Zeichnung wird fuer die Vorhersage genauso verschoben
static void centerByMass(const GrayScaleImage source, GrayScaleImage target)
{
    double sum = 0;
    double centerX = 0;
    double centerY = 0;

    memset(target.buffer, 0, target.width * target.height * sizeof(GrayScalePixelType));

    for(unsigned int i = 0; i < source.width * source.height; i++)
    {
        sum += source.buffer[i];
        centerX += (double)(i % source.width) * source.buffer[i];
        centerY += (double)(i / source.width) * source.buffer[i];
    }

    if(sum > 0)
    {
        int shiftX = (int)lround((source.width - 1) / 2.0 - centerX / sum);
        int shiftY = (int)lround((source.height - 1) / 2.0 - centerY / sum);

        for(int y = 0; y < (int)source.height; y++)
        {
            for(int x = 0; x < (int)source.width; x++)
            {
                int newX = x + shiftX;
                int newY = y + shiftY;

                if(newX >= 0 && newY >= 0 && newX < (int)target.width && newY < (int)target.height)
                    target.buffer[newY * target.width + newX] = source.buffer[y * source.width + x];
            }
        }
    }
}

static void predictDrawing(DrawingCanvas *canvas, ClassifierFunctionType classify, const void *classifierContext)
{
    GrayScaleImage centered = canvas->image;
    centered.buffer = (GrayScalePixelType *)malloc(canvas->image.width * canvas->image.height * sizeof(GrayScalePixelType));
    canvas->prediction = -1;

    if(centered.buffer != NULL)
    {
        centerByMass(canvas->image, centered);
        canvas->prediction = classify(&centered, classifierContext);
        free(centered.buffer);
    }
}

static void updateDrawingCanvas(DrawingCanvas *canvas, const MnistVisualization *container, ClassifierFunctionType classify, const void *classifierContext, TextLabel *predictionLabel)
{
    int isDrawing = IsMouseButtonDown(MOUSE_BUTTON_LEFT);
    int isErasing = IsMouseButtonDown(MOUSE_BUTTON_RIGHT);

    if(IsKeyReleased(KEY_C))
        clearDrawingCanvas(canvas);

    if(isDrawing || isErasing)
    {
        Vector2 mouse = GetMousePosition();
        Vector2 position = {(mouse.x - container->position.x) / container->pixelSize.x, (mouse.y - container->position.y) / container->pixelSize.y};

        if(canvas->isStroking)
            paintLine(canvas, canvas->lastPosition, position, isErasing);
        else
            paintAt(canvas, position, isErasing);

        canvas->lastPosition = position;
        canvas->isStroking = 1;
        canvas->isEmpty = 0;
        predictDrawing(canvas, classify, classifierContext);
    }
    else
        canvas->isStroking = 0;

    if(canvas->isEmpty)
        setLabelText(predictionLabel, "Draw a digit ...");
    else if(canvas->prediction < 0)
        setLabelText(predictionLabel, "Prediction failed ...");
    else
        snprintf(predictionLabel->text, MAX_TEXT_LEN, "Predicted label:  %d", canvas->prediction);
}

static void drawTextLabel(const TextLabel *label)
{
    DrawText(label->text, label->position.x, label->position.y, label->fontSize, label->color);
}

static void drawDigit(const GrayScaleImage image, Vector2 position, Vector2 pixelSize)
{
    float xOffset = 0;
    float yOffset = position.y - pixelSize.y;
    Color color = {0};
    color.a = 255;

    for(int i = 0; i < image.width * image.height; i++)
    {
        if(i % image.width == 0)
        {
            xOffset = position.x;
            yOffset += pixelSize.y;
        }

        color.r = image.buffer[i];
        color.b = image.buffer[i];
        color.g = image.buffer[i];

        DrawRectangle(xOffset, yOffset, pixelSize.x, pixelSize.y, color);
        xOffset += pixelSize.x;
    }
}

static void drawAll(const MnistVisualization *container, const TextLabel *navigationLabel, const TextLabel *predictionLabel)
{
    BeginDrawing();
    ClearBackground(BLACK);
    
    drawDigit(container->series->images[container->currentIdx], container->position, container->pixelSize);
    drawTextLabel(navigationLabel);
    drawTextLabel(predictionLabel);
        
    EndDrawing();
}

static void drawCanvasAll(const DrawingCanvas *canvas, const MnistVisualization *container, const TextLabel *navigationLabel, const TextLabel *predictionLabel)
{
    // Hilfsrahmen: MNIST-Ziffern fuellen etwa die inneren 20x20 von 28x28 Pixeln
    const float margin = canvas->image.width / 7.0f;
    Rectangle guide = {container->position.x + margin * container->pixelSize.x, container->position.y + margin * container->pixelSize.y,
                       (canvas->image.width - 2 * margin) * container->pixelSize.x, (canvas->image.height - 2 * margin) * container->pixelSize.y};

    BeginDrawing();
    ClearBackground(BLACK);

    drawDigit(canvas->image, container->position, container->pixelSize);
    DrawRectangleLinesEx(guide, 1, DARKGRAY);
    drawTextLabel(navigationLabel);
    drawTextLabel(predictionLabel);

    EndDrawing();
}

static int checkUserInput()
{
    int inputResult = 0;

    if(IsKeyReleased(KEY_LEFT))
        inputResult = -1;
    else if(IsKeyReleased(KEY_RIGHT))
        inputResult = 1;
    
    return inputResult;
}

static void updateDisplayContainer(MnistVisualization *container, int updateDirection)
{
    int newIndex = (int)container->currentIdx + updateDirection;

    if(newIndex < 0)
        newIndex = 0;
    else if(newIndex >= container->series->count)
        newIndex = container->series->count - 1;

    container->currentIdx = newIndex;
}

static void updatePredictionLabel(TextLabel *predictionLabel, unsigned char trueLabel, unsigned char predictedLabel)
{
    snprintf(predictionLabel->text, MAX_TEXT_LEN, "True label:        %u\nPredicted label:  %u", trueLabel, predictedLabel);
}

static void update(MnistVisualization *container, TextLabel *predictionLabel, int updateDirection)
{
    updateDisplayContainer(container, updateDirection);
    updatePredictionLabel(predictionLabel, container->series->labels[container->currentIdx], container->predictions[container->currentIdx]);
}

void showMnist(unsigned int windowWidth, unsigned int windowHeight, const GrayScaleImageSeries *series, const unsigned char predictions[], ClassifierFunctionType classify, const void *classifierContext)
{
    const Vector2 windowSize = {windowWidth, windowHeight};
    VisualizationMode mode = BROWSE_MODE;

    MnistVisualization *container = createVisualizationContainer(series, predictions, windowSize);
    DrawingCanvas *canvas = container != NULL ? createDrawingCanvas(series->images[0].width, series->images[0].height) : NULL;
    TextLabel *navigationLabel = createTextLabel(BROWSE_TEXT, 20, WHITE);
    TextLabel *predictionLabel = createTextLabel("", 20, WHITE);
    predictionLabel->position.x = 10;
    predictionLabel->position.y = windowSize.y - 50;

    if(container != NULL && canvas != NULL && classify != NULL && navigationLabel != NULL && predictionLabel != NULL)
    {
        InitWindow(windowSize.x, windowSize.y, "MNIST Browser");

        SetTargetFPS(60);

        while (!WindowShouldClose())
        {
            if(IsKeyReleased(KEY_D))
            {
                mode = mode == BROWSE_MODE ? DRAW_MODE : BROWSE_MODE;
                setLabelText(navigationLabel, mode == BROWSE_MODE ? BROWSE_TEXT : DRAW_TEXT);
            }

            if(mode == BROWSE_MODE)
            {
                int updateDirection = checkUserInput();
                update(container, predictionLabel, updateDirection);
                drawAll(container, navigationLabel, predictionLabel);
            }
            else
            {
                updateDrawingCanvas(canvas, container, classify, classifierContext, predictionLabel);
                drawCanvasAll(canvas, container, navigationLabel, predictionLabel);
            }
        }
    }

    CloseWindow();

    free(container);
    freeDrawingCanvas(canvas);
    free(navigationLabel);
    free(predictionLabel);
}