feat: add per-track segmentation init mode

Author: lideen

.changeset/kind-turtles-dance.md

@@ -0,0 +1,5 @@
+---
+"mp4box": minor
+---
+
+feat: add per-track initialization segments for segmentation

README.md

@@ -270,28 +270,29 @@ mp4boxfile.onSegment = function (id, user, buffer, sampleNumber, last) {
 };
 ```
 
-#### initializeSegmentation()
+#### initializeSegmentation(mode)
 
-Indicates that the application is ready to receive segments. Returns an array of objects containing the following properties:
+Indicates that the application is ready to receive segments.
+`mode` can be `'combined'` (default) or `'per-track'`.
 
-- **id**: Number, the track id
-- **user**: Object, the caller of the segmentation for this track, as given in [setSegmentOptions](##setsegmentoptionstrack_id-user-options)
-- **buffer**: ArrayBuffer, the initialization segment for this track.
-- **sampleNumber**: Number, sample number of the last sample in the segment, plus 1.
-- **last**: Boolean, indication if this is the last segment to be received.
+By default, it returns a single initialization segment containing all tracks configured with [setSegmentOptions](#setsegmentoptionstrack_id-user-options):
+
+```json
+{
+  "tracks": [
+    { "id": 2, "user": "[SourceBuffer]" },
+    { "id": 3, "user": "[SourceBuffer]" }
+  ],
+  "buffer": "[ArrayBuffer]"
+}
+```
+
+If called with `'per-track'`, it returns one initialization segment per fragmented track:
 
 ```json
 [
-  {
-    "id": 2,
-    "buffer": "[ArrayBuffer]",
-    "user": "[SourceBuffer]"
-  },
-  {
-    "id": 3,
-    "buffer": "[ArrayBuffer]",
-    "user": "[SourceBuffer]"
-  }
+  { "id": 2, "buffer": "[ArrayBuffer]", "user": "[SourceBuffer]" },
+  { "id": 3, "buffer": "[ArrayBuffer]", "user": "[SourceBuffer]" }
 ]
 ```
 

entries/types.ts

@@ -119,6 +119,23 @@ export interface FragmentedTrack<TUser> {
     accumulatedSize: number;
   };
 }
+
+export interface SegmentationInitializationTrack<TUser> {
+  id: number;
+  user: TUser;
+}
+
+export interface SegmentationInitialization<TUser> {
+  tracks: Array<SegmentationInitializationTrack<TUser>>;
+  buffer: ArrayBuffer;
+}
+
+export interface SegmentationInitializationPerTrack<
+  TUser,
+> extends SegmentationInitializationTrack<TUser> {
+  buffer: ArrayBuffer;
+}
+
 export interface ExtractedTrack<TUser> {
   id: number;
   user: TUser;

src/isofile.ts

@@ -93,6 +93,8 @@ import type {
   Item,
   Movie,
   Output,
+  SegmentationInitialization,
+  SegmentationInitializationPerTrack,
   Sample,
   SampleEntryFourCC,
   SubSample,
@@ -1303,42 +1305,65 @@ export class ISOFile<TSegmentUser = unknown, TSampleUser = unknown> {
   }
 
   /** @bundle isofile-write.js */
-  initializeSegmentation() {
+  initializeSegmentation(mode?: 'combined'): SegmentationInitialization<TSegmentUser>;
+  initializeSegmentation(
+    mode: 'per-track',
+  ): Array<SegmentationInitializationPerTrack<TSegmentUser>>;
+  initializeSegmentation(
+    mode?: 'combined' | 'per-track',
+  ):
+    | SegmentationInitialization<TSegmentUser>
+    | Array<SegmentationInitializationPerTrack<TSegmentUser>> {
     if (!this.onSegment) {
       Log.warn('MP4Box', 'No segmentation callback set!');
     }
+    if (mode !== undefined && mode !== 'combined' && mode !== 'per-track') {
+      throw new Error(`Invalid segmentation mode: ${mode}`);
+    }
     if (!this.isFragmentationInitialized) {
       this.isFragmentationInitialized = true;
       this.resetTables();
     }
 
-    // Create the moov that will hold all the tracks
-    const moov = new moovBox();
-    moov.addBox(this.moov.mvhd);
-
-    // Add the tracks we want to fragment
-    for (let i = 0; i < this.fragmentedTracks.length; i++) {
-      const trak = this.getTrackById(this.fragmentedTracks[i].id);
+    const tracksToInitialize: Array<{ id: number; user: TSegmentUser; trak: trakBox }> = [];
+    for (const fragmentedTrack of this.fragmentedTracks) {
+      const trak = this.getTrackById(fragmentedTrack.id);
       if (!trak) {
         Log.warn(
           'ISOFile',
-          `Track with id ${this.fragmentedTracks[i].id} not found, skipping fragmentation initialization`,
+          `Track with id ${fragmentedTrack.id} not found, skipping fragmentation initialization`,
         );
         continue;
       }
-      moov.addBox(trak);
+      tracksToInitialize.push({ id: fragmentedTrack.id, user: fragmentedTrack.user, trak });
+    }
+
+    const fragmentDuration = this.moov?.mvex?.mehd.fragment_duration;
+
+    if (mode === 'per-track') {
+      return tracksToInitialize.map(({ id, user, trak }) => {
+        const moov = new moovBox();
+        moov.addBox(this.moov.mvhd);
+        moov.addBox(trak);
+
+        return {
+          id,
+          user,
+          buffer: ISOFile.writeInitializationSegment(this.ftyp, moov, fragmentDuration),
+        };
+      });
+    }
+
+    // Create the moov that will hold all selected fragmented tracks
+    const moov = new moovBox();
+    moov.addBox(this.moov.mvhd);
+    for (const track of tracksToInitialize) {
+      moov.addBox(track.trak);
     }
 
     return {
-      tracks: moov.traks.map((trak, i) => ({
-        id: trak.tkhd.track_id,
-        user: this.fragmentedTracks[i].user,
-      })),
-      buffer: ISOFile.writeInitializationSegment(
-        this.ftyp,
-        moov,
-        this.moov?.mvex?.mehd.fragment_duration,
-      ),
+      tracks: tracksToInitialize.map(({ id, user }) => ({ id, user })),
+      buffer: ISOFile.writeInitializationSegment(this.ftyp, moov, fragmentDuration),
     };
   }
 

tests/segmentation.test.ts

@@ -31,7 +31,7 @@ describe('File Segmentation', () => {
     const init = mp4.initializeSegmentation();
 
     // Write the initialization segments to the output stream
-    out.insertBuffer(init.buffer);
+    out.insertBuffer(MP4BoxBuffer.fromArrayBuffer(init.buffer, offset));
     offset += init.buffer.byteLength;
     saveBufferToFile(init.buffer, task.id, true);
 
@@ -81,7 +81,7 @@ describe('File Segmentation', () => {
     const init = mp4.initializeSegmentation();
 
     // Write the initialization segments to the output stream
-    out.insertBuffer(init.buffer);
+    out.insertBuffer(MP4BoxBuffer.fromArrayBuffer(init.buffer, offset));
     offset += init.buffer.byteLength;
     saveBufferToFile(init.buffer, task.id, true);
 
@@ -128,7 +128,7 @@ describe('File Segmentation', () => {
     const init = mp4.initializeSegmentation();
 
     // Write the initialization segments to the output stream
-    out.insertBuffer(init.buffer);
+    out.insertBuffer(MP4BoxBuffer.fromArrayBuffer(init.buffer, offset));
     offset += init.buffer.byteLength;
     saveBufferToFile(init.buffer, task.id, true);
 
@@ -178,7 +178,7 @@ describe('File Segmentation', () => {
     const init = mp4.initializeSegmentation();
 
     // Write the initialization segments to the output stream
-    out.insertBuffer(init.buffer);
+    out.insertBuffer(MP4BoxBuffer.fromArrayBuffer(init.buffer, offset));
     offset += init.buffer.byteLength;
     saveBufferToFile(init.buffer, task.id, true);
 
@@ -291,4 +291,59 @@ describe('File Segmentation', () => {
     expect(fragTrack.state.accumulatedSize).toBe(0);
     expect(fragTrack.segmentStream).toBeUndefined();
   });
+
+  it('with one init segment per fragmented track', async () => {
+    const { testFile } = getFilePath('isobmff', '01_simple.mp4');
+    const { mp4 } = await loadAndGetInfo(testFile, true, true);
+
+    mp4.setSegmentOptions(101, undefined, {
+      nbSamples: 50,
+      rapAlignement: false,
+    });
+    mp4.setSegmentOptions(201, undefined, {
+      nbSamples: 50,
+      rapAlignement: false,
+    });
+
+    const initSegments = mp4.initializeSegmentation('per-track');
+    expect(initSegments.length).toBe(2);
+
+    for (const initSegment of initSegments) {
+      const initMp4 = createFile();
+      initMp4.appendBuffer(MP4BoxBuffer.fromArrayBuffer(initSegment.buffer, 0));
+      initMp4.flush();
+
+      const info = initMp4.getInfo();
+      expect(info.tracks.length).toBe(1);
+      expect(info.tracks[0].id).toBe(initSegment.id);
+      expect(initMp4.moov.mvex.trexs.length).toBe(1);
+      expect(initMp4.moov.mvex.trexs[0].track_id).toBe(initSegment.id);
+    }
+  });
+
+  it('with explicit combined mode', async () => {
+    const { testFile } = getFilePath('isobmff', '01_simple.mp4');
+    const { mp4 } = await loadAndGetInfo(testFile, true, true);
+
+    mp4.setSegmentOptions(101, undefined, {
+      nbSamples: 50,
+      rapAlignement: false,
+    });
+    mp4.setSegmentOptions(201, undefined, {
+      nbSamples: 50,
+      rapAlignement: false,
+    });
+
+    const defaultInit = mp4.initializeSegmentation();
+    const combinedInit = mp4.initializeSegmentation('combined');
+
+    expect(combinedInit.tracks.map(track => track.id)).toEqual(
+      defaultInit.tracks.map(track => track.id),
+    );
+
+    const initMp4 = createFile();
+    initMp4.appendBuffer(MP4BoxBuffer.fromArrayBuffer(combinedInit.buffer, 0));
+    initMp4.flush();
+    expect(initMp4.getInfo().tracks.length).toBe(2);
+  });
 });