786 lines
16 KiB
Python
786 lines
16 KiB
Python
"""
|
||
Lab010. Оптимизация изображения для слабого радиоканала.
|
||
|
||
Программа формирует несколько вариантов одного кадра:
|
||
|
||
1. Исходный JPEG.
|
||
2. Цветной кадр до 640x480.
|
||
3. Серый кадр до 640x480.
|
||
4. Серый кадр до 320x240, JPEG quality 30.
|
||
5. Серый кадр до 320x240, JPEG quality 15.
|
||
6. Карта контуров 320x240.
|
||
7. Серый кадр с наложенными контурами.
|
||
|
||
Для каждого варианта рассчитываются:
|
||
|
||
- размер JPEG;
|
||
- число фрагментов;
|
||
- полный объём DATA + ACK;
|
||
- эффективность протокола;
|
||
- время передачи на нескольких скоростях;
|
||
- время при удвоении трафика из-за повторов.
|
||
"""
|
||
|
||
from csv import DictWriter
|
||
from hashlib import sha256
|
||
from math import ceil
|
||
from pathlib import Path
|
||
|
||
from PIL import (
|
||
Image,
|
||
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/lab009_source.jpg"
|
||
)
|
||
|
||
OUTPUT_DIRECTORY = Path(
|
||
"data/processed/lab010"
|
||
)
|
||
|
||
FRAGMENT_DATA_SIZE = 512
|
||
|
||
IMAGE_ID_BASE = 2026071300
|
||
|
||
BITRATES_KBPS = [
|
||
5,
|
||
10,
|
||
20,
|
||
50,
|
||
100,
|
||
]
|
||
|
||
|
||
# ============================================================
|
||
# Вспомогательные функции
|
||
# ============================================================
|
||
|
||
def resize_inside(
|
||
image: Image.Image,
|
||
maximum_size: tuple[int, int],
|
||
) -> Image.Image:
|
||
"""
|
||
Уменьшить изображение с сохранением пропорций.
|
||
|
||
Изображение не растягивается и не искажается.
|
||
"""
|
||
|
||
resized_image = image.copy()
|
||
|
||
resized_image.thumbnail(
|
||
maximum_size,
|
||
Image.Resampling.LANCZOS,
|
||
)
|
||
|
||
return resized_image
|
||
|
||
|
||
def save_jpeg(
|
||
image: Image.Image,
|
||
output_path: Path,
|
||
quality: int,
|
||
) -> None:
|
||
"""
|
||
Сохранить изображение в JPEG.
|
||
"""
|
||
|
||
output_path.parent.mkdir(
|
||
parents=True,
|
||
exist_ok=True,
|
||
)
|
||
|
||
image.save(
|
||
output_path,
|
||
format="JPEG",
|
||
quality=quality,
|
||
optimize=True,
|
||
)
|
||
|
||
|
||
def format_bytes(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:
|
||
"""
|
||
Рассчитать параметры передачи одного JPEG.
|
||
"""
|
||
|
||
image_bytes = file_path.read_bytes()
|
||
|
||
fragments = split_image_bytes(
|
||
image_bytes=image_bytes,
|
||
image_id=image_id,
|
||
fragment_data_size=FRAGMENT_DATA_SIZE,
|
||
)
|
||
|
||
total_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,
|
||
)
|
||
|
||
total_data_packet_bytes += len(packet)
|
||
|
||
ack_packet = build_packet(
|
||
payload=b"",
|
||
message_type=MESSAGE_TYPE_ACK,
|
||
sequence_number=0,
|
||
)
|
||
|
||
total_ack_bytes = (
|
||
len(fragments)
|
||
* len(ack_packet)
|
||
)
|
||
|
||
total_radio_bytes = (
|
||
total_data_packet_bytes
|
||
+ total_ack_bytes
|
||
)
|
||
|
||
efficiency_percent = (
|
||
len(image_bytes)
|
||
/ total_radio_bytes
|
||
* 100
|
||
)
|
||
|
||
transfer_times = {}
|
||
|
||
for bitrate_kbps in BITRATES_KBPS:
|
||
|
||
bitrate_bits_per_second = (
|
||
bitrate_kbps * 1000
|
||
)
|
||
|
||
transfer_times[bitrate_kbps] = (
|
||
total_radio_bytes
|
||
* 8
|
||
/ bitrate_bits_per_second
|
||
)
|
||
|
||
return {
|
||
"file_size_bytes": len(image_bytes),
|
||
"fragment_count": len(fragments),
|
||
"data_packet_bytes": total_data_packet_bytes,
|
||
"ack_bytes": total_ack_bytes,
|
||
"total_radio_bytes": total_radio_bytes,
|
||
"efficiency_percent": efficiency_percent,
|
||
"transfer_times": transfer_times,
|
||
"sha256": sha256(image_bytes).hexdigest(),
|
||
}
|
||
|
||
|
||
# ============================================================
|
||
# Проверка исходной фотографии
|
||
# ============================================================
|
||
|
||
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:
|
||
|
||
source_image.load()
|
||
|
||
source_rgb = source_image.convert("RGB")
|
||
|
||
original_width, original_height = (
|
||
source_rgb.size
|
||
)
|
||
|
||
|
||
# ============================================================
|
||
# Формирование вариантов изображения
|
||
# ============================================================
|
||
|
||
variant_paths = []
|
||
|
||
# ------------------------------------------------------------
|
||
# Вариант 0. Исходный JPEG без изменения
|
||
# ------------------------------------------------------------
|
||
|
||
variant_paths.append(
|
||
(
|
||
"00_original",
|
||
SOURCE_PATH,
|
||
)
|
||
)
|
||
|
||
|
||
# ------------------------------------------------------------
|
||
# Вариант 1. Цветной кадр до 640x480, quality 40
|
||
# ------------------------------------------------------------
|
||
|
||
color_640 = resize_inside(
|
||
source_rgb,
|
||
(640, 480),
|
||
)
|
||
|
||
color_640_path = (
|
||
OUTPUT_DIRECTORY
|
||
/ "01_color_640_q40.jpg"
|
||
)
|
||
|
||
save_jpeg(
|
||
color_640,
|
||
color_640_path,
|
||
quality=40,
|
||
)
|
||
|
||
variant_paths.append(
|
||
(
|
||
"01_color_640_q40",
|
||
color_640_path,
|
||
)
|
||
)
|
||
|
||
|
||
# ------------------------------------------------------------
|
||
# Вариант 2. Серый кадр до 640x480, quality 40
|
||
# ------------------------------------------------------------
|
||
|
||
gray_640 = ImageOps.grayscale(
|
||
color_640
|
||
)
|
||
|
||
gray_640_path = (
|
||
OUTPUT_DIRECTORY
|
||
/ "02_gray_640_q40.jpg"
|
||
)
|
||
|
||
save_jpeg(
|
||
gray_640,
|
||
gray_640_path,
|
||
quality=40,
|
||
)
|
||
|
||
variant_paths.append(
|
||
(
|
||
"02_gray_640_q40",
|
||
gray_640_path,
|
||
)
|
||
)
|
||
|
||
|
||
# ------------------------------------------------------------
|
||
# Вариант 3. Серый кадр до 320x240, quality 30
|
||
# ------------------------------------------------------------
|
||
|
||
color_320 = resize_inside(
|
||
source_rgb,
|
||
(320, 240),
|
||
)
|
||
|
||
gray_320 = ImageOps.grayscale(
|
||
color_320
|
||
)
|
||
|
||
gray_320_q30_path = (
|
||
OUTPUT_DIRECTORY
|
||
/ "03_gray_320_q30.jpg"
|
||
)
|
||
|
||
save_jpeg(
|
||
gray_320,
|
||
gray_320_q30_path,
|
||
quality=30,
|
||
)
|
||
|
||
variant_paths.append(
|
||
(
|
||
"03_gray_320_q30",
|
||
gray_320_q30_path,
|
||
)
|
||
)
|
||
|
||
|
||
# ------------------------------------------------------------
|
||
# Вариант 4. Серый кадр до 320x240, quality 15
|
||
# ------------------------------------------------------------
|
||
|
||
gray_320_q15_path = (
|
||
OUTPUT_DIRECTORY
|
||
/ "04_gray_320_q15.jpg"
|
||
)
|
||
|
||
save_jpeg(
|
||
gray_320,
|
||
gray_320_q15_path,
|
||
quality=15,
|
||
)
|
||
|
||
variant_paths.append(
|
||
(
|
||
"04_gray_320_q15",
|
||
gray_320_q15_path,
|
||
)
|
||
)
|
||
|
||
|
||
# ------------------------------------------------------------
|
||
# Вариант 5. Только контуры
|
||
# ------------------------------------------------------------
|
||
|
||
edge_map = gray_320.filter(
|
||
ImageFilter.FIND_EDGES
|
||
)
|
||
|
||
edge_map = ImageOps.autocontrast(
|
||
edge_map
|
||
)
|
||
|
||
# Инвертируем: белый фон, тёмные контуры.
|
||
edge_map = ImageOps.invert(
|
||
edge_map
|
||
)
|
||
|
||
edges_path = (
|
||
OUTPUT_DIRECTORY
|
||
/ "05_edges_320_q40.jpg"
|
||
)
|
||
|
||
save_jpeg(
|
||
edge_map,
|
||
edges_path,
|
||
quality=40,
|
||
)
|
||
|
||
variant_paths.append(
|
||
(
|
||
"05_edges_320_q40",
|
||
edges_path,
|
||
)
|
||
)
|
||
|
||
|
||
# ------------------------------------------------------------
|
||
# Вариант 6. Серый кадр с усиленными контурами
|
||
# ------------------------------------------------------------
|
||
|
||
gray_edges_overlay = Image.blend(
|
||
gray_320,
|
||
edge_map,
|
||
alpha=0.25,
|
||
)
|
||
|
||
gray_edges_path = (
|
||
OUTPUT_DIRECTORY
|
||
/ "06_gray_edges_320_q30.jpg"
|
||
)
|
||
|
||
save_jpeg(
|
||
gray_edges_overlay,
|
||
gray_edges_path,
|
||
quality=30,
|
||
)
|
||
|
||
variant_paths.append(
|
||
(
|
||
"06_gray_edges_320_q30",
|
||
gray_edges_path,
|
||
)
|
||
)
|
||
|
||
|
||
# ============================================================
|
||
# Расчёт параметров всех вариантов
|
||
# ============================================================
|
||
|
||
results = []
|
||
|
||
for variant_number, (
|
||
variant_name,
|
||
variant_path,
|
||
) in enumerate(variant_paths):
|
||
|
||
with Image.open(variant_path) as variant_image:
|
||
|
||
width, height = variant_image.size
|
||
mode = variant_image.mode
|
||
|
||
transfer_result = estimate_transfer(
|
||
file_path=variant_path,
|
||
image_id=IMAGE_ID_BASE + variant_number,
|
||
)
|
||
|
||
result = {
|
||
"name": variant_name,
|
||
"path": str(variant_path),
|
||
"width": width,
|
||
"height": height,
|
||
"mode": mode,
|
||
**transfer_result,
|
||
}
|
||
|
||
results.append(result)
|
||
|
||
|
||
# ============================================================
|
||
# Создание сравнительной картинки
|
||
# ============================================================
|
||
|
||
CONTACT_SHEET_COLUMNS = 3
|
||
CONTACT_SHEET_CELL_WIDTH = 360
|
||
CONTACT_SHEET_CELL_HEIGHT = 290
|
||
|
||
contact_sheet_rows = ceil(
|
||
len(results)
|
||
/ CONTACT_SHEET_COLUMNS
|
||
)
|
||
|
||
contact_sheet = Image.new(
|
||
"RGB",
|
||
(
|
||
CONTACT_SHEET_COLUMNS
|
||
* CONTACT_SHEET_CELL_WIDTH,
|
||
contact_sheet_rows
|
||
* CONTACT_SHEET_CELL_HEIGHT,
|
||
),
|
||
"white",
|
||
)
|
||
|
||
draw = ImageDraw.Draw(
|
||
contact_sheet
|
||
)
|
||
|
||
for result_index, result in enumerate(results):
|
||
|
||
column = (
|
||
result_index
|
||
% CONTACT_SHEET_COLUMNS
|
||
)
|
||
|
||
row = (
|
||
result_index
|
||
// CONTACT_SHEET_COLUMNS
|
||
)
|
||
|
||
cell_x = (
|
||
column
|
||
* CONTACT_SHEET_CELL_WIDTH
|
||
)
|
||
|
||
cell_y = (
|
||
row
|
||
* CONTACT_SHEET_CELL_HEIGHT
|
||
)
|
||
|
||
with Image.open(result["path"]) as variant_image:
|
||
|
||
preview = variant_image.convert("RGB")
|
||
|
||
preview.thumbnail(
|
||
(330, 220),
|
||
Image.Resampling.LANCZOS,
|
||
)
|
||
|
||
paste_x = (
|
||
cell_x
|
||
+ (
|
||
CONTACT_SHEET_CELL_WIDTH
|
||
- preview.width
|
||
)
|
||
// 2
|
||
)
|
||
|
||
paste_y = (
|
||
cell_y + 35
|
||
)
|
||
|
||
contact_sheet.paste(
|
||
preview,
|
||
(paste_x, paste_y),
|
||
)
|
||
|
||
label = (
|
||
f"{result['name']}\n"
|
||
f"{format_bytes(result['file_size_bytes'])}, "
|
||
f"{result['fragment_count']} fragments"
|
||
)
|
||
|
||
draw.text(
|
||
(cell_x + 10, cell_y + 8),
|
||
label,
|
||
fill="black",
|
||
)
|
||
|
||
contact_sheet_path = (
|
||
OUTPUT_DIRECTORY
|
||
/ "lab010_comparison.jpg"
|
||
)
|
||
|
||
contact_sheet.save(
|
||
contact_sheet_path,
|
||
format="JPEG",
|
||
quality=90,
|
||
)
|
||
|
||
|
||
# ============================================================
|
||
# Сохранение таблицы CSV
|
||
# ============================================================
|
||
|
||
csv_path = (
|
||
OUTPUT_DIRECTORY
|
||
/ "lab010_results.csv"
|
||
)
|
||
|
||
csv_fieldnames = [
|
||
"name",
|
||
"path",
|
||
"width",
|
||
"height",
|
||
"mode",
|
||
"file_size_bytes",
|
||
"fragment_count",
|
||
"total_radio_bytes",
|
||
"efficiency_percent",
|
||
"time_5_kbps",
|
||
"time_10_kbps",
|
||
"time_20_kbps",
|
||
"time_50_kbps",
|
||
"time_100_kbps",
|
||
"time_20_kbps_with_retries_x2",
|
||
]
|
||
|
||
with csv_path.open(
|
||
"w",
|
||
newline="",
|
||
encoding="utf-8-sig",
|
||
) as csv_file:
|
||
|
||
writer = DictWriter(
|
||
csv_file,
|
||
fieldnames=csv_fieldnames,
|
||
)
|
||
|
||
writer.writeheader()
|
||
|
||
for result in results:
|
||
|
||
writer.writerow(
|
||
{
|
||
"name": result["name"],
|
||
"path": result["path"],
|
||
"width": result["width"],
|
||
"height": result["height"],
|
||
"mode": result["mode"],
|
||
"file_size_bytes": (
|
||
result["file_size_bytes"]
|
||
),
|
||
"fragment_count": (
|
||
result["fragment_count"]
|
||
),
|
||
"total_radio_bytes": (
|
||
result["total_radio_bytes"]
|
||
),
|
||
"efficiency_percent": (
|
||
f"{result['efficiency_percent']:.2f}"
|
||
),
|
||
"time_5_kbps": (
|
||
result["transfer_times"][5]
|
||
),
|
||
"time_10_kbps": (
|
||
result["transfer_times"][10]
|
||
),
|
||
"time_20_kbps": (
|
||
result["transfer_times"][20]
|
||
),
|
||
"time_50_kbps": (
|
||
result["transfer_times"][50]
|
||
),
|
||
"time_100_kbps": (
|
||
result["transfer_times"][100]
|
||
),
|
||
"time_20_kbps_with_retries_x2": (
|
||
result["transfer_times"][20]
|
||
* 2
|
||
),
|
||
}
|
||
)
|
||
|
||
|
||
# ============================================================
|
||
# Вывод результатов
|
||
# ============================================================
|
||
|
||
print(
|
||
"=== Lab010. Оптимизация изображения ==="
|
||
)
|
||
|
||
print("\nИсходное разрешение:")
|
||
|
||
print(
|
||
original_width,
|
||
"x",
|
||
original_height,
|
||
)
|
||
|
||
|
||
print("\nСравнение вариантов:")
|
||
|
||
print(
|
||
f"{'Вариант':<28}"
|
||
f"{'Размер':>12}"
|
||
f"{'Фрагм.':>9}"
|
||
f"{'10 кбит/с':>13}"
|
||
f"{'20 кбит/с':>13}"
|
||
f"{'50 кбит/с':>13}"
|
||
f"{'20 кбит/с x2':>16}"
|
||
)
|
||
|
||
print("-" * 104)
|
||
|
||
for result in results:
|
||
|
||
print(
|
||
f"{result['name']:<28}"
|
||
f"{format_bytes(result['file_size_bytes']):>12}"
|
||
f"{result['fragment_count']:>9}"
|
||
f"{format_duration(result['transfer_times'][10]):>13}"
|
||
f"{format_duration(result['transfer_times'][20]):>13}"
|
||
f"{format_duration(result['transfer_times'][50]):>13}"
|
||
f"{format_duration(result['transfer_times'][20] * 2):>16}"
|
||
)
|
||
|
||
|
||
# ============================================================
|
||
# Итоговые кандидаты
|
||
# ============================================================
|
||
|
||
smallest_result = min(
|
||
results,
|
||
key=lambda item: item["file_size_bytes"],
|
||
)
|
||
|
||
operator_candidate = next(
|
||
result
|
||
for result in results
|
||
if result["name"] == "03_gray_320_q30"
|
||
)
|
||
|
||
|
||
print("\nСамый маленький файл:")
|
||
|
||
print(
|
||
smallest_result["name"],
|
||
"-",
|
||
format_bytes(
|
||
smallest_result["file_size_bytes"]
|
||
),
|
||
)
|
||
|
||
|
||
print("\nБазовый кандидат для оператора:")
|
||
|
||
print(
|
||
operator_candidate["name"],
|
||
"-",
|
||
format_bytes(
|
||
operator_candidate["file_size_bytes"]
|
||
),
|
||
"-",
|
||
format_duration(
|
||
operator_candidate[
|
||
"transfer_times"
|
||
][20]
|
||
),
|
||
"при 20 кбит/с без повторов",
|
||
)
|
||
|
||
|
||
print("\nСравнительная картинка:")
|
||
|
||
print(contact_sheet_path)
|
||
|
||
|
||
print("\nТаблица результатов CSV:")
|
||
|
||
print(csv_path)
|
||
|
||
|
||
# ============================================================
|
||
# Проверки
|
||
# ============================================================
|
||
|
||
assert len(results) == 7
|
||
|
||
assert contact_sheet_path.exists()
|
||
|
||
assert csv_path.exists()
|
||
|
||
assert all(
|
||
result["fragment_count"] > 0
|
||
for result in results
|
||
)
|
||
|
||
print(
|
||
"\nПроверка пройдена: "
|
||
"варианты кадра созданы и рассчитаны."
|
||
)
|