From d3840103fcfb5394e92ae1a27bb41858b04886cd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?C=C3=A9dric=20Verstraeten?= <cedric@verstraeten.io>
Date: Thu, 5 Jun 2025 13:28:10 +0000
Subject: [PATCH] Add VPS NALUs support in IPCamera configuration and MP4
 handling for improved video processing

---
 machinery/src/capture/main.go        | 3 ++-
 machinery/src/components/Kerberos.go | 1 +
 machinery/src/models/Config.go       | 1 +
 machinery/src/video/mp4.go           | 6 ++++--
 4 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/machinery/src/capture/main.go b/machinery/src/capture/main.go
index d89b6dd..84f1e50 100644
--- a/machinery/src/capture/main.go
+++ b/machinery/src/capture/main.go
@@ -240,9 +240,10 @@ func HandleRecordStream(queue *packets.Queue, configDirectory string, configurat
 					// Get SPS and PPS NALUs from the camera.
 					spsNALUS := configuration.Config.Capture.IPCamera.SPSNALUs
 					ppsNALUS := configuration.Config.Capture.IPCamera.PPSNALUs
+					vpsNALUS := configuration.Config.Capture.IPCamera.VPSNALUs
 
 					// Create a video file, and set the dimensions.
-					mp4Video = video.NewMP4(fullName, spsNALUS, ppsNALUS)
+					mp4Video = video.NewMP4(fullName, spsNALUS, ppsNALUS, vpsNALUS)
 					mp4Video.SetWidth(width)
 					mp4Video.SetHeight(height)
 
diff --git a/machinery/src/components/Kerberos.go b/machinery/src/components/Kerberos.go
index 06cc3de..84c8090 100644
--- a/machinery/src/components/Kerberos.go
+++ b/machinery/src/components/Kerberos.go
@@ -170,6 +170,7 @@ func RunAgent(configDirectory string, configuration *models.Configuration, commu
 	// Set the SPS and PPS values in the configuration.
 	configuration.Config.Capture.IPCamera.SPSNALUs = [][]byte{videoStream.SPS}
 	configuration.Config.Capture.IPCamera.PPSNALUs = [][]byte{videoStream.PPS}
+	configuration.Config.Capture.IPCamera.VPSNALUs = [][]byte{videoStream.VPS}
 
 	// Define queues for the main and sub stream.
 	var queue *packets.Queue
diff --git a/machinery/src/models/Config.go b/machinery/src/models/Config.go
index ffc7f70..8e246fa 100644
--- a/machinery/src/models/Config.go
+++ b/machinery/src/models/Config.go
@@ -91,6 +91,7 @@ type IPCamera struct {
 	ONVIFPassword string   `json:"onvif_password" bson:"onvif_password"`
 	SPSNALUs      [][]byte `json:"sps_nalus,omitempty" bson:"sps_nalus,omitempty"`
 	PPSNALUs      [][]byte `json:"pps_nalus,omitempty" bson:"pps_nalus,omitempty"`
+	VPSNALUs      [][]byte `json:"pps_nalus,omitempty" bson:"pps_nalus,omitempty"`
 }
 
 // USBCamera configuration, such as the device path (/dev/video*)
diff --git a/machinery/src/video/mp4.go b/machinery/src/video/mp4.go
index 2bebb75..9ac1c32 100644
--- a/machinery/src/video/mp4.go
+++ b/machinery/src/video/mp4.go
@@ -42,6 +42,7 @@ type MP4 struct {
 	Start              bool
 	SPSNALUs           [][]byte // SPS NALUs for H264
 	PPSNALUs           [][]byte // PPS NALUs for H264
+	VPSNALUs           [][]byte // VPS NALUs for H264
 	FreeBoxSize        int64
 	MoofBoxes          int64             // Number of moof boxes in the file
 	MoofBoxSizes       []int64           // Sizes of each moof box
@@ -58,7 +59,7 @@ type MP4 struct {
 }
 
 // NewMP4 creates a new MP4 object
-func NewMP4(fileName string, spsNALUs [][]byte, ppsNALUs [][]byte) *MP4 {
+func NewMP4(fileName string, spsNALUs [][]byte, ppsNALUs [][]byte, vpsNALUs [][]byte) *MP4 {
 
 	init := mp4ff.NewMP4Init()
 
@@ -92,6 +93,7 @@ func NewMP4(fileName string, spsNALUs [][]byte, ppsNALUs [][]byte) *MP4 {
 		Writer:      bufferedWriter,
 		SPSNALUs:    spsNALUs,
 		PPSNALUs:    ppsNALUs,
+		VPSNALUs:    vpsNALUs,
 	}
 }
 
@@ -331,7 +333,7 @@ func (mp4 *MP4) Close(config *models.Config) {
 	} else if mp4.VideoTrackName == "H265" || mp4.VideoTrackName == "HEV1" {
 		init.AddEmptyTrack(videoTimescale, "video", "und")
 		includePS := true
-		err := init.Moov.Traks[0].SetHEVCDescriptor("hev1", [][]byte{}, mp4.SPSNALUs, mp4.PPSNALUs, [][]byte{}, includePS)
+		err := init.Moov.Traks[0].SetHEVCDescriptor("hev1", mp4.VPSNALUs, mp4.SPSNALUs, mp4.PPSNALUs, [][]byte{}, includePS)
 		if err != nil {
 			//panic(err)
 		}