thermalcam_decoder/decode.py
2024-02-15 00:30:45 +02:00

184 lines
5.2 KiB
Python

#!/usr/bin/env python3
import argparse
from pathlib import Path
import os
import subprocess
from io import BytesIO
import numpy as np
from tqdm import tqdm
from datetime import datetime
import pandas as pd
import pcapng
from struct import unpack
from PIL import Image
# Create the parser
parser = argparse.ArgumentParser(description="Process a pcap file.")
parser.add_argument("--live", action="store_true", help="Process images live")
# Add an argument for the pcap file, with a default value
parser.add_argument('input_file', nargs='?', default='in.pcap', help='The pcap file to process')
# Parse the arguments
args = parser.parse_args()
# TODO - probably a better way to do this
def live_capture():
while True:
# TODO: This is wrong:
# during scapy.all.sniff construction we lose about 27 packets.
# we work around it by capturing 2* packets(frame), so we just lose
# about half the frames.
cap = scapy.all.sniff(iface="enp1s0f0", count=32 * 2)
for pkt in cap:
l = len(bytes(pkt))
if l != 6972:
continue
data = bytes(pkt)
with open('udp.bytes', 'ab+') as fd:
fd.write(bytes(pkt.payload))
yield data
if args.live:
print('live stream, import scapy')
import scapy.all
print('open stream')
stream = live_capture()
start = datetime.now()
timestamp = start.strftime('%Y%m%d_%H%M%S')
basename = f'live_{timestamp}'
blocks = tqdm(stream)
else:
# Now use args.input_file as the file to process
input_file = args.input_file
basename = os.path.splitext(os.path.basename(input_file))[0]
stream = open(input_file, 'rb')
# Read packets from a pcap file
scanner = pcapng.scanner.FileScanner(stream)
blocks = tqdm(scanner)
def rightsize(it):
for i, obj in enumerate(it):
if isinstance(obj, bytes):
l = len(obj)
data = obj
else:
if not hasattr(obj, 'packet_len'):
continue
l = obj.packet_len
data = obj.packet_data
if l != 6972:
continue
yield data
def removestart(it):
"Remove the UDP header from the packets"
for x in it:
yield x[0x2A:]
# 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
def parsed(it):
for x in it:
yield parse(x)
# Function to group data into frames
def frames(it):
current = []
#otherdata = []
for obj in it:
if obj['part'] == 0:
if len(current) > 0:
yield b"".join(current)
current = []
#otherdata = []
current.append(obj["data"])
#otherdata.append(obj)
if len(current) > 0:
yield b"".join(current)
WIDTH = 384
HEIGHT = 288
def skip_bad_frames(it, width=WIDTH, height=HEIGHT):
for frame in it:
if len(frame) != width * height * 2: # 16 bpp
# Will be fixed when we stopped doing restarts
#print(f'{len(frame)} != {width} * {height} * 2')
continue
yield frame
def iterimages(it, width=WIDTH, height=HEIGHT, pixelformat=">H"):
for frame in it:
yield Image.fromarray(np.frombuffer(frame, dtype=pixelformat).reshape(width, height))
# Get frames and convert them to images
frames = skip_bad_frames(frames(parsed(removestart(rightsize(blocks)))))
# Create the directory for frames if not exists
frame_dir = f"frames/{basename}"
if not os.path.exists(frame_dir):
os.makedirs(frame_dir)
# Save each image as a PNG file
if not args.live:
images = iterimages(it=frames)
for i, img in enumerate(images):
img.save(f'frames/{basename}/{basename}_{i:04}.png')
ffmpeg_input = f"frames/{basename}/{basename}_%04d.png"
command = [
"ffmpeg",
"-y", # Overwrite output file without asking
"-hide_banner", # Hide banner
"-loglevel", "info", # Log level
"-f", "image2", # Input format
"-framerate", "25", # Framerate
"-i", ffmpeg_input, # Input file pattern
"-vf", "transpose=1", # Video filter for transposing
"-s", "384x288", # Size of one frame
"-vcodec", "libopenh264", # Video codec
"-pix_fmt", "yuv420p", # Pixel format: YUV 4:2:0
"thermal.mp4", # Output file in MP4 container
]
subprocess.run(command)
print("to play: ffplay thermal.mp4")
else:
# TODO: to video via ffmpeg; right now just a single png
# of the last frame
def todo_live_ffmpeg():
output = 'to_ffmpeg'
# live: write to named pipe
if not Path(output).exists():
print(f'making fifo at {output}')
os.mkfifo(output)
fd = open(output, 'wb')
for frame in frames:
fd.write(frame)
for frame in frames:
print('.')
Image.fromarray(np.frombuffer(frame, dtype='>H').reshape(WIDTH, HEIGHT)).save(f'live.new.png')
os.rename('live.new.png', 'live.png')