encoder.py

import numpy as np
import math
import argparse
from PIL import Image
from scipy.io import wavfile
import common

class SSTVEncoder:
    def __init__(self, image_source, width=320, height=240, sample_rate=common.SAMPLE_RATE):
        self.image_source = image_source
        self.config = common.get_mode_config(width, height, sample_rate)
        self.phase = 0.0
        self.audio_buffer = []

    def load_and_process_image(self):
        if isinstance(self.image_source, str):
            img = Image.open(self.image_source).convert('RGB')
        elif isinstance(self.image_source, Image.Image):
            img = self.image_source.convert('RGB')
        elif isinstance(self.image_source, np.ndarray):
            img = Image.fromarray(self.image_source).convert('RGB')
        else:
            raise ValueError("Unsupported image source type")
            
        img = img.resize((self.config.width, self.config.height), Image.Resampling.LANCZOS)
        return np.array(img)

    def add_tone(self, freq, duration_samples):
        # Generate tone with continuous phase
        t = np.arange(duration_samples)
        phase_increment = 2 * np.pi * freq / self.config.sample_rate
        phases = self.phase + phase_increment * t
        
        # Update phase for next segment
        self.phase = phases[-1] + phase_increment 
        
        tone = np.sin(phases)
        self.audio_buffer.append(tone)

    def add_scan_line_tone(self, pixels, duration_samples, channel_idx):
        # pixels: array of pixel values for the channel
        # Map pixel values to frequency range 1500Hz - 2300Hz
        
        freqs = common.BLACK_FREQ + (pixels * (common.WHITE_FREQ - common.BLACK_FREQ) / 255.0)
        
        # Generate samples with varying frequency (FM)
        t_indices = np.linspace(0, len(pixels) - 1, duration_samples)
        interpolated_freqs = np.interp(t_indices, np.arange(len(pixels)), freqs)
        
        # Integrate frequency to get phase
        # phase[n] = phase[n-1] + 2*pi*f[n]/Fs
        phase_increments = 2 * np.pi * interpolated_freqs / self.config.sample_rate
        cumulative_phases = np.cumsum(phase_increments) + self.phase
        
        self.phase = cumulative_phases[-1]
        
        tone = np.sin(cumulative_phases)
        self.audio_buffer.append(tone)

    def generate_header(self):
        # Calibration Header
        # Leader tone
        self.add_tone(common.LEADER_TONE_FREQ, self.config.leader_tone_samples)
        # Break
        self.add_tone(common.BREAK_FREQ, self.config.break_samples)
        # Leader tone
        self.add_tone(common.LEADER_TONE_FREQ, self.config.leader_tone_samples)
        # VIS Start bit
        self.add_tone(common.BREAK_FREQ, self.config.vis_bit_samples)
        
        # VIS Code 8 (Robot36): 00001000 (LSB first) -> 0, 0, 0, 1, 0, 0, 0, 0
        vis_code = 8
        parity = 0
        for i in range(7):
            bit = (vis_code >> i) & 1
            parity ^= bit
            freq = common.VIS_BIT_1_FREQ if bit else common.VIS_BIT_0_FREQ
            self.add_tone(freq, self.config.vis_bit_samples)
            
        # Parity bit
        freq = common.VIS_BIT_1_FREQ if parity else common.VIS_BIT_0_FREQ
        self.add_tone(freq, self.config.vis_bit_samples)
        
        # Stop bit
        self.add_tone(common.BREAK_FREQ, self.config.vis_bit_samples)

    def write_wav(self, file_object):
        img_array = self.load_and_process_image()
        
        # Convert to YUV
        r = img_array[:,:,0].astype(float)
        g = img_array[:,:,1].astype(float)
        b = img_array[:,:,2].astype(float)
        
        y, u, v = common.rgb_to_yuv(r, g, b)
        
        # Generate Audio
        self.generate_header()
        
        for line in range(self.config.height):
            # Sync
            self.add_tone(common.SYNC_FREQ, self.config.sync_samples)
            # Sync Porch
            self.add_tone(common.PORCH_FREQ, self.config.sync_porch_samples)
            
            # Y Scan
            self.add_scan_line_tone(y[line], self.config.y_scan_samples, 0)
            
            if line % 2 == 0:
                # Even Line: Separator (1500) -> Porch -> V Scan
                self.add_tone(common.EVEN_SEPARATOR_FREQ, self.config.separator_samples)
                self.add_tone(common.PORCH_SEPARATOR_FREQ, self.config.porch_samples)
                self.add_scan_line_tone(v[line], self.config.uv_scan_samples, 2)
            else:
                # Odd Line: Separator (2300) -> Porch -> U Scan
                self.add_tone(common.ODD_SEPARATOR_FREQ, self.config.separator_samples)
                self.add_tone(common.PORCH_SEPARATOR_FREQ, self.config.porch_samples)
                self.add_scan_line_tone(u[line], self.config.uv_scan_samples, 1)

        # Concatenate and save
        full_signal = np.concatenate(self.audio_buffer)
        
        # Normalize to 16-bit PCM range
        # full_signal is -1.0 to 1.0
        audio_data = (full_signal * 32767).astype(np.int16)
        
        wavfile.write(file_object, self.config.sample_rate, audio_data)

    def encode(self, output_path):
        with open(output_path, 'wb') as f:
            self.write_wav(f)
        print(f"SSTV Audio saved to {output_path}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Python Robot36 SSTV Encoder")
    parser.add_argument("input_image", help="Path to input image")
    parser.add_argument("output_wav", help="Path to output WAV file")
    parser.add_argument("--width", type=int, default=320, help="Output width (default: 320)")
    parser.add_argument("--height", type=int, default=240, help="Output height (default: 240)")
    
    args = parser.parse_args()
    
    encoder = SSTVEncoder(args.input_image, args.width, args.height)
    encoder.encode(args.output_wav)