Repair project layout and fix Lab024B spectrum detection
This commit is contained in:
LittleSam129
354
protocol/image_fragments.py
Normal file
354
protocol/image_fragments.py
Normal file
@@ -0,0 +1,354 @@
|
||||
"""
|
||||
Фрагментация и сборка изображений SDR Rover Link.
|
||||
|
||||
Изображение рассматривается как обычная последовательность байтов.
|
||||
Модулю неважно, используется JPEG, PNG или другой формат.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
import struct
|
||||
from typing import Iterable
|
||||
|
||||
from protocol.packet import MAX_PAYLOAD_SIZE
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Формат заголовка фрагмента изображения
|
||||
# ============================================================
|
||||
|
||||
# I — идентификатор изображения, 4 байта
|
||||
# H — номер фрагмента, 2 байта
|
||||
# H — общее количество фрагментов, 2 байта
|
||||
# I — размер исходного изображения, 4 байта
|
||||
FRAGMENT_HEADER_FORMAT = ">IHHI"
|
||||
|
||||
FRAGMENT_HEADER_SIZE = struct.calcsize(
|
||||
FRAGMENT_HEADER_FORMAT
|
||||
)
|
||||
|
||||
MAX_FRAGMENT_DATA_SIZE = (
|
||||
MAX_PAYLOAD_SIZE - FRAGMENT_HEADER_SIZE
|
||||
)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Ошибки
|
||||
# ============================================================
|
||||
|
||||
class ImageFragmentError(ValueError):
|
||||
"""Ошибка структуры фрагмента изображения."""
|
||||
|
||||
|
||||
class MissingFragmentsError(ImageFragmentError):
|
||||
"""Отсутствуют необходимые фрагменты изображения."""
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Представление одного фрагмента
|
||||
# ============================================================
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ImageFragment:
|
||||
"""
|
||||
Один фрагмент изображения.
|
||||
|
||||
Атрибуты
|
||||
--------
|
||||
image_id:
|
||||
Идентификатор изображения.
|
||||
|
||||
fragment_index:
|
||||
Номер фрагмента. Нумерация начинается с нуля.
|
||||
|
||||
fragment_count:
|
||||
Общее количество фрагментов.
|
||||
|
||||
original_size:
|
||||
Размер исходного изображения в байтах.
|
||||
|
||||
data:
|
||||
Часть исходного изображения.
|
||||
"""
|
||||
|
||||
image_id: int
|
||||
fragment_index: int
|
||||
fragment_count: int
|
||||
original_size: int
|
||||
data: bytes
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Разбиение изображения
|
||||
# ============================================================
|
||||
|
||||
def split_image_bytes(
|
||||
image_bytes: bytes,
|
||||
image_id: int,
|
||||
fragment_data_size: int = 512,
|
||||
) -> list[ImageFragment]:
|
||||
"""
|
||||
Разбить изображение на фрагменты.
|
||||
|
||||
Параметры
|
||||
--------
|
||||
image_bytes:
|
||||
Полное содержимое файла изображения.
|
||||
|
||||
image_id:
|
||||
Идентификатор изображения от 0 до 4294967295.
|
||||
|
||||
fragment_data_size:
|
||||
Максимальный размер данных одного фрагмента.
|
||||
|
||||
Возвращает
|
||||
----------
|
||||
Список объектов ImageFragment.
|
||||
"""
|
||||
|
||||
if not isinstance(image_bytes, (bytes, bytearray)):
|
||||
raise TypeError(
|
||||
"image_bytes должен иметь тип bytes или bytearray"
|
||||
)
|
||||
|
||||
image_bytes = bytes(image_bytes)
|
||||
|
||||
if not image_bytes:
|
||||
raise ValueError(
|
||||
"Изображение не должно быть пустым"
|
||||
)
|
||||
|
||||
if not 0 <= image_id <= 0xFFFFFFFF:
|
||||
raise ValueError(
|
||||
"image_id должен находиться "
|
||||
"в диапазоне 0...4294967295"
|
||||
)
|
||||
|
||||
if not 1 <= fragment_data_size <= MAX_FRAGMENT_DATA_SIZE:
|
||||
raise ValueError(
|
||||
"Некорректный размер фрагмента"
|
||||
)
|
||||
|
||||
fragment_count = (
|
||||
len(image_bytes)
|
||||
+ fragment_data_size
|
||||
- 1
|
||||
) // fragment_data_size
|
||||
|
||||
if fragment_count > 65535:
|
||||
raise ValueError(
|
||||
"Изображение требует более 65535 фрагментов"
|
||||
)
|
||||
|
||||
fragments = []
|
||||
|
||||
for fragment_index in range(fragment_count):
|
||||
|
||||
start = fragment_index * fragment_data_size
|
||||
end = start + fragment_data_size
|
||||
|
||||
fragment_data = image_bytes[start:end]
|
||||
|
||||
fragments.append(
|
||||
ImageFragment(
|
||||
image_id=image_id,
|
||||
fragment_index=fragment_index,
|
||||
fragment_count=fragment_count,
|
||||
original_size=len(image_bytes),
|
||||
data=fragment_data,
|
||||
)
|
||||
)
|
||||
|
||||
return fragments
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Преобразование фрагмента в байты
|
||||
# ============================================================
|
||||
|
||||
def encode_image_fragment(
|
||||
fragment: ImageFragment,
|
||||
) -> bytes:
|
||||
"""
|
||||
Преобразовать ImageFragment в полезную нагрузку пакета.
|
||||
"""
|
||||
|
||||
if not isinstance(fragment, ImageFragment):
|
||||
raise TypeError(
|
||||
"fragment должен иметь тип ImageFragment"
|
||||
)
|
||||
|
||||
if not 0 <= fragment.fragment_index < fragment.fragment_count:
|
||||
raise ImageFragmentError(
|
||||
"Некорректный номер фрагмента"
|
||||
)
|
||||
|
||||
if not fragment.data:
|
||||
raise ImageFragmentError(
|
||||
"Фрагмент не должен быть пустым"
|
||||
)
|
||||
|
||||
if len(fragment.data) > MAX_FRAGMENT_DATA_SIZE:
|
||||
raise ImageFragmentError(
|
||||
"Данные фрагмента слишком велики"
|
||||
)
|
||||
|
||||
header = struct.pack(
|
||||
FRAGMENT_HEADER_FORMAT,
|
||||
fragment.image_id,
|
||||
fragment.fragment_index,
|
||||
fragment.fragment_count,
|
||||
fragment.original_size,
|
||||
)
|
||||
|
||||
return header + fragment.data
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Преобразование байтов обратно во фрагмент
|
||||
# ============================================================
|
||||
|
||||
def decode_image_fragment(
|
||||
payload: bytes,
|
||||
) -> ImageFragment:
|
||||
"""
|
||||
Разобрать полезную нагрузку пакета изображения.
|
||||
"""
|
||||
|
||||
if not isinstance(payload, (bytes, bytearray)):
|
||||
raise TypeError(
|
||||
"payload должен иметь тип bytes или bytearray"
|
||||
)
|
||||
|
||||
payload = bytes(payload)
|
||||
|
||||
if len(payload) <= FRAGMENT_HEADER_SIZE:
|
||||
raise ImageFragmentError(
|
||||
"Фрагмент слишком короткий"
|
||||
)
|
||||
|
||||
header = payload[:FRAGMENT_HEADER_SIZE]
|
||||
fragment_data = payload[FRAGMENT_HEADER_SIZE:]
|
||||
|
||||
(
|
||||
image_id,
|
||||
fragment_index,
|
||||
fragment_count,
|
||||
original_size,
|
||||
) = struct.unpack(
|
||||
FRAGMENT_HEADER_FORMAT,
|
||||
header,
|
||||
)
|
||||
|
||||
if fragment_count == 0:
|
||||
raise ImageFragmentError(
|
||||
"Количество фрагментов не может быть равно нулю"
|
||||
)
|
||||
|
||||
if fragment_index >= fragment_count:
|
||||
raise ImageFragmentError(
|
||||
"Номер фрагмента выходит за допустимый диапазон"
|
||||
)
|
||||
|
||||
if original_size == 0:
|
||||
raise ImageFragmentError(
|
||||
"Размер исходного изображения равен нулю"
|
||||
)
|
||||
|
||||
return ImageFragment(
|
||||
image_id=image_id,
|
||||
fragment_index=fragment_index,
|
||||
fragment_count=fragment_count,
|
||||
original_size=original_size,
|
||||
data=fragment_data,
|
||||
)
|
||||
|
||||
|
||||
# ============================================================
|
||||
# Сборка изображения
|
||||
# ============================================================
|
||||
|
||||
def reassemble_image(
|
||||
fragments: Iterable[ImageFragment],
|
||||
) -> bytes:
|
||||
"""
|
||||
Собрать исходное изображение из фрагментов.
|
||||
|
||||
Фрагменты могут поступать в произвольном порядке.
|
||||
"""
|
||||
|
||||
fragments = list(fragments)
|
||||
|
||||
if not fragments:
|
||||
raise MissingFragmentsError(
|
||||
"Список фрагментов пуст"
|
||||
)
|
||||
|
||||
first_fragment = fragments[0]
|
||||
|
||||
image_id = first_fragment.image_id
|
||||
fragment_count = first_fragment.fragment_count
|
||||
original_size = first_fragment.original_size
|
||||
|
||||
fragments_by_index = {}
|
||||
|
||||
for fragment in fragments:
|
||||
|
||||
if fragment.image_id != image_id:
|
||||
raise ImageFragmentError(
|
||||
"Фрагменты относятся к разным изображениям"
|
||||
)
|
||||
|
||||
if fragment.fragment_count != fragment_count:
|
||||
raise ImageFragmentError(
|
||||
"Различается общее количество фрагментов"
|
||||
)
|
||||
|
||||
if fragment.original_size != original_size:
|
||||
raise ImageFragmentError(
|
||||
"Различается размер исходного изображения"
|
||||
)
|
||||
|
||||
existing_fragment = fragments_by_index.get(
|
||||
fragment.fragment_index
|
||||
)
|
||||
|
||||
if existing_fragment is not None:
|
||||
|
||||
if existing_fragment.data != fragment.data:
|
||||
raise ImageFragmentError(
|
||||
"Получены разные данные "
|
||||
"для одного номера фрагмента"
|
||||
)
|
||||
|
||||
# Полностью одинаковый повторный фрагмент
|
||||
# можно безопасно проигнорировать.
|
||||
continue
|
||||
|
||||
fragments_by_index[
|
||||
fragment.fragment_index
|
||||
] = fragment
|
||||
|
||||
missing_indexes = [
|
||||
index
|
||||
for index in range(fragment_count)
|
||||
if index not in fragments_by_index
|
||||
]
|
||||
|
||||
if missing_indexes:
|
||||
raise MissingFragmentsError(
|
||||
"Отсутствуют фрагменты: "
|
||||
+ ", ".join(map(str, missing_indexes))
|
||||
)
|
||||
|
||||
assembled_bytes = b"".join(
|
||||
fragments_by_index[index].data
|
||||
for index in range(fragment_count)
|
||||
)
|
||||
|
||||
if len(assembled_bytes) != original_size:
|
||||
raise ImageFragmentError(
|
||||
"Размер собранного изображения "
|
||||
"не совпадает с исходным размером"
|
||||
)
|
||||
|
||||
return assembled_bytes
|
||||
Reference in New Issue
Block a user