import gradio as gr
import numpy as np
from pysstv.color import Robot36
from pysstv.utils import load_wav
from io import BytesIO
import matplotlib.pyplot as plt

def decode_sstv(file):
    try:
        # Загрузка WAV файла
        wav_data, sample_rate = load_wav(file.name)
        
        # Используем декодер Robot36
        sstv = Robot36(wav_data, sample_rate)
        image = sstv.decode()

        # Конвертируем изображение для отображения
        fig, ax = plt.subplots()
        ax.imshow(image)
        ax.axis('off')
        buf = BytesIO()
        plt.savefig(buf, format="png")
        buf.seek(0)
        return buf
    except Exception as e:
        return f"Ошибка: {e}"

# Интерфейс Gradio
iface = gr.Interface(
    fn=decode_sstv,
    inputs=gr.File(label="Загрузите WAV файл (SSTV)"),
    outputs=gr.Image(type="file", label="Декодированное изображение"),
    title="SSTV Декодер",
    description="Загрузите WAV файл с SSTV сигналом для декодирования изображения."
)

# Запуск интерфейса
iface.launch()