Add vu-from-video

Signed-off-by: Mattias Andrée <maandree@kth.se>

7 files changed, 296 insertions, 3 deletions

diff --git a/src/stream.h b/src/stream.h
index b59527a..027db4a 100644
--- a/src/stream.h
+++ b/src/stream.h
@@ -2,6 +2,8 @@
 #include <stddef.h>
 #include <stdio.h>
 
+#define STREAM_HEAD_MAX (3 * 3 * sizeof(size_t) + sizeof(((struct stream *)0)->pixfmt) + 10)
+
 #define einit_stream(...)      eninit_stream(1, __VA_ARGS__)
 #define eset_pixel_size(...)   enset_pixel_size(1, __VA_ARGS__)
 #define eread_stream(...)      enread_stream(1, __VA_ARGS__)
diff --git a/src/util/colour.h b/src/util/colour.h
index ff40841..854705a 100644
--- a/src/util/colour.h
+++ b/src/util/colour.h
@@ -29,6 +29,30 @@ srgb_decode(double t)
 }
 
 static inline void
+yuv_to_srgb(double y, double u, double v, double *r, double *g, double *b)
+{
+#define MULTIPLY(CY, CU, CV)  ((CY) * y + (CU) * u + (CV) * v)
+	*r = MULTIPLY(1,  0.00028328010485821202317155420580263580632163211703,  1.14070449590558520291949662350816652178764343261719);
+	*g = MULTIPLY(1, -0.39630886669497211727275498560629785060882568359375, -0.58107364288228224857846271333983168005943298339844);
+	*b = MULTIPLY(1,  2.03990003507541306504435851820744574069976806640625,  0.00017179031692307700847528739718228507626918144524);
+#undef MULTIPLY
+}
+
+static inline void
+srgb_to_yuv(double r, double g, double b, double *y, double *u, double *v)
+{
+#define MULTIPLY(CR, CG, CB) ((CR) * r + (CG) * g + (CB) * b)
+	*y = MULTIPLY(0.299, 0.587, 0.114);
+	*u = MULTIPLY(-0.14662756598240470062854967636667424812912940979004,
+		      -0.28771586836102963635752871596196200698614120483398,
+		       0.43434343434343436474165400795754976570606231689453);
+	*v = MULTIPLY( 0.61456892577224520035628074765554629266262054443359,
+		      -0.51452282157676354490405401520547457039356231689453,
+		      -0.10004610419548178035231700278018251992762088775635);
+#undef MULTIPLY
+}
+
+static inline void
 ciexyz_to_srgb(double x, double y, double z, double *r, double *g, double *b)
 {
 #define MULTIPLY(CX, CY, CZ)  ((CX) * x + (CY) * y + (CZ) * z)
diff --git a/src/vu-concat.c b/src/vu-concat.c
index e354962..95676e2 100644
--- a/src/vu-concat.c
+++ b/src/vu-concat.c
@@ -49,7 +49,7 @@ concat_to_file(int argc, char *argv[], char *output_file)
 	struct stream stream, refstream;
 	int first = 0;
 	int fd = eopen(output_file, O_RDWR | O_CREAT | O_TRUNC, 0666);
-	char head[3 * 3 * sizeof(size_t) + sizeof(stream.pixfmt) + 10];
+	char head[STREAM_HEAD_MAX];
 	ssize_t headlen, size = 0;
 	char *data;
 
diff --git a/src/vu-from-video.c b/src/vu-from-video.c
new file mode 100644
index 0000000..96bd83d
--- /dev/null
+++ b/src/vu-from-video.c
@@ -0,0 +1,265 @@
+/* See LICENSE file for copyright and license details. */
+#include "stream.h"
+#include "util.h"
+
+#if defined(HAVE_PRCTL)
+# include <sys/prctl.h>
+#endif
+#include <sys/mman.h>
+#include <sys/stat.h>
+#include <sys/wait.h>
+#include <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <unistd.h>
+
+USAGE("[-r frame-rate] [-w width -h height] input-file output-file")
+
+static void
+read_metadata(FILE *fp, char *fname, size_t *width, size_t *height, size_t *frames)
+{
+	char *line = NULL;
+	size_t size = 0;
+	ssize_t len;
+	int have_width = !width, have_height = !height, have_frames = 0;
+	char *p;
+
+	while ((len = getline(&line, &size, fp)) != -1) {
+		if (len && line[len - 1])
+			line[--len] = '\0';
+		p = strchr(line, '=') + 1;
+		if (width && strstr(line, "width=") == line) {
+			if (tozu(p, 1, SIZE_MAX, width))
+				eprintf("invalid width: %s\n", p);
+			have_width = 1;
+		} else if (height && strstr(line, "height=") == line) {
+			if (tozu(p, 1, SIZE_MAX, height))
+				eprintf("invalid height: %s\n", p);
+			have_height = 1;
+		} else if (strstr(line, "nb_read_frames=") == line) {
+			if (tozu(p, 0, SIZE_MAX, frames))
+				eprintf("invalid frame count: %s\n", p);
+			have_frames = 1;
+		}
+	}
+
+	if (ferror(fp))
+		eprintf("getline %s:", fname);
+	free(line);
+
+	if (have_width + have_height + have_frames < 3)
+		eprintf("could not get all required metadata\n");
+}
+
+static void
+get_metadata(char *file, size_t *width, size_t *height, size_t *frames)
+{
+	FILE *fp;
+	int fd, pipe_rw[2];
+	pid_t pid;
+	int status;
+
+	if (pipe(pipe_rw))
+		eprintf("pipe:");
+
+	pid = fork();
+	if (pid == -1)
+		eprintf("fork:");
+
+	if (!pid) {
+#if defined(HAVE_PRCTL) && defined(PR_SET_PDEATHSIG)
+		prctl(PR_SET_PDEATHSIG, SIGKILL);
+#endif
+		fd = eopen(file, O_RDONLY);
+		if (dup2(fd, STDIN_FILENO) == -1)
+			eprintf("dup2:");
+		close(fd);
+		close(pipe_rw[0]);
+		if (dup2(pipe_rw[1], STDOUT_FILENO) == -1)
+			eprintf("dup2:");
+		close(pipe_rw[1]);
+		execlp("ffprobe", "ffprobe", "-v", "quiet", "-show_streams",
+		       "-select_streams", "v", "-count_frames", "-", NULL);
+		eprintf("exec ffprobe:");
+	}
+
+	close(pipe_rw[1]);
+	fp = fdopen(pipe_rw[0], "rb");
+	if (!fp)
+		eprintf("fdopen <subprocess>:");
+	read_metadata(fp, file, width, height, frames);
+	fclose(fp);
+	close(pipe_rw[0]);
+
+	if (waitpid(pid, &status, 0) == -1)
+		eprintf("waitpid:");
+	if (status)
+		exit(1);
+}
+
+static void
+convert_segment(char *buf, size_t n, int fd, char *file)
+{
+	typedef double pixel_t[4];
+	size_t i, ptr;
+	double y, u, v, max = (double)UINT16_MAX;
+	double r, g, b;
+	pixel_t pixels[1024];
+	for (ptr = i = 0; ptr < n; ptr += 8, i++) {
+		pixels[i][3] = ntohs(((uint16_t *)(buf + ptr))[0]) / max;
+		y = ntohs(((uint16_t *)(buf + ptr))[1]) / max;
+		u = ntohs(((uint16_t *)(buf + ptr))[2]) / max;
+		v = ntohs(((uint16_t *)(buf + ptr))[3]) / max;
+		yuv_to_srgb(y, u, v, &r, &g, &b);
+		r = srgb_decode(r);
+		g = srgb_decode(g);
+		b = srgb_decode(b);
+		srgb_to_ciexyz(r, g, b, pixels[i] + 0, pixels[i] + 1, pixels[i] + 2);
+		if (++i == 1024) {
+			i = 0;
+			ewriteall(fd, pixels, sizeof(pixels), file);
+		}
+	}
+	if (i)
+		ewriteall(fd, pixels, i * sizeof(*pixels), file);
+}
+
+static void
+convert(char *infile, int outfd, char *outfile, size_t width, size_t height, char *frame_rate)
+{
+	char geometry[2 * 3 * sizeof(size_t) + 2], *cmd[13], buf[BUFSIZ];
+	int status, fd, pipe_rw[2];
+	size_t i = 0, n, ptr;
+	ssize_t r;
+	pid_t pid;
+
+	cmd[i++] = "ffmpeg";
+	cmd[i++] = "-i", cmd[i++] = "-";
+	cmd[i++] = "-f", cmd[i++] = "rawvideo";
+	cmd[i++] = "-pix_fmt", cmd[i++] = "ayuv64le";
+	if (width && height) {
+		sprintf(geometry, "%zux%zu", width, height);
+		cmd[i++] = "-s:v", cmd[i++] = geometry;
+	}
+	if (frame_rate)
+		cmd[i++] = "-r", cmd[i++] = frame_rate;
+	cmd[i++] = "-";
+	cmd[i++] = NULL;
+
+	if (pipe(pipe_rw))
+		eprintf("pipe:");
+
+	pid = fork();
+	if (pid == -1)
+		eprintf("fork:");
+
+	if (!pid) {
+#if defined(HAVE_PRCTL) && defined(PR_SET_PDEATHSIG)
+		prctl(PR_SET_PDEATHSIG, SIGKILL);
+#endif
+		fd = eopen(infile, O_RDONLY);
+		if (dup2(fd, STDIN_FILENO) == -1)
+			eprintf("dup2:");
+		close(fd);
+		close(pipe_rw[0]);
+		if (dup2(pipe_rw[1], STDOUT_FILENO) == -1)
+			eprintf("dup2:");
+		close(pipe_rw[1]);
+		execvp("ffmpeg", cmd);
+		eprintf("exec ffmpeg:");
+	}
+
+	close(pipe_rw[1]);
+
+	for (ptr = 0;;) {
+		r = read(pipe_rw[0], buf + ptr, sizeof(buf) - ptr);
+		if (r < 0)
+			eprintf("read <subprocess>:");
+		if (r == 0)
+			break;
+		ptr += (size_t)r;
+		n = ptr - (ptr % 8);
+		convert_segment(buf, n, outfd, outfile);
+		memmove(buf, buf + n, ptr -= n);
+	}
+	if (ptr)
+		eprintf("<subprocess>: incomplete frame\n");
+
+	close(pipe_rw[0]);
+	if (waitpid(pid, &status, 0) == -1)
+		eprintf("waitpid:");
+	if (status)
+		exit(1);
+}
+
+int
+main(int argc, char *argv[])
+{
+	size_t width = 0, height = 0, frames;
+	char head[STREAM_HEAD_MAX];
+	char *frame_rate = NULL;
+	char *infile;
+	char *outfile;
+	char *data;
+	ssize_t headlen;
+	size_t length;
+	int outfd, status;
+	pid_t pid;
+	struct stat st;
+
+	ARGBEGIN {
+	case 'r':
+		frame_rate = EARG();
+		break;
+	case 'w':
+		width = etozu_flag('w', EARG(), 1, SIZE_MAX);
+		break;
+	case 'h':
+		height = etozu_flag('h', EARG(), 1, SIZE_MAX);
+		break;
+	default:
+		usage();
+	} ARGEND;
+
+	if (argc != 2 || !width != !height)
+		usage();
+
+	infile = argv[0];
+	outfile = argv[1];
+
+	outfd = eopen(outfile, O_RDWR | O_CREAT | O_TRUNC, 0666);
+
+	pid = fork();
+	if (pid == -1)
+		eprintf("fork:");
+
+	if (!pid) {
+#if defined(HAVE_PRCTL) && defined(PR_SET_PDEATHSIG)
+		prctl(PR_SET_PDEATHSIG, SIGKILL);
+#endif
+		convert(infile, outfd, outfile, width, height, frame_rate);
+		exit(0);
+	}
+
+	get_metadata(infile, width ? NULL : &width, height ? NULL : &height, &frames);
+
+	if (waitpid(pid, &status, 0) == -1)
+		eprintf("waitpid:");
+	if (status)
+		exit(1);
+
+	if (fstat(outfd, &st))
+		eprintf("fstat %s:", outfile);
+	length = (size_t)(st.st_size);
+
+	sprintf(head, "%zu %zu %zu %s\n%cuivf%zn", frames, width, height, "xyza", 0, &headlen);
+	ewriteall(outfd, head, (size_t)headlen, outfile);
+	data = mmap(0, length + (size_t)headlen, PROT_READ | PROT_WRITE, MAP_PRIVATE, outfd, 0);
+	memmove(data + headlen, data, length);
+	memcpy(data, head, (size_t)headlen);
+	munmap(data, length + (size_t)headlen);
+
+	close(outfd);
+	return 0;
+}
diff --git a/src/vu-read-head.c b/src/vu-read-head.c
index 0050852..afb2fbc 100644
--- a/src/vu-read-head.c
+++ b/src/vu-read-head.c
@@ -10,7 +10,7 @@ USAGE("")
 int
 main(int argc, char *argv[])
 {
-	char buf[4 + 3 * 3 * sizeof(size_t) + sizeof(((struct stream *)0)->pixfmt)];
+	char buf[STREAM_HEAD_MAX];
 	char magic[] = {'\0', 'u', 'i', 'v', 'f'};
 	char b, *p;
 	size_t i, ptr;
diff --git a/src/vu-rewrite-head.c b/src/vu-rewrite-head.c
index 6609eb9..0de1716 100644
--- a/src/vu-rewrite-head.c
+++ b/src/vu-rewrite-head.c
@@ -15,7 +15,7 @@ USAGE("[-h] file [(frames | 'auto') [(width | 'same') (height | 'same') [format
 static void
 rewrite(struct stream *stream, int frames_auto)
 {
-	char head[3 * 3 * sizeof(size_t) + sizeof(stream->pixfmt) + 10];
+	char head[STREAM_HEAD_MAX];
 	ssize_t headlen;
 	size_t frame_size, frame_count, length;
 	struct stat st;
diff --git a/src/vu-split.c b/src/vu-split.c
index b1ed398..ff6ce38 100644
--- a/src/vu-split.c
+++ b/src/vu-split.c
@@ -47,6 +47,8 @@ main(int argc, char *argv[])
 	for (i = 0; i < parts; i++) {
 		fd = eopen(argv[i * 2], O_WRONLY | O_CREAT | O_TRUNC, 0666);
 		fp = fdopen(fd, "wb");
+		if (!fp)
+			eprintf("fdopen %s:", argv[i * 2]);
 
 		stream.frames = ends[i] - (i ? ends[i - 1] : 0);
 		fprint_stream_head(fp, &stream);