sstv-decoder-prototype/render.py

131 lines
2.7 KiB
Python

#!/usr/bin/env python3
import matplotlib.pyplot as plt
import struct
from PIL import Image
IMG_WIDTH = 640
SYNC_LENGTH = 105
SYNC_THRESH = 100
PORCH_LENGTH = 10
LINE_LENGTH = SYNC_LENGTH + PORCH_LENGTH + 4 * IMG_WIDTH
MAX_DEV = 600
CENTER = 1750
MIN_FREQ = 1200
MAX_FREQ = 2300
COLOR_LOW = 1500
COLOR_HIGH = MAX_FREQ
def to_freq(sample):
freq = CENTER + sample * MAX_DEV
freq = max(MIN_FREQ, freq)
freq = min(MAX_FREQ, freq)
return freq
def to_color(sample):
sample = (sample - 1500) / (COLOR_HIGH - COLOR_LOW)
sample = int(sample * 255)
sample = max(sample, 0)
sample = min(sample, 255)
return sample
def ycbcr_to_rgb(y, cb, cr):
r = int(y + 1.402 * (cr - 128))
g = int(y - 3.44136 * (cb - 128) - 0.714136 * (cr - 128))
b = int(y + 1.772 * (cb - 128))
return (r,g,b)
last_sync = False
def is_sync(backlog):
global last_sync
count = 0
for sample in backlog:
if sample < COLOR_LOW:
count += 1
res = False
if count > SYNC_THRESH and not last_sync:
res = True
last_sync = count > SYNC_THRESH
return res
def render_lines(line):
pixels = [[(0,0,0)] * IMG_WIDTH, [(0,0,0)] * IMG_WIDTH]
# Strip porch
porch = len(line) - SYNC_LENGTH - 4 * IMG_WIDTH
if porch < 0:
return pixels
line = line[PORCH_LENGTH:]
for i in range(0, IMG_WIDTH):
y0 = to_color(line[i])
cr = to_color(line[IMG_WIDTH + i])
cb = to_color(line[2 * IMG_WIDTH + i])
y1 = to_color(line[3 * IMG_WIDTH + i])
pixels[0][i] = y0, cb, cr
pixels[1][i] = y1, cb, cr
return pixels
def main():
data_file = open("demod.dat", 'rb')
samples = []
syncs = []
pixels = []
backlog = [0.0] * SYNC_LENGTH
line = []
bytes = data_file.read(4)
while len(bytes) == 4:
sample = struct.unpack('f', bytes)[0]
sample = to_freq(sample)
samples += [sample]
backlog = backlog[1:] + [sample]
line += [sample]
if is_sync(backlog) or len(line) > LINE_LENGTH:
if len(line) > SYNC_LENGTH + 4 * IMG_WIDTH:
syncs += [1000]
print("%d \t %d" % (len(line), len(line) - LINE_LENGTH))
pixels += render_lines(line)
else:
print("Short line dropped")
line = []
else:
syncs += [0]
bytes = data_file.read(4)
plt.plot(samples)
plt.plot(syncs)
plt.show()
img = Image.new('YCbCr', (IMG_WIDTH, len(pixels)))
img_pix = img.load()
for x in range(0, IMG_WIDTH):
for y in range(0, len(pixels)):
img_pix[x,y] = pixels[y][x]
img = img.convert('RGB')
img.save('output.png')
if __name__ == '__main__':
main()