main, lsb и dct под все типы сообщений и картинок

core/dct.py

@@ -1,5 +1,6 @@
 """
 Модуль DCT стеганографии с квантованием (JPEG-style)
+Оптимизирован для длинных сообщений
 """
 
 import numpy as np
@@ -9,16 +10,16 @@ from .utils import text_to_bits, bits_to_text
 
 BLOCK_SIZE = 8
 
-# Стандартная таблица квантования JPEG для качества 50%
+# Более мягкая таблица квантования (качество ~70-80%)
 QUANT_TABLE = np.array([
-    [16, 11, 10, 16, 24, 40, 51, 61],
-    [12, 12, 14, 19, 26, 58, 60, 55],
-    [14, 13, 16, 24, 40, 57, 69, 56],
-    [14, 17, 22, 29, 51, 87, 80, 62],
-    [18, 22, 37, 56, 68, 109, 103, 77],
-    [24, 35, 55, 64, 81, 104, 113, 92],
-    [49, 64, 78, 87, 103, 121, 120, 101],
-    [72, 92, 95, 98, 112, 100, 103, 99]
+    [8, 6, 5, 8, 12, 20, 26, 31],
+    [6, 6, 7, 10, 13, 29, 30, 28],
+    [7, 7, 8, 12, 20, 29, 35, 28],
+    [7, 9, 11, 15, 26, 44, 40, 31],
+    [9, 11, 19, 28, 34, 55, 52, 39],
+    [12, 18, 28, 32, 41, 52, 57, 46],
+    [25, 32, 39, 44, 52, 61, 60, 51],
+    [36, 46, 48, 49, 56, 50, 52, 50]
 ])
 
 
@@ -40,7 +41,8 @@ def _get_zigzag_order() -> list:
 
 
 ZIGZAG_ORDER = _get_zigzag_order()
-MID_FREQ_INDICES = ZIGZAG_ORDER[1:40]
+# Используем больше коэффициентов (1-50)
+MID_FREQ_INDICES = ZIGZAG_ORDER[1:50]
 
 
 def _dct_quantize(block: np.ndarray) -> np.ndarray:
@@ -93,28 +95,21 @@ def _embed_bits_in_block(block_quant: np.ndarray, bits: list, bit_index: int) ->
 
 def encode_dct(image_path: str, message: str, output_path: str) -> bool:
     """Скрывает сообщение в изображении методом DCT."""
-    img = Image.open(image_path).convert('L')
+    img = Image.open(image_path).convert('RGB')
     pixels = np.array(img, dtype=np.float64)
-    height, width = pixels.shape
+    height, width, channels = pixels.shape
 
-    # Кодируем сообщение в байты
-    message_bytes = message.encode('utf-8')
-    msg_length = len(message_bytes)
-    
-    # Формат: [длина: 32 бита] + [сообщение]
-    # 32 бита = 4 байта, достаточно для сообщений до 4 ГБ
+    msg_bytes = message.encode('utf-8')
+    msg_length = len(msg_bytes)
     length_bits = format(msg_length, '032b')
     message_bits = text_to_bits(message)
-    
-    # Собираем всё вместе
     all_bits = length_bits + message_bits
     bit_list = [int(b) for b in all_bits]
     total_bits = len(bit_list)
 
-    # Проверяем вместимость
     blocks_per_row = width // BLOCK_SIZE
     blocks_per_col = height // BLOCK_SIZE
-    max_bits = blocks_per_row * blocks_per_col * len(MID_FREQ_INDICES)
+    max_bits = blocks_per_row * blocks_per_col * len(MID_FREQ_INDICES) * 3
 
     if total_bits > max_bits:
         print(f"Ошибка: сообщение слишком большое.")
@@ -129,16 +124,20 @@ def encode_dct(image_path: str, message: str, output_path: str) -> bool:
             if bit_index >= total_bits:
                 break
             
-            block = pixels[i:i+BLOCK_SIZE, j:j+BLOCK_SIZE]
-            quant_block = _dct_quantize(block)
-            quant_block, bit_index = _embed_bits_in_block(quant_block, bit_list, bit_index)
-            new_block = _idct_dequantize(quant_block)
-            modified_pixels[i:i+BLOCK_SIZE, j:j+BLOCK_SIZE] = new_block
+            for c in range(3):
+                if bit_index >= total_bits:
+                    break
+                    
+                block = pixels[i:i+BLOCK_SIZE, j:j+BLOCK_SIZE, c]
+                quant_block = _dct_quantize(block)
+                quant_block, bit_index = _embed_bits_in_block(quant_block, bit_list, bit_index)
+                new_channel = _idct_dequantize(quant_block)
+                modified_pixels[i:i+BLOCK_SIZE, j:j+BLOCK_SIZE, c] = new_channel
             
         if bit_index >= total_bits:
             break
 
-    result_img = Image.fromarray(modified_pixels.astype(np.uint8), mode='L')
+    result_img = Image.fromarray(modified_pixels.astype(np.uint8), mode='RGB')
     result_img.save(output_path)
     print(f"Успешно! Спрятано {total_bits} бит ({total_bits // 8} байт)")
     return True
@@ -146,21 +145,20 @@ def encode_dct(image_path: str, message: str, output_path: str) -> bool:
 
 def decode_dct(image_path: str) -> str:
     """Извлекает сообщение из изображения методом DCT."""
-    img = Image.open(image_path).convert('L')
+    img = Image.open(image_path).convert('RGB')
     pixels = np.array(img, dtype=np.float64)
-    height, width = pixels.shape
+    height, width, channels = pixels.shape
 
-    # Извлекаем все биты
     all_bits = []
     
     for i in range(0, height - BLOCK_SIZE + 1, BLOCK_SIZE):
         for j in range(0, width - BLOCK_SIZE + 1, BLOCK_SIZE):
-            block = pixels[i:i+BLOCK_SIZE, j:j+BLOCK_SIZE]
-            quant_block = _dct_quantize(block)
-            bits_from_block = _extract_bits_from_block(quant_block, len(MID_FREQ_INDICES))
-            all_bits.extend(bits_from_block)
+            for c in range(3):
+                channel_block = pixels[i:i+BLOCK_SIZE, j:j+BLOCK_SIZE, c]
+                quant_block = _dct_quantize(channel_block)
+                bits_from_block = _extract_bits_from_block(quant_block, len(MID_FREQ_INDICES))
+                all_bits.extend(bits_from_block)
 
-    # Читаем длину сообщения (первые 32 бита)
     if len(all_bits) < 32:
         return ""
     
@@ -168,14 +166,11 @@ def decode_dct(image_path: str) -> str:
     length_str = ''.join(str(b) for b in length_bits)
     msg_length = int(length_str, 2)
     
-    # Проверяем, что длина разумная
-    if msg_length > len(all_bits) // 8:
+    if msg_length > (len(all_bits) - 32) // 8:
         return ""
     
-    # Читаем сообщение
     message_bits = all_bits[32:32 + msg_length * 8]
     
-    # Дополняем до кратности 8 (хотя уже должно быть кратно)
     remainder = len(message_bits) % 8
     if remainder != 0:
         message_bits.extend([0] * (8 - remainder))