thermalcam_decoder/decode.py

from os import system
import numpy as np
from tqdm import tqdm
import pandas as pd
import pcapng
from struct import unpack
from PIL import Image

# Read packets from a pcap file
scanner = pcapng.scanner.FileScanner(open('in.pcap', 'rb'))
blocks = list(tqdm(scanner))


# Helper function to safely get an attribute from an object
def tryget(obj, att):
    if hasattr(obj, att):
        return getattr(obj, att)
    return None


# Create a DataFrame with packet lengths
df = pd.DataFrame([{'index': i, 'length': tryget(obj, 'packet_len')} for i, obj in enumerate(blocks)])


# Filter and extract data packets of a specific length
data = [blocks[i] for i in df[df.length == 6972.0].index]

# Remove the UDP header from the packets
raw = [d.packet_data[0x2a:] for d in data]


# Function to parse packet data
def parse(data):
    hdr = 4 + 2 * 7  # Header length
    # Unpack data into variables
    c1, c2, part, a, ffaa, b, c, d = unpack('>Lhhhhhhh', data[:hdr])
    ret = locals()
    del ret['data']
    del ret['hdr']
    ret['data'] = data[hdr:]
    return ret


# Parse each packet and create a DataFrame
df = pd.DataFrame([parse(d) for d in raw])
df2 = df[[c for c in df.columns if c != 'data']]


# Function to group data into frames
def getframes(df):
    frames = []
    current = []
    for i, row in df.iterrows():
        if row.part == 0:
            if len(current) > 0:
                frames.append(current)
            current = []
        current.append(row['data'])
    if len(current) > 0:
        frames.append(current)
    return [b''.join(parts) for parts in frames]


# Function to convert binary frame data into images
def image16(frame, width, height, pixelformat='>H'):
    return [
        Image.fromarray(np.frombuffer(frame, dtype=pixelformat).reshape(width, height)) for frame in frames
        if len(frame) == 2 * width * height
    ]


# Get frames and convert them to images
frames = getframes(df)
images = image16(frames, 384, 288)

# Save each image as a PNG file
for i, img in enumerate(tqdm(images)):
    img.save(f'{i:04}.png')

# Produce a video from the saved images
system('ffmpeg -hide_banner -loglevel info -y -f image2 -framerate 25 -i %04d.png -vf "transpose=1" -s 384x288 -vcodec libx264 -pix_fmt yuv420p thermal.mp4')
print('to play: ffplay thermal.mp4')