Files
SDR-Rover/tests/lab011_scene_modes.py

525 lines
10 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""
Lab011. Проверка облегчённых режимов изображения
на сцене, похожей на вид с камеры ровера.
Создаются варианты:
1. Grayscale 320x240, JPEG quality 30.
2. Grayscale 320x240, JPEG quality 15.
3. Grayscale 320x240 с наложением контуров, quality 30.
Для каждого варианта рассчитываются:
- размер JPEG;
- количество радиопакетов;
- время передачи при разных скоростях.
"""
from math import ceil
from pathlib import Path
from PIL import (
Image,
ImageChops,
ImageDraw,
ImageFilter,
ImageOps,
)
from protocol.image_fragments import (
encode_image_fragment,
split_image_bytes,
)
from protocol.packet import (
MESSAGE_TYPE_ACK,
MESSAGE_TYPE_IMAGE_FRAGMENT,
build_packet,
)
# ============================================================
# Настройки
# ============================================================
SOURCE_PATH = Path(
"data/raw/lab011_scene.jpg"
)
OUTPUT_DIRECTORY = Path(
"data/processed/lab011"
)
FRAME_SIZE = (320, 240)
FRAGMENT_DATA_SIZE = 512
IMAGE_ID_BASE = 2026071400
BITRATES_KBPS = [
5,
10,
20,
50,
]
# ============================================================
# Вспомогательные функции
# ============================================================
def prepare_frame(
image: Image.Image,
) -> Image.Image:
"""
Привести изображение к фиксированному кадру 320x240.
Пропорции сохраняются.
Недостающие области заполняются чёрным.
"""
image = ImageOps.exif_transpose(image)
image = image.convert("RGB")
return ImageOps.pad(
image,
FRAME_SIZE,
method=Image.Resampling.LANCZOS,
color="black",
centering=(0.5, 0.5),
)
def save_jpeg(
image: Image.Image,
path: Path,
quality: int,
) -> None:
"""
Сохранить изображение в JPEG.
"""
path.parent.mkdir(
parents=True,
exist_ok=True,
)
image.save(
path,
format="JPEG",
quality=quality,
optimize=True,
)
def format_size(
byte_count: int,
) -> str:
"""
Представить размер в КиБ.
"""
return f"{byte_count / 1024:.2f} КиБ"
def format_duration(
seconds: float,
) -> str:
"""
Представить длительность передачи.
"""
if seconds < 1:
return f"{seconds * 1000:.0f} мс"
if seconds < 60:
return f"{seconds:.2f} с"
minutes = int(seconds // 60)
remaining_seconds = seconds % 60
return (
f"{minutes} мин "
f"{remaining_seconds:.1f} с"
)
def estimate_transfer(
file_path: Path,
image_id: int,
) -> dict:
"""
Рассчитать полный объём DATA + ACK.
"""
image_bytes = file_path.read_bytes()
fragments = split_image_bytes(
image_bytes=image_bytes,
image_id=image_id,
fragment_data_size=FRAGMENT_DATA_SIZE,
)
data_packet_bytes = 0
for fragment in fragments:
fragment_payload = encode_image_fragment(
fragment
)
packet = build_packet(
payload=fragment_payload,
message_type=MESSAGE_TYPE_IMAGE_FRAGMENT,
sequence_number=fragment.fragment_index,
)
data_packet_bytes += len(packet)
ack_packet = build_packet(
payload=b"",
message_type=MESSAGE_TYPE_ACK,
sequence_number=0,
)
ack_bytes = (
len(fragments)
* len(ack_packet)
)
total_radio_bytes = (
data_packet_bytes
+ ack_bytes
)
times = {}
for bitrate_kbps in BITRATES_KBPS:
times[bitrate_kbps] = (
total_radio_bytes
* 8
/ (bitrate_kbps * 1000)
)
return {
"file_size": len(image_bytes),
"fragment_count": len(fragments),
"total_radio_bytes": total_radio_bytes,
"times": times,
}
# ============================================================
# Проверка исходного файла
# ============================================================
if not SOURCE_PATH.exists():
raise FileNotFoundError(
f"Не найден файл: {SOURCE_PATH}"
)
OUTPUT_DIRECTORY.mkdir(
parents=True,
exist_ok=True,
)
# ============================================================
# Подготовка базового кадра
# ============================================================
with Image.open(SOURCE_PATH) as source_image:
original_size = source_image.size
color_frame = prepare_frame(
source_image
)
gray_frame = ImageOps.grayscale(
color_frame
)
# ============================================================
# Вариант 1. Grayscale, quality 30
# ============================================================
gray_q30_path = (
OUTPUT_DIRECTORY
/ "01_gray_320_q30.jpg"
)
save_jpeg(
gray_frame,
gray_q30_path,
quality=30,
)
# ============================================================
# Вариант 2. Grayscale, quality 15
# ============================================================
gray_q15_path = (
OUTPUT_DIRECTORY
/ "02_gray_320_q15.jpg"
)
save_jpeg(
gray_frame,
gray_q15_path,
quality=15,
)
# ============================================================
# Вариант 3. Grayscale + контуры
# ============================================================
edge_map = gray_frame.filter(
ImageFilter.FIND_EDGES
)
edge_map = ImageOps.autocontrast(
edge_map
)
# Оставляем преимущественно сильные контуры.
binary_edges = edge_map.point(
lambda value: 255 if value >= 45 else 0
)
# После инверсии контуры становятся чёрными,
# а фон — белым.
dark_edges = ImageOps.invert(
binary_edges
)
# Сохраняем исходный серый фон и добавляем
# поверх него тёмные линии контуров.
gray_with_edges = ImageChops.darker(
gray_frame,
dark_edges
)
gray_edges_path = (
OUTPUT_DIRECTORY
/ "03_gray_edges_320_q30.jpg"
)
save_jpeg(
gray_with_edges,
gray_edges_path,
quality=30,
)
# ============================================================
# Список вариантов
# ============================================================
variants = [
(
"gray_320_q30",
gray_q30_path,
),
(
"gray_320_q15",
gray_q15_path,
),
(
"gray_edges_320_q30",
gray_edges_path,
),
]
# ============================================================
# Расчёт параметров
# ============================================================
results = []
for variant_index, (
variant_name,
variant_path,
) in enumerate(variants):
result = estimate_transfer(
file_path=variant_path,
image_id=IMAGE_ID_BASE + variant_index,
)
results.append(
{
"name": variant_name,
"path": variant_path,
**result,
}
)
# ============================================================
# Создание сравнительной картинки
# ============================================================
preview_variants = [
(
"SOURCE PREVIEW",
color_frame,
None,
)
]
for result in results:
with Image.open(result["path"]) as image:
preview_variants.append(
(
result["name"],
image.convert("RGB").copy(),
result,
)
)
columns = 2
cell_width = 440
cell_height = 340
rows = ceil(
len(preview_variants)
/ columns
)
comparison_image = Image.new(
"RGB",
(
columns * cell_width,
rows * cell_height,
),
"white",
)
draw = ImageDraw.Draw(
comparison_image
)
for index, (
label,
preview,
result,
) in enumerate(preview_variants):
column = index % columns
row = index // columns
x = column * cell_width
y = row * cell_height
preview = preview.copy()
preview.thumbnail(
(400, 270),
Image.Resampling.NEAREST,
)
paste_x = (
x
+ (cell_width - preview.width) // 2
)
comparison_image.paste(
preview,
(paste_x, y + 55),
)
if result is None:
text = (
f"{label}\n"
f"original: "
f"{original_size[0]}x{original_size[1]}"
)
else:
text = (
f"{label}\n"
f"{format_size(result['file_size'])}, "
f"{result['fragment_count']} fragments, "
f"{format_duration(result['times'][20])} "
f"at 20 kbps"
)
draw.multiline_text(
(x + 10, y + 10),
text,
fill="black",
spacing=4,
)
comparison_path = (
OUTPUT_DIRECTORY
/ "lab011_comparison.jpg"
)
comparison_image.save(
comparison_path,
format="JPEG",
quality=90,
)
# ============================================================
# Вывод результатов
# ============================================================
print(
"=== Lab011. Режимы кадра для ровера ==="
)
print("\nИсходное разрешение:")
print(
original_size[0],
"x",
original_size[1],
)
print("\nРезультаты:")
print(
f"{'Вариант':<27}"
f"{'Размер':>12}"
f"{'Пакеты':>10}"
f"{'10 кбит/с':>13}"
f"{'20 кбит/с':>13}"
f"{'50 кбит/с':>13}"
)
print("-" * 88)
for result in results:
print(
f"{result['name']:<27}"
f"{format_size(result['file_size']):>12}"
f"{result['fragment_count']:>10}"
f"{format_duration(result['times'][10]):>13}"
f"{format_duration(result['times'][20]):>13}"
f"{format_duration(result['times'][50]):>13}"
)
print("\nСравнительная картинка:")
print(comparison_path)
print(
"\nПроверка пройдена: "
"режимы кадра созданы и рассчитаны."
)