import numpy as np
from pydub import AudioSegment
from matplotlib import pyplot as plt
from pysstv.color import Robot36
import gradio as gr

def decode_sstv(audio_file):
    # Чтение аудиофайла
    audio = AudioSegment.from_file(audio_file)
    samples = np.array(audio.get_array_of_samples(), dtype=np.int16)

    # Декодирование SSTV сигнала с передачей numpy-массива
    # На этот раз, создаем изображение после преобразования образца
    sstv = Robot36(samples, samples_per_sec=audio.frame_rate, bits=16, color_depth=8)
    sstv_img = sstv.decode()
    img = np.array(sstv_img)

    # Сохранение и отображение изображения
    plt.imshow(img)
    plt.axis('off')
    img_path = 'decoded_image.png'
    plt.savefig(img_path, bbox_inches='tight', pad_inches=0)
    return img_path

iface = gr.Interface(
    fn=decode_sstv,
    inputs=gr.Audio(type="filepath", label="Upload SSTV Audio File"),
    outputs=gr.Image(label="Decoded Image"),
    title="SSTV Decoder",
    description="Upload an SSTV audio file to decode the image."
)

if __name__ == "__main__":
    iface.launch()