fix: cosine matrix

dct.py

@@ -6,6 +6,7 @@ from PIL import Image
 
 Img = npt.NDArray[np.uint8]
 
+
 def load_img(path: str) -> Img:
     img = Image.open(path).convert("L")
 
@@ -36,11 +37,10 @@ def cos_mat(n: int) -> Img:
 
     for i in range(n):
         for j in range(n):
-
             if i == 0:
-                mat[i, j] = 1 / n
+                mat[i, j] = 1 / np.sqrt(n)
             else:
-                mat[i, j] = np.sqrt(2 / n) *  np.cos((2 * j + 1) * i * np.pi \
+                mat[i, j] = np.sqrt(2 / n) * np.cos((2 * j + 1) * i * np.pi
                                                     / (2 * n))
 
     return mat
@@ -59,11 +59,10 @@ def dct(padded_img: Img, kernel_size: int) -> Img:
 
     return final_img
 
+
 def idct(transformed_img: Img, kernel_size: int) -> Img:
     c_mat = cos_mat(kernel_size)
-    c_t = np.transpose(c_mat)
 
-    # Compute the inverse DCT matrix
     inv_c_mat = np.linalg.inv(c_mat)
     inv_c_t = np.transpose(inv_c_mat)
 
@@ -72,10 +71,12 @@ def idct(transformed_img: Img, kernel_size: int) -> Img:
     for i in range(0, transformed_img.shape[0], kernel_size):
         for j in range(0, transformed_img.shape[1], kernel_size):
             sub_mat = sliding_window(transformed_img, i, j, kernel_size)
-            final_img[i:i+kernel_size, j:j+kernel_size] = inv_c_mat @ sub_mat @ inv_c_t
+            final_img[i:i+kernel_size, j:j +
+                      kernel_size] = inv_c_mat @ sub_mat @ inv_c_t
 
     return final_img
 
+
 def quantized(kernel_size: int, quality: int) -> Img:
     quant = np.zeros((kernel_size, kernel_size))
 
@@ -92,7 +93,8 @@ def quantize_dct(dct_img: Img, kernel_size: int, quality: int) -> Img:
 
     for i in range(0, dct_img.shape[0], kernel_size):
         for j in range(0, dct_img.shape[1], kernel_size):
-            final_mat[i:i+kernel_size, j:j+kernel_size] = np.round(dct_img[i:i+kernel_size, j:j+kernel_size] / quant)
+            final_mat[i:i+kernel_size, j:j+kernel_size] = np.round(
+                dct_img[i:i+kernel_size, j:j+kernel_size] / quant)
 
     return final_mat
 
@@ -104,16 +106,23 @@ def pad_image(img: Img, kernel_size: int):
     pad_rows = (kernel_size - rows % kernel_size) % kernel_size
     pad_cols = (kernel_size - cols % kernel_size) % kernel_size
 
-    return np.pad(img, ((0, pad_rows), (0, pad_cols)), mode='constant', constant_values=0)
+    return np.pad(img, ((0, pad_rows), (0, pad_cols)), mode='constant',
+                  constant_values=0)
+
 
 def plot_images(images, titles):
-    fig, axes = plt.subplots(1, 4, figsize=(12, 3))  # 1 row, 4 columns
+    fig, axes = plt.subplots(2, 2, figsize=(12, 9))
+    axes = axes.flatten()
+
     for ax, img, title in zip(axes, images, titles):
         ax.imshow(img, cmap='gray')
         ax.set_title(title)
         ax.axis('off')
+
+    plt.tight_layout()
     plt.show()
 
+
 def main():
     parser = argparse.ArgumentParser(prog='DCT converter')
 
@@ -130,15 +139,14 @@ def main():
     dct_img = dct(padded, kernel_size)
     inv_dct = idct(dct_img, kernel_size)
 
-    q_dct = quantize_dct(dct_img, kernel_size, 2)
+    q_dct = quantize_dct(dct_img, kernel_size, 3)
     inv_q_dct = idct(q_dct, kernel_size)
 
     images = [dct_img, inv_dct, q_dct, inv_q_dct]
-    titles = ["DCT Image", "Inverse DCT", "Quantized DCT", "Inverse Quantized DCT"]
+    titles = ["DCT", "Inverse DCT", "Quantized DCT", "Inverse Quantized DCT"]
 
     plot_images(images, titles)
 
 
-
 if __name__ == "__main__":
     main()