jpeg_encoder.py

import cv2
import numpy as np
import tables
import jpeg_helper as helper
import time

class JPEGEncoder():
    def __init__(self, quality, progress_callback=None):
        self.dc_tables = [helper.dht_to_table(tables.standard_huffman_dc0),
                          helper.dht_to_table(tables.standard_huffman_dc1)]
        self.ac_tables = [helper.dht_to_table(tables.standard_huffman_ac0),
                          helper.dht_to_table(tables.standard_huffman_ac1)]
        self.last_dc = {'y': 0, 'cb': 0, 'cr': 0}
        self.buffer = helper.Byte_Buffer()
        self.progress_callback = progress_callback
        
        # Generate and store quantization tables based on quality
        self.luminance_quant, self.chrominance_quant = self.quant_tables_quality(quality)
        self.total_blocks = 0
        self.processed_blocks = 0
    
    #Loads image file using OpenCV
    @staticmethod 
    def load_image(file_path):
        image= cv2.imread(file_path, cv2.IMREAD_COLOR)
        if image is None:
            raise ValueError(f"Could not load image from {file_path}")
        else:    
            return image
    
    #Resizes image to ensure its dimensions are multiples of 8
    @staticmethod
    def prepare_image(image):
        """Resize image to ensure its dimensions are multiples of 8."""
        h, w = image.shape[:2]
        if h % 8 == 0 and w % 8 == 0:
            return image
        new_h, new_w = (h // 8) * 8, (w // 8) * 8
        return cv2.resize(image, (new_w, new_h), cv2.INTER_CUBIC)
    
    #splits image into 8x8 blocks
    @staticmethod
    def block_split(image, h, w):
        blocks = []
        for i in range(0, h, 8):
            for j in range(0, w, 8):
                block = image[i:min(i+8, h), j:min(j+8, w)]
                # Pad the block if it's smaller than 8x8
                if block.shape[:2] != (8, 8):
                    pad_h = 8 - block.shape[0]
                    pad_w = 8 - block.shape[1]
                    block = np.pad(block, ((0, pad_h), (0, pad_w), (0, 0)), mode='edge')
                blocks.append(block)
        return blocks
    
    @staticmethod
    def quant_tables_quality(quality=100):
        if quality < 50:
            scale = 5000 / quality
        else:
            scale = 200 - 2 * quality

        luminance_quant = np.floor((tables.standard_luminance * scale + 50) / 100)
        chrominance_quant = np.floor((tables.standard_chrominance * scale + 50) / 100)

        # Ensure values are between 1 and 255
        luminance_quant = np.clip(luminance_quant, 1, 255).astype(np.uint8)
        chrominance_quant = np.clip(chrominance_quant, 1, 255).astype(np.uint8)

        return luminance_quant, chrominance_quant

    #processes 8x8 block of image, applies DCT, quantization, zigzag scan and run-length encoding to encode block
    def process_block(self, block, is_y_channel, channel):
        # Ensure block is 2D (single channel)
        if block.ndim == 3:
            block = block[:, :, 0 if channel == 'y' else (1 if channel == 'cb' else 2)]
        
        block = block.astype(np.float32) - 128  # Center values around 0

        block_dct = cv2.dct(block)  # Apply DCT to block

        quant_table = self.luminance_quant if is_y_channel else self.chrominance_quant
        block_quant = helper.quantize(block_dct, quant_table)

        block_zz = helper.zigzag(block_quant)

        
        dc = block_zz[0] - self.last_dc[channel]

        ac = helper.run_length(block_zz[1:])


        
        block_buffer = helper.Byte_Buffer()
        block_buffer.append_buffer(helper.write_num(-1, dc, self.dc_tables[not is_y_channel]))


        for zero_count, value in zip(ac[::2], ac[1::2]):

            block_buffer.append_buffer(helper.write_num(zero_count, value, self.ac_tables[not is_y_channel]))

        
        self.last_dc[channel] = block_zz[0]
        return block_buffer

    #Encodes image using JPEG compression steps and writes buffers to jpg file
    def encode(self, image, output_file):
        """Encode an image using JPEG compression with improved block splitting."""
        h, w = image.shape[:2] # Get image dimensions


        self.total_blocks = ((h + 7) // 8) * ((w + 7) // 8) 
        self.processed_blocks = 0

        # Pad image if necessary, using this code increses file size but we are using prepare_image to resize image as alternative approach
        # if h % 8 != 0 or w % 8 != 0:
        #     pad_h = (8 - h % 8) % 8
        #     pad_w = (8 - w % 8) % 8
        #     image = np.pad(image, ((0, pad_h), (0, pad_w), (0, 0)), mode='edge')

        with open(output_file, 'wb') as fp:
            #general headers writes common headers of jpeg image.
            self.buffer.append_buffer(helper.general_headers(h, w, self.luminance_quant, self.chrominance_quant))
            self.buffer.flush(fp, header=True)

                    # block splitting 8x8
            blocks = self.block_split(image,h,w)

            for block in blocks:
                # Iterate over channels ('y', 'cb', 'cr') and process each separately
                for k, channel in enumerate(['y', 'cb', 'cr']):
                    # Extract the appropriate channel from the block
                    channel_block = block[:, :, k if k == 0 else 3-k]  # Adjust if channel order is different
                    # Process the block
                    block_buffer = self.process_block(channel_block, k == 0, channel)
                    self.buffer.append_buffer(block_buffer)

                self.buffer.flush(fp)
                
                self.processed_blocks += 1
                self.update_progress()

            if self.buffer.temp:
                self.buffer.append_str('0' * (8 - len(self.buffer.temp)))
                self.buffer.flush(fp)
            
            fp.write(b'\xff\xd9')  # EOI

        self.update_progress(force=True)
    def update_progress(self, force=False):
        if self.progress_callback:
            progress = (self.processed_blocks / self.total_blocks) * 100
            if force or self.processed_blocks % max(1, self.total_blocks // 100) == 0:
                self.progress_callback(progress)