decode.c

#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>

typedef struct {
    uint8_t counts[17];
    uint8_t symbols[256];

    int mincode[17];
    int maxcode[17];
    int valptr[17];
} HuffmanTable;

typedef struct {
    int coefficients[64];
} Block;

typedef struct {
    int v[16][16][3];
} MCU;

int width = 160;
int height = 90;
uint8_t* buffer;
FILE* f;
uint8_t q_tables[4][64];
HuffmanTable dc_tables[4];
HuffmanTable ac_tables[4];
uint8_t bit_buffer = 0;
uint8_t bits_left = 0;
HuffmanTable* use_dc[4];
HuffmanTable* use_ac[4];
int last_dc[4] = {0, 0, 0, 0};
uint8_t component_qt[4];

void init_buffer(int width, int height) {
    buffer = malloc(width * height * 3);
    memset(buffer, 0, width * height * 3);
}

void flush_buffer(const char* ppm_path, const char* png_path) {
    FILE *f = fopen(ppm_path, "w");
    fprintf(f, "P3\n");
    fprintf(f, "%d %d\n", width, height);
    fprintf(f, "255\n");

    for (int i = 0; i < width * height * 3; i++) {
        uint8_t b = buffer[i];
        fprintf(f, "%d ", b);
    }
    fclose(f);
    char s[256];
    sprintf(s, "convert %s %s", ppm_path, png_path);
    system(s);
}

void set_pixel(int x, int y, int r, int g, int b) {
   int offset = (y * width + x) * 3;
   buffer[offset + 0] = r;
   buffer[offset + 1] = g;
   buffer[offset + 2] = b;
}

uint8_t read_byte() {
    uint8_t b = fgetc(f);
    return b;
}

uint16_t read_word() {
    uint8_t b0 = fgetc(f);
    uint8_t b1 = fgetc(f);
    return ((uint16_t)b0 << 8) | b1;
}

int read_bit() {
    if (!bits_left) {
        bit_buffer = read_byte();
        bits_left = 8;
    }
    int value = bit_buffer >> 7;
    bit_buffer <<= 1;
    bits_left--;
    return value;
}

int read_huffman_symbol(HuffmanTable* ht) {
    int code = 0;
    for (int i = 1; i <= 16; i++) {
        code <<= 1;
        code |= read_bit();
        if (code >= ht->mincode[i] && code <= ht->maxcode[i]) {
            int entry = ht->valptr[i];
            entry += code - ht->mincode[i];
            return ht->symbols[entry];
        }
    }
}

int read_n_bits_sign_extend(int n) {
    int zdwh = 0; // Zahl, die wir haben
    for (int i = 0; i < n; i++) {
        uint8_t bit = read_bit();
        zdwh <<= 1;
        zdwh |= bit;
    }
    // result is in zdwh now, now sign extend if MSB is 1
    if (!(zdwh >> (n - 1))) zdwh -= (1 << n) - 1;
    return zdwh;
}

void assert(int condition, const char* message) {
    if (!condition) {
        printf("Assertion failed: NOT! %s\n", message);
        exit(1);
    }
}

void parse_app0() {
    uint16_t length = read_word();
    fseek(f, length - 2, SEEK_CUR);
}

void parse_dqt() {
    uint16_t length = read_word();
    assert(length == 67, "DQT: exactly one quantization table present");
    uint8_t pf = read_byte();
    uint8_t pq = pf >> 4; // 4 obere Bits
    uint8_t tq = pf & 15; // 00001111 => 4 untere Bits bleiben übrig
    assert(pq == 0, "DQT: 8 bit precision");
    assert(tq < 4, "DQT: slot within range");
    fread(q_tables[tq], 64, 1, f);
    printf("DQT: Read quantization table #%d\n", tq);
}

void parse_sof0() {
    uint16_t length = read_word();
    assert(length == 17, "SOF0: found expected segment length");
    uint8_t p = read_byte();
    assert(p == 8, "SOF0: sample precision is 8 bits");
    height = read_word();
    width = read_word();
    init_buffer(width, height);
    uint8_t nf = read_byte();
    assert(nf == 3, "SOF0: number of image components = 3");

    // now parse the image components
    assert(read_byte() == 0x01, "component 1");
    assert(read_byte() == 0x22, "component 1: subsampling 2:2");
    component_qt[1] = read_byte();

    assert(read_byte() == 0x02, "component 2");
    assert(read_byte() == 0x11, "component 2: subsampling 1:1");
    component_qt[2] = read_byte();

    assert(read_byte() == 0x03, "component 3");
    assert(read_byte() == 0x11, "component 3: subsampling 1:1");
    component_qt[3] = read_byte();

    printf("Read SOF0 marker: image is %d x %d pixels\n", width, height);
}

void parse_dht() {
    uint16_t length = read_word();
    uint8_t pf = read_byte();
    uint8_t tc = pf >> 4;
    uint8_t th = pf & 15;
    assert(tc < 2, "DHT: DC or AC table detected it has");
    assert(th < 4, "DHT: not exceeded bounds 0..3 it has");
    HuffmanTable* ht = (tc == 0) ? (&dc_tables[th]) : (&ac_tables[th]);
    printf("DHT: Reading %s table #%d.\n", tc == 0 ? "DC" : "AC", th);
    int symbol_count = 0;
    for (int i = 1; i <= 16; i++) {
        uint8_t count = read_byte();
        ht->counts[i] = count;
        symbol_count += count;
    }
    printf("We have %d symbols in %s table %d.\n", symbol_count, tc == 0 ? "DC" : "AC", th);
    fread(ht->symbols, 1, symbol_count, f);
    int code = 0;
    int p = 0;
    for (int i = 1; i <= 16; i++) {
        if (ht->counts[i] == 0) {
            ht->mincode[i] = -1;
            ht->maxcode[i] = -1;
            ht->valptr[i]  = -1;
        } else {
            ht->mincode[i] = code;
            ht->maxcode[i] = code + ht->counts[i] - 1;
            ht->valptr[i] = p;
        }
        p += ht->counts[i];
        code += ht->counts[i];
        code <<= 1;
    }
}

int read_dc_diff(HuffmanTable* ht)  {
    int symbol = read_huffman_symbol(ht);
    int dc_diff = read_n_bits_sign_extend(symbol);
    return dc_diff;
}

void read_block(Block* block, int channel) {
    // zero block
    memset(block->coefficients, 0, sizeof(block->coefficients));
    // Read DC
    int dc_diff = read_dc_diff(use_dc[channel]);
    int dc = last_dc[channel] + dc_diff;
    last_dc[channel] = dc;
    block->coefficients[0] = dc;
    // Read AC
    int offset = 1;
    while (offset < 64) {
        uint8_t symbol = read_huffman_symbol(use_ac[channel]);
        // printf("AC symbol: %02x\n", symbol);
        if (symbol == 0x00) {
            // End Of Block
            offset = 64;
        } else if (symbol == 0xF0) {
            // ZRL
            offset += 16;
        } else {
            uint8_t leading_zeroes = symbol >> 4;
            uint8_t bits = symbol & 0xf;
            offset += leading_zeroes;
            int ac = read_n_bits_sign_extend(bits);
            // printf("AC coefficient: %d\n", ac);
            block->coefficients[offset] = ac;
            offset++;
        }
    }

    // de-quantize block
    int quant_table_index = component_qt[channel];
    for (int i = 0; i < 64; i++) block->coefficients[i] *= q_tables[quant_table_index][i];

    // perform iDCT
    // no, ackshually, we'll just fill the entire block
    // with the average value
    // HUE HUE HUE...
    int avg = block->coefficients[0] / 8 + 128;
    if (avg < 0) avg = 0;
    if (avg > 255) avg = 255;
    for (int i = 0; i < 64; i++) block->coefficients[i] = avg;
}

void parse_sos() {
    uint16_t length = read_word();
    assert(length == 12, "SOS: segment length 12 it is");
    uint8_t ns = read_byte();
    assert(ns == 3, "SOS: 3 image components it has");
    for (int i = 0; i < ns; i++) {
        uint8_t cs = read_byte();
        assert(cs < 4, "SOS: component selector is 0..7");
        uint8_t pf = read_byte();
        uint8_t dc = pf >> 4;
        uint8_t ac = pf & 15;
        use_dc[cs] = &dc_tables[dc];
        use_ac[cs] = &ac_tables[ac];
    }
    uint8_t ss = read_byte();
    uint8_t se = read_byte();
    uint8_t sa = read_byte();
    assert(ss == 0, "SOS: spectrum start == 0");
    assert(se == 63, "SOS: spectrum end == 63");
    printf("Reading huffman codes...\n");
    // // Read DC
    // int dc_diff = read_dc_diff(use_dc[1]);
    // int dc = last_dc[1] + dc_diff;
    // last_dc[1] = dc;
    // printf("DC coefficient is %d\n", dc);
    // // Read AC
    int mcu_y = 0, mcu_x = 0;
    while (mcu_y < height) {
        Block y00, y10, y01, y11, cb, cr;
        read_block(&y00, 1);
        read_block(&y10, 1);
        read_block(&y01, 1);
        read_block(&y11, 1);
        read_block(&cb, 2);
        read_block(&cr, 3);

        // assemble MCU
        MCU mcu;
        memset(mcu.v, 0, sizeof(mcu.v));
        for (int y = 0; y < 8; y++) {
            for (int x = 0; x < 8; x++) {
                mcu.v[y][x][0] = y00.coefficients[y * 8 + x];
                mcu.v[y][x + 8][0] = y10.coefficients[y * 8 + x];
                mcu.v[y + 8][x][0] = y01.coefficients[y * 8 + x];
                mcu.v[y + 8][x + 8][0] = y11.coefficients[y * 8 + x];
                for (int dy = 0; dy < 2; dy++) {
                    for (int dx = 0; dx < 2; dx++) {
                        mcu.v[y * 2 + dy][x * 2 + dx][1] = cb.coefficients[y * 8 + x];
                        mcu.v[y * 2 + dy][x * 2 + dx][2] = cr.coefficients[y * 8 + x];
                    }
                }
            }
        }

        // draw pixels

        mcu_x += 16;
        if (mcu_x >= width) {
            mcu_x = 0;
            mcu_y += 16;
        }
    }
}

int main(int argc, const char** argv) {
    printf("Hello and welcome to the awesome JPEG decoder!\n");
    // init_buffer(width, height);
    // set_pixel(width / 2, height / 2, 0, 255, 0);

    f = fopen(argv[1], "r");

    uint16_t marker = read_word();
    printf("read marker: %x\n", marker);
    assert(marker == 0xffd8, "JPEG marker found");

    while (1) {
        marker = read_word();
        printf("read marker: %x\n", marker);
        switch(marker) {
            case 0xFFE0:
                parse_app0();
                break;
            case 0xFFDB:
                parse_dqt();
                break;
            case 0xFFC0:
                parse_sof0();
                break;
            case 0xFFC4:
                parse_dht();
                break;
            case 0xFFDA:
                parse_sos();
                break;
            default:
                printf("HELP!\n");
                exit(1);
        }
    }

    fclose(f);

    flush_buffer("out.ppm", "out.png");
    free(buffer);
    return 0;
}