jpeg.py

import jpeg_headers
from typing import Dict, Tuple, List
from time import time
import math
import tkinter

# References
#   https://www.w3.org/Graphics/JPEG/itu-t81.pdf
#   https://impulseadventure.com

IMAGE_TO_OPEN = "test_images/lena.jpg"

start_time = time()


class PhotoDisplay:
    def __init__(self, tk, height, width, rgb_array: List[List[Tuple[int, int, int]]]):
        self.width = width
        self.height = height

        self.i = tkinter.PhotoImage(width=width, height=height)
        pixels = " ".join(("{" + " ".join(("#%02x%02x%02x" % rgb_array[color_row][color_col]
                                           for color_col in range(width))) + "}" for color_row in range(height)))
        self.i.put(pixels, (0, 0, self.width, self.height))

        self.canvas = tkinter.Canvas(tk, width=self.width, height=self.height)
        self.canvas.pack()
        self.canvas.create_image(0, 0, image=self.i, anchor=tkinter.NW)

        self.canvas.bind("<ButtonPress-1>", self.scroll_start)
        self.canvas.bind("<B1-Motion>", self.scroll_move)

    def scroll_start(self, event):
        self.canvas.scan_mark(event.x, event.y)

    def scroll_move(self, event):
        self.canvas.scan_dragto(event.x, event.y, gain=1)


def bit_from_bytearray(arr: bytearray, bit_idx: int, order: str) -> int:
    if order == "little":
        return (arr[bit_idx // 8] & (0b1 << (bit_idx % 8))) >> (bit_idx % 8)
    else:
        return (arr[bit_idx // 8] & (0b1 << (7 - (bit_idx % 8)))) >> (7 - (bit_idx % 8))


def bits_from_bytearray(arr: bytearray, start_idx: int, num_bits: int, order: str) -> int:
    out = 0
    for bit_idx in range(start_idx, start_idx + num_bits):
        out = (out << 1) | bit_from_bytearray(arr, bit_idx, order)
    return out


def get_signed_value(bits: int, num_bits: int) -> int:
    # EXTEND
    # Table F.1
    # Figure F.12

    # if output should be positive, output our bits
    # if output should be negative, output the maximum negative value plus our bits
    if bits < 2**(num_bits - 1):  # Check if bits is less than the middle value
        min_val = (-1 << num_bits) + 1
        return min_val + bits

    return bits


def get_next_huffman_value(data: bytearray, data_pos: int, huff_table: Dict[Tuple[int, int], int],
                           debug_print=False) -> Tuple[int, int]:
    # DECODE
    encoded_bits = bit_from_bytearray(data, data_pos, "big")
    start_bit = data_pos
    curr_pos = data_pos + 1

    while (encoded_bits, curr_pos - start_bit) not in huff_table:
        encoded_bits = (encoded_bits << 1) | bit_from_bytearray(scan_data, curr_pos, "big")
        curr_pos += 1

    num_bits = curr_pos - start_bit

    if debug_print:
        print(f"encoded: {encoded_bits:0{num_bits}b}, length: {num_bits}")
    return huff_table[(encoded_bits, num_bits)], num_bits


def ycbcr_to_rgb(lum, chrom_blue, chrom_red) -> Tuple[int, int, int]:
    # https://www.impulseadventure.com/photo/jpeg-color-space.html
    red = chrom_red * (2 - 2 * 0.299) + lum
    blue = chrom_blue * (2 - 2 * 0.114) + lum
    green = (lum - 0.114 * blue - 0.299 * red) / 0.587
    return min(255, max(0, round(red + 128))), min(255, max(0, round(green + 128))), min(255, max(0, round(blue + 128)))


# Define JPEG segment names according to offset from 0xFFC0
# Table B.1
seg_names = ["SOF0 - Baseline DCT; Huffman", "SOF1 - Extended sequential DCT; Huffman",
             "SOF2 - Progressive DCT; Huffman", "SOF3 - Lossless (sequential); Huffman",
             "DHT - Define Huffman table(s)", "SOF5 - Differential sequential DCT; Huffman",
             "SOF6 - Differential progressive DCT; Huffman", "SOF7 - Differential lossless (sequential); Huffman",
             "JPG - Reserved for JPEG extensions", "SOF9 - Extended sequential DCT; Arithmetic",
             "SOF10 - Progressive DCT; Arithmetic", "SOF11 - Lossless (sequential); Arithmetic",
             "DAC - Define arithmetic coding conditioning(s)", "SOF13 - Differential sequential DCT; Arithmetic",
             "SOF14 - Differential progressive DCT; Arithmetic",
             "SOF15 - Differential lossless (sequential); Arithmetic",
             "RST0 - Restart with modulo 8 count 0", "RST0 - Restart with modulo 8 count 1",
             "RST0 - Restart with modulo 8 count 2", "RST0 - Restart with modulo 8 count 3",
             "RST0 - Restart with modulo 8 count 4", "RST0 - Restart with modulo 8 count 5",
             "RST0 - Restart with modulo 8 count 6", "RST0 - Restart with modulo 8 count 7",
             "SOI - Start of image", "EOI - End of image", "SOS - Start of scan", "DQT - Define quantization table(s)",
             "DNL - Define number of lines", "DRI - Define restart interval", "DHP - Define hierarchical progression",
             "EXP - Expand reference components", "JFIF header"]
seg_names.extend(["[Reserved: Application segments]"] * 16)
seg_names.extend(["[Reserved: JPEG extension]"] * 14)
seg_names.extend(["Comment", "[Invalid]"])

# Create necessary variables to store header information
huff_tables = {}
quant_tables = {}
sof: jpeg_headers.StartOfFrame

# Precalculate IDCT constants (A.3.3)
idct_lookup = []
for y in range(8):
    idct_row = []
    for x in range(8):
        uv_matrix = []
        for u in range(8):
            uv_row = []
            for v in range(8):
                cu = (1 / math.sqrt(2)) if u == 0 else 1
                cv = (1 / math.sqrt(2)) if v == 0 else 1
                uv_row.append(cu * cv *
                              math.cos(((2 * x + 1) * u * math.pi) / 16) * math.cos(((2 * y + 1) * v * math.pi) / 16))
            uv_matrix.append(uv_row)
        idct_row.append(uv_matrix)
    idct_lookup.append(idct_row)

with open(IMAGE_TO_OPEN, "rb") as f:
    block_id_bytes = f.read(2)
    while block_id_bytes:
        block_id = int.from_bytes(block_id_bytes, byteorder="big")
        pos = f.tell() - 2

        if block_id < 0xFFC0:
            print("Segment ID expected, not found.")
            break  # JPEG_SEG_ERR

        print(f"*** Marker: {seg_names[block_id - 0xFFC0]}, (0x{block_id:04X}) ***")
        print(f"OFFSET: {pos} (0x{pos:X})\n")

        if block_id == 0xFFD9:    # -------------------------------------- End of Image
            print(f"Decoding process took {time() - start_time} seconds")

            # Show decoded RGB array
            t = tkinter.Tk()
            t.resizable(False, False)
            t.title("decoded image")
            display = PhotoDisplay(t, sof.num_lines, sof.samples_per_line, image_rgb)
            t.mainloop()

            break
        elif block_id == 0xFFDA:  # -------------------------------------- Start of Scan
            sos = jpeg_headers.read_start_of_scan(f)
            assert sof is not None

            # Get maximum sampling factors and number of MCUs
            max_h_sampling_factor = 0
            max_v_sampling_factor = 0
            for component in sof.components:
                max_h_sampling_factor = max(max_h_sampling_factor, component.h_sampling_factor)
                max_v_sampling_factor = max(max_v_sampling_factor, component.v_sampling_factor)

            mcu_size_x = 8 * max_h_sampling_factor
            mcu_size_y = 8 * max_v_sampling_factor

            num_mcu_x = math.ceil(sof.samples_per_line / mcu_size_x)
            num_mcu_y = math.ceil(sof.num_lines / mcu_size_y)

            print(f"MCU Size: {mcu_size_x} x {mcu_size_y}")
            print(f"{num_mcu_x} MCU cols, {num_mcu_y} MCU rows\n")

            # Read data that's left
            start_pos = f.tell()
            f.seek(0, 2)
            end_pos = f.tell()

            f.seek(start_pos)
            scan_data = bytearray(f.read(end_pos - start_pos))

            # Scan for next marker and remove all stuff bytes in scan data
            marker_pos = None
            marker_pos_diff = 0

            for i in range(len(scan_data) - 2, 0, -1):  # scan backwards
                marker_code = scan_data[i:i + 2]

                # Remove stuff byte
                if marker_code == b'\xFF\x00':
                    scan_data.pop(i + 1)
                    marker_pos_diff += 1

                # Marker found if 0xFF exists without a stuff byte after
                elif marker_code > b'\xFF\x00':
                    if b'\xFF\xD0' <= marker_code <= b'\xFF\xD7':
                        # TODO: RSTm (restart marker) found
                        # remove the 0xFF and continue
                        print('restart marker found, skipping...')
                        scan_data.pop(i)
                        marker_pos_diff += 1
                    else:
                        marker_pos = i
                        marker_pos_diff = 0

            assert marker_pos is not None

            # Set scan data from start to next marker
            scan_data = scan_data[:marker_pos - marker_pos_diff]
            f.seek(start_pos + marker_pos)

            print("Scan data: (after bitstuff removed)")
            print("  " + "".join(f"{b:02x} " +
                                 ("\n  " if (idx + 1) % 36 == 0 else "") for idx, b in enumerate(scan_data[:720])))
            if len(scan_data) > 720:
                print("WARNING: Dump truncated.")
            print()
            # block_id_bytes = f.read(2)
            # continue

            # Initialize array storing final RGB values
            image_rgb = [[(0, 0, 0) for _ in range(sof.samples_per_line)] for _ in range(sof.num_lines)]

            # F.2.1.2
            # Figure E.9
            # Figure E.10
            curr_bit = 0

            # Initialize DC predictions for each component to zero
            predictions = [0 for _ in range(len(sos.components))]

            # Loop through all MCUs in image
            for mcu_row in range(num_mcu_y):
                print(f"Processing MCU row {mcu_row}")
                for mcu_col in range(num_mcu_x):
                    debug = False
                    if debug:
                        print(f"\nMCU {mcu_row}, {mcu_col}")

                    # List of MCUs (could contain luminance, chrominance blue, chrominance red or just one)
                    mcu_arr = []

                    # Decode each MCU
                    for component_idx, component in enumerate(sos.components):
                        frame_component = None
                        for c in sof.components:
                            if c.identifier == component.selector:
                                frame_component = c
                                break
                        assert frame_component is not None

                        quant_table = quant_tables[frame_component.quant_table_dest]
                        dc_huff_table = huff_tables[component.dc_table, 0]
                        ac_huff_table = huff_tables[component.ac_table, 1]

                        # Initialize 2D array for MCU
                        mcu = [[0 for _ in range(8 * frame_component.h_sampling_factor)]
                               for _ in range(8 * frame_component.v_sampling_factor)]

                        # Go through all data units in order specified by A.2.3
                        for data_unit_row in range(frame_component.v_sampling_factor):
                            for data_unit_col in range(frame_component.h_sampling_factor):
                                if debug:
                                    if frame_component.quant_table_dest == 0:
                                        print("Lum")
                                    else:
                                        print("Chr")

                                # Decode DC coefficient
                                dc_code, length = get_next_huffman_value(scan_data, curr_bit, dc_huff_table, debug)
                                curr_bit += length

                                # RECEIVE
                                # A.3.5, F.2.1.3.1
                                additional_bits = bits_from_bytearray(scan_data, curr_bit, dc_code, "big")
                                curr_bit += dc_code

                                diff = get_signed_value(additional_bits, dc_code)
                                abs_dc_value = predictions[component_idx] + diff

                                predictions[component_idx] = abs_dc_value

                                if debug:
                                    start_byte = (curr_bit - dc_code - length) // 8
                                    print(f"val: {diff}, coeff: 00=DC")
                                    print(f"val_bits: {additional_bits:0{dc_code}b}")
                                    print(f"data: 0x "
                                          f"{' '.join(hex(b)[2:] for b in scan_data[start_byte:start_byte + 4])}\n")

                                # Start decoding DCT matrix
                                dct_coeffs = [0 for _ in range(64)]
                                dct_coeffs[0] = abs_dc_value

                                # Decode AC coefficients, F.2.2.2
                                # Figure F.13
                                k = 0
                                while k != 63:
                                    k += 1

                                    # rs (8 bits)-> rrrrssss
                                    rs, length = get_next_huffman_value(scan_data, curr_bit, ac_huff_table, debug)
                                    curr_bit += length

                                    rrrr = rs >> 4  # Skip
                                    ssss = rs & 0b1111  # Coded length

                                    if ssss == 0:
                                        if rrrr == 15:
                                            k += 15
                                            continue
                                        else:
                                            if debug:
                                                print("EOB")
                                            break

                                    k += rrrr

                                    # Decode_ZZ
                                    # ZZ(k) = RECEIVE(ssss)
                                    additional_bits = bits_from_bytearray(scan_data, curr_bit, ssss, "big")
                                    curr_bit += ssss

                                    # ZZ(k) = EXTEND(ZZ(k), ssss)
                                    v = get_signed_value(additional_bits, ssss)

                                    if debug:
                                        start_byte = (curr_bit - ssss - length) // 8
                                        print(f"val: {v}, coeff: {k-rrrr:02d}..{k:02d}, skip: {rrrr}")
                                        print(f"val_bits: {additional_bits:0{ssss}b}")
                                        print(f"data: 0x "
                                              f"{' '.join(hex(b)[2:] for b in scan_data[start_byte:start_byte+4])}\n")

                                    dct_coeffs[k] = v

                                # Zig-zag reorder AC and DC coefficient list into DCT matrix
                                # Multiply each value by its corresponding value in the quantization table to dequantize
                                dct_matrix = [[0 for _ in range(8)] for _ in range(8)]
                                for i, coeff in enumerate(dct_coeffs):
                                    row, col = jpeg_headers.zigzag[i]
                                    dct_matrix[row][col] = coeff * quant_table[row][col]

                                # Perform IDCT (A.3.3)
                                for y in range(8):
                                    for x in range(8):
                                        val = 0
                                        for u in range(8):
                                            for v in range(8):
                                                val += idct_lookup[y][x][u][v] * dct_matrix[v][u]
                                        val /= 4

                                        # Assign value in MCU
                                        mcu[(data_unit_row * 8) + y][(data_unit_col * 8) + x] = val

                        # Expand MCU to maximum MCU size by duplicating values vertically or horizontally
                        horiz_multiplier = max_h_sampling_factor // frame_component.h_sampling_factor
                        vert_multiplier = max_v_sampling_factor // frame_component.v_sampling_factor

                        if vert_multiplier > 1 or horiz_multiplier > 1:
                            mcu = [[val for val in row for _ in range(horiz_multiplier)]
                                   for row in mcu for _ in range(vert_multiplier)]

                        # Append MCU to MCU array
                        mcu_arr.append(mcu)

                    # TODO: Handle images with one component
                    # Convert all Y, Cb, and Cr component values to RGB and store them in array
                    for i in range(mcu_size_y):
                        # Break if MCU goes past y bounds of image
                        if (mcu_row * mcu_size_y) + i >= sof.num_lines:
                            break

                        for j in range(mcu_size_x):
                            # Break if MCU goes past x bounds of image
                            if (mcu_col * mcu_size_x) + j >= sof.samples_per_line:
                                break

                            image_rgb[(mcu_row * mcu_size_y) + i][(mcu_col * mcu_size_x) + j] = \
                                ycbcr_to_rgb(mcu_arr[0][i][j], mcu_arr[1][i][j], mcu_arr[2][i][j])

        elif block_id == 0xFFD8:  # -------------------------------------- Start of Image
            pass
        elif block_id == 0xFFC4:  # -------------------------------------- Define Huffman Table
            for table in jpeg_headers.define_huffman_table(f):
                huff_tables[table.dest_id, table.table_class] = table.huff_data
        elif block_id == 0xFFDB:  # -------------------------------------- Define Quantization Table
            for table in jpeg_headers.define_quantization_table(f):
                quant_tables[table.dest_id] = table.table
        elif block_id == 0xFFC0 or block_id == 0xFFC1:  # ---------------- Start of Frame
            sof = jpeg_headers.read_start_of_frame(f)
        else:  # All other segments have length specified at the start, skip for now
            size = int.from_bytes(f.read(2), byteorder="big")
            f.seek(size - 2, 1)

        block_id_bytes = f.read(2)