Add support for multiple partition video, and refactor the parsing of…

… ubnt_ubvinfo into a separate UbvInfoParser which generates a list of UbvPartition containing a list of FrameDataRef.

Also adds some basic unit tests, using sample output from ubvinfo

pom.xml

@@ -20,6 +20,15 @@
 		</repository>
 	</repositories>
 
+	<dependencies>
+		<dependency>
+			<groupId>junit</groupId>
+			<artifactId>junit</artifactId>
+			<version>4.12</version>
+			<scope>test</scope>
+		</dependency>
+	</dependencies>
+
 	<build>
 		<plugins>
 			<plugin>

src/main/java/FrameDataRef.java

@@ -0,0 +1,12 @@
+public class FrameDataRef
+{
+	public final int offset;
+	public final int size;
+
+
+	public FrameDataRef(final int offset, final int size)
+	{
+		this.offset = offset;
+		this.size = size;
+	}
+}

src/main/java/Remux.java

@@ -1,12 +1,12 @@
-import java.io.*;
+import java.io.File;
 import java.io.FileInputStream;
 import java.io.FileOutputStream;
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.nio.channels.FileChannel;
 import java.nio.file.Files;
-import java.util.ArrayList;
 import java.util.List;
+import java.util.stream.Stream;
 
 public class Remux
 {
@@ -17,59 +17,40 @@ public static void main(String[] args) throws Exception
 	{
 		final File inputFile = new File(args[0]);
 
-		final File parsedFile = new File(inputFile.getParent(), inputFile.getName() + ".txt");
-
-		final List<String> ubvinfoOutput;
-		if (parsedFile.exists())
-		{
-			// ubvinfo parse output already available
-			ubvinfoOutput = Files.readAllLines(parsedFile.toPath());
-		}
-		else
+		final List<UbvPartition> partitions;
 		{
-			ubvinfoOutput = ubvinfo(inputFile);
-		}
-
-		final int partitions = (int) ubvinfoOutput
-				                             .stream()
-				                             .filter(s -> s.equals("----------- PARTITION START -----------"))
-				                             .count();
+			final File parsedFile = new File(inputFile.getParent(), inputFile.getName() + ".txt");
 
-		if (partitions != 1)
-			throw new IllegalArgumentException(
-					"Input file contains multiple partitions, code does not currently handle this (discontinuities). Partition found: " +
-					partitions);
-
-
-		// TODO detect multiple partitions and create multiple output files
-
-		final File h264Stream = new File(inputFile.getParent(), inputFile.getName() + ".h264");
-
-		List<int[]> frames = new ArrayList<>();
-		boolean firstLine = true;
-		for (String line : ubvinfoOutput)
-		{
-			if (firstLine)
+			if (parsedFile.exists())
 			{
-				firstLine = false;
+				System.out.println("Cached ubnt_ubvinfo is available, using that instead of invoking ubnt_ubvinfo locally");
+				// ubvinfo parse output already available
+				partitions = UbvInfoParser.parse(Files.lines(parsedFile.toPath()));
 			}
 			else
 			{
-				if (Character.isWhitespace(line.charAt(0)))
-				{
-					String[] fields = line.split(" +", 7);
-
-					final int[] data = {Integer.parseInt(fields[4]), Integer.parseInt(fields[5])};
-					frames.add(data);
-				}
+				System.out.println("Invoking ubnt_ubvinfo on local machine...");
+				Stream<String> ubvinfoOutput = ubvinfo(inputFile);
+				partitions = UbvInfoParser.parse(ubvinfoOutput);
 			}
 		}
 
-		extractPrimitiveVideoStream(inputFile, h264Stream, frames);
+		System.out.println("Frame info parsed, extracting video frames...");
+
+		final boolean multipartition = partitions.size() > 1;
+
+		for (UbvPartition partition : partitions)
+		{
+			final String outputBasename = (multipartition) ? (inputFile.getName() + ".p" + partition.index) : inputFile.getName();
+
+			final File h264Stream = new File(inputFile.getParent(), outputBasename + ".h264");
+
+			extractPrimitiveVideoStream(inputFile, h264Stream, partition.frames);
+		}
 	}
 
 
-	private static List<String> ubvinfo(final File inputFile) throws IOException, InterruptedException
+	private static Stream<String> ubvinfo(final File inputFile) throws IOException, InterruptedException
 	{
 		// Write ubvinfo to a temporary file
 		File tempFile = null;
@@ -91,7 +72,7 @@ private static List<String> ubvinfo(final File inputFile) throws IOException, In
 			if (process.exitValue() != 0)
 				throw new IllegalArgumentException("ubvinfo failed!");
 
-			return Files.readAllLines(tempFile.toPath());
+			return Files.lines(tempFile.toPath());
 		}
 		finally
 		{
@@ -103,7 +84,7 @@ private static List<String> ubvinfo(final File inputFile) throws IOException, In
 
 	private static int extractPrimitiveVideoStream(final File inputFile,
 	                                               final File h264Stream,
-	                                               final List<int[]> frames) throws IOException
+	                                               final List<FrameDataRef> frames) throws IOException
 	{
 		final boolean video = true;
 
@@ -119,10 +100,10 @@ private static int extractPrimitiveVideoStream(final File inputFile,
 				{
 					try (FileChannel oc = fos.getChannel())
 					{
-						for (int[] offsetAndLength : frames)
+						for (FrameDataRef dataref : frames)
 						{
-							final int frameOffset = offsetAndLength[0];
-							final int frameLength = offsetAndLength[1];
+							final int frameOffset = dataref.offset;
+							final int frameLength = dataref.size;
 
 							// TODO for video: process NALs (ubv format uses length prefixes rather than NAL separators)
 							// TODO for audio: simply extract the raw data

src/main/java/UbvInfoParser.java

@@ -0,0 +1,55 @@
+import java.util.ArrayList;
+import java.util.Iterator;
+import java.util.List;
+import java.util.stream.Stream;
+
+public class UbvInfoParser
+{
+	private static final String PARTITION_START = "----------- PARTITION START -----------";
+	private static final int FIELD_OFFSET = 4;
+	private static final int FIELD_SIZE = 5;
+
+
+	/**
+	 * Takes the output of ubvinfo, expected to start with the following line:
+	 * <pre>Type   TID  KF           OFFSET     SIZE          DTS   CTS              WC     CR</pre>
+	 *
+	 * @param lines
+	 * @return
+	 */
+	public static List<UbvPartition> parse(Stream<String> lines)
+	{
+		final List<UbvPartition> results = new ArrayList<>();
+
+		UbvPartition current = null;
+		boolean firstLine = true;
+
+		final Iterator<String> iterator = lines.iterator();
+		while (iterator.hasNext())
+		{
+			final String line = iterator.next();
+
+			if (firstLine)
+			{
+				// Skip the first line (column headers) explicitly
+				firstLine = false;
+			}
+			else if (line.equals(PARTITION_START))
+			{
+				// Start a new partition
+				current = new UbvPartition(results.size() + 1);
+				results.add(current);
+			}
+			else if (Character.isWhitespace(line.charAt(0)))
+			{
+				final String[] fields = line.split(" +", 7);
+				final int offset = Integer.parseInt(fields[FIELD_OFFSET]);
+				final int size = Integer.parseInt(fields[FIELD_SIZE]);
+
+				current.frames.add(new FrameDataRef(offset, size));
+			}
+		}
+
+		return results;
+	}
+}

src/main/java/UbvPartition.java

@@ -0,0 +1,14 @@
+import java.util.ArrayList;
+import java.util.List;
+
+public class UbvPartition
+{
+	public final int index;
+	public final List<FrameDataRef> frames = new ArrayList<>();
+
+
+	public UbvPartition(final int index)
+	{
+		this.index = index;
+	}
+}

src/test/java/UbvInfoParserTest.java

@@ -0,0 +1,42 @@
+import org.junit.Test;
+
+import java.io.BufferedReader;
+import java.io.IOException;
+import java.io.InputStreamReader;
+import java.util.List;
+import java.util.stream.Stream;
+import java.util.zip.GZIPInputStream;
+
+import static org.junit.Assert.assertEquals;
+
+public class UbvInfoParserTest
+{
+	@Test
+	public void testSinglePartitionFile() throws IOException
+	{
+		final List<UbvPartition> partitions = UbvInfoParser.parse(readCompressedInfoFile(
+				"/FCECFFFFFFFF_2_rotating_1596209441895.ubv.txt.gz"));
+
+		assertEquals("expected partition count", 1, partitions.size());
+		assertEquals("frames in partition 1", 239194, partitions.get(0).frames.size());
+	}
+
+
+	@Test
+	public void testParseMultiPartition() throws IOException
+	{
+		final List<UbvPartition> partitions = UbvInfoParser.parse(readCompressedInfoFile(
+				"/F09FFFFFFFFF_2_timelapse_1556890741069.ubv.txt.gz"));
+
+		assertEquals("expected partition count", 44, partitions.size());
+		assertEquals("frames in partition 1", 366, partitions.get(0).frames.size());
+		assertEquals("partition 1 frame 1 offset", 96, partitions.get(0).frames.get(0).offset);
+		assertEquals("partition 1 frame 1 size", 2218, partitions.get(0).frames.get(0).size);
+	}
+
+
+	private Stream<String> readCompressedInfoFile(final String resource) throws IOException
+	{
+		return new BufferedReader(new InputStreamReader(new GZIPInputStream(getClass().getResourceAsStream(resource)))).lines();
+	}
+}