From ec7978e6b2942476a17acc8e67cdbd42e10c55a7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Andreas=20=C3=96man?= <andreas@lonelycoder.com>
Date: Wed, 1 Sep 2010 16:28:13 +0000
Subject: [PATCH] Write aspect ratio in recorded MKV files

---
 src/dvr/mkmux.c              |  7 ++++++
 src/packet.h                 |  3 +++
 src/parser_h264.c            | 42 +++++++++++++++++++++++++++++++++---
 src/parsers.c                | 31 ++++++++++++++++++++++++--
 src/plumbing/globalheaders.c | 21 ++++++++++++++----
 src/streaming.h              |  4 +++-
 src/tvhead.h                 |  3 +++
 7 files changed, 101 insertions(+), 10 deletions(-)

diff --git a/src/dvr/mkmux.c b/src/dvr/mkmux.c
index d6a71c40..cbb003fd 100644
--- a/src/dvr/mkmux.c
+++ b/src/dvr/mkmux.c
@@ -269,6 +269,13 @@ mk_build_tracks(mk_mux_t *mkm, const struct streaming_start *ss)
       ebml_append_uint(vi, 0xb0, ssc->ssc_width);
       ebml_append_uint(vi, 0xba, ssc->ssc_height);
 
+      if(ssc->ssc_aspect_num && ssc->ssc_aspect_den) {
+	ebml_append_uint(vi, 0x54b2, 3); // Display width/height is in DAR
+	ebml_append_uint(vi, 0x54b0, ssc->ssc_aspect_num);
+	ebml_append_uint(vi, 0x54ba, ssc->ssc_aspect_den);
+
+      }
+
       ebml_append_master(t, 0xe0, vi);
     }
 
diff --git a/src/packet.h b/src/packet.h
index 0398c112..0ed40ad5 100644
--- a/src/packet.h
+++ b/src/packet.h
@@ -50,6 +50,9 @@ typedef struct th_pkt {
   uint8_t pkt_channels;
   uint8_t pkt_sri;
 
+  uint16_t pkt_aspect_num;
+  uint16_t pkt_aspect_den;
+
   pktbuf_t *pkt_payload;
   pktbuf_t *pkt_header;
 
diff --git a/src/parser_h264.c b/src/parser_h264.c
index 71e6d775..19bc3757 100644
--- a/src/parser_h264.c
+++ b/src/parser_h264.c
@@ -99,6 +99,8 @@ typedef struct h264_private {
     char fixed_rate;
     int units_in_tick;
     int time_scale;
+    uint16_t aspect_num;
+    uint16_t aspect_den;
   } sps[256];
   
   struct {
@@ -109,14 +111,44 @@ typedef struct h264_private {
 
 
 
+static const uint16_t h264_aspect[17][2] = {
+ {0, 1},
+ {1, 1},
+ {12, 11},
+ {10, 11},
+ {16, 11},
+ {40, 33},
+ {24, 11},
+ {20, 11},
+ {32, 11},
+ {80, 33},
+ {18, 11},
+ {15, 11},
+ {64, 33},
+ {160,99},
+ {4, 3},
+ {3, 2},
+ {2, 1},
+};
+
 
 static int 
 decode_vui(h264_private_t *p, bitstream_t *bs, int sps_id)
 {
+  p->sps[sps_id].aspect_num = 0;
+  p->sps[sps_id].aspect_den = 1;
+
   if(read_bits1(bs)) {
-    if(read_bits(bs, 8) == 255) {
-      read_bits(bs, 16);
-      read_bits(bs, 16);
+    int aspect = read_bits(bs, 8);
+
+    if(aspect == 255) {
+      uint16_t num = read_bits(bs, 16);
+      uint16_t den = read_bits(bs, 16);
+      p->sps[sps_id].aspect_num = num;
+      p->sps[sps_id].aspect_den = den;
+    } else if(aspect < 17) {
+      p->sps[sps_id].aspect_num =  h264_aspect[aspect][0];
+      p->sps[sps_id].aspect_den =  h264_aspect[aspect][1];
     }
   }
 
@@ -357,5 +389,9 @@ h264_decode_slice_header(th_stream_t *st, bitstream_t *bs, int *pkttype,
     parser_set_stream_vsize(st, p->sps[sps_id].width, 
 			    p->sps[sps_id].height *
 			    (2 - p->sps[sps_id].mbs_only_flag));
+
+  st->st_aspect_num = p->sps[sps_id].aspect_num;
+  st->st_aspect_den = p->sps[sps_id].aspect_den;
+
   return 0;
 }
diff --git a/src/parsers.c b/src/parsers.c
index 5ae0c38f..124941f2 100644
--- a/src/parsers.c
+++ b/src/parsers.c
@@ -757,6 +757,26 @@ parser_set_stream_vsize(th_stream_t *st, int width, int height)
 }
 
 
+static const uint8_t mpeg2_aspect[16][2]={
+    {0,1},
+    {1,1},
+    {4,3},
+    {16,9},
+    {221,100},
+    {0,1},
+    {0,1},
+    {0,1},
+    {0,1},
+    {0,1},
+    {0,1},
+    {0,1},
+    {0,1},
+    {0,1},
+    {0,1},
+    {0,1},
+};
+
+
 /**
  * Parse mpeg2video sequence start
  */
@@ -764,14 +784,18 @@ static int
 parse_mpeg2video_seq_start(th_transport_t *t, th_stream_t *st,
 			   bitstream_t *bs)
 {
-  int v, width, height;
+  int v, width, height, aspect;
 
   if(bs->len < 61)
     return 1;
   
   width = read_bits(bs, 12);
   height = read_bits(bs, 12);
-  skip_bits(bs, 4);
+  aspect = read_bits(bs, 4);
+
+  st->st_aspect_num = mpeg2_aspect[aspect][0];
+  st->st_aspect_den = mpeg2_aspect[aspect][1];
+
   st->st_frame_duration = mpeg2video_framedurations[read_bits(bs, 4)];
 
   v = read_bits(bs, 18) * 400;
@@ -1143,6 +1167,9 @@ parser_deliver(th_transport_t *t, th_stream_t *st, th_pkt_t *pkt)
 	 pktbuf_len(pkt->pkt_payload));
 #endif
 
+  pkt->pkt_aspect_num = st->st_aspect_num;
+  pkt->pkt_aspect_den = st->st_aspect_den;
+
   //  avgstat_add(&st->st_rate, pkt->pkt_payloadlen, dispatch_clock);
 
   /**
diff --git a/src/plumbing/globalheaders.c b/src/plumbing/globalheaders.c
index c35f8d8d..d99c56c4 100644
--- a/src/plumbing/globalheaders.c
+++ b/src/plumbing/globalheaders.c
@@ -67,6 +67,13 @@ apply_header(streaming_start_component_t *ssc, th_pkt_t *pkt)
     ssc->ssc_sri      = pkt->pkt_sri;
   }
 
+  if(SCT_ISVIDEO(ssc->ssc_type)) {
+    if(pkt->pkt_aspect_num && pkt->pkt_aspect_den) {
+      ssc->ssc_aspect_num = pkt->pkt_aspect_num;
+      ssc->ssc_aspect_den = pkt->pkt_aspect_den;
+    }
+  }
+
   if(ssc->ssc_gh != NULL)
     return;
 
@@ -82,6 +89,7 @@ apply_header(streaming_start_component_t *ssc, th_pkt_t *pkt)
 
   case SCT_H264:
   case SCT_MPEG2VIDEO:
+
     if(pkt->pkt_header != NULL) {
       ssc->ssc_gh = pkt->pkt_header;
       pktbuf_ref_inc(ssc->ssc_gh);
@@ -95,12 +103,17 @@ apply_header(streaming_start_component_t *ssc, th_pkt_t *pkt)
  *
  */
 static int
-header_complete(streaming_start_component_t *ssc)
+header_complete(streaming_start_component_t *ssc, int not_so_picky)
 {
   if((SCT_ISAUDIO(ssc->ssc_type) || SCT_ISVIDEO(ssc->ssc_type)) &&
      ssc->ssc_frameduration == 0)
     return 0;
 
+  if(SCT_ISVIDEO(ssc->ssc_type)) {
+    if(!not_so_picky && (ssc->ssc_aspect_num == 0 || ssc->ssc_aspect_den == 0))
+      return 0;
+  }
+
   if(SCT_ISAUDIO(ssc->ssc_type) &&
      (ssc->ssc_sri == 0 || ssc->ssc_channels == 0))
     return 0;
@@ -121,16 +134,16 @@ headers_complete(globalheaders_t *gh, int64_t qd)
 {
   streaming_start_t *ss = gh->gh_ss;
   streaming_start_component_t *ssc;
-  int i;
+  int i, threshold = qd > (MAX_SCAN_TIME * 90);
 
   assert(ss != NULL);
  
   for(i = 0; i < ss->ss_num_components; i++) {
     ssc = &ss->ss_components[i];
 
-    if(!header_complete(ssc)) {
+    if(!header_complete(ssc, threshold)) {
 
-      if(qd > (MAX_SCAN_TIME * 90)) {
+      if(threshold) {
 	ssc->ssc_disabled = 1;
       } else {
 	return 0;
diff --git a/src/streaming.h b/src/streaming.h
index ac878e17..48b0c865 100644
--- a/src/streaming.h
+++ b/src/streaming.h
@@ -33,10 +33,12 @@ typedef struct streaming_start_component {
   uint16_t ssc_pid;
   int16_t ssc_width;
   int16_t ssc_height;
+  int16_t ssc_aspect_num;
+  int16_t ssc_aspect_den;
   uint8_t ssc_sri;
   uint8_t ssc_channels;
   uint8_t ssc_disabled;
-
+  
   pktbuf_t *ssc_gh;
 
   int ssc_frameduration;
diff --git a/src/tvhead.h b/src/tvhead.h
index f208936a..e374d389 100644
--- a/src/tvhead.h
+++ b/src/tvhead.h
@@ -380,6 +380,9 @@ typedef struct th_stream {
   streaming_component_type_t st_type;
   int st_index;
 
+  uint16_t st_aspect_num;
+  uint16_t st_aspect_den;
+
   char st_lang[4];           /* ISO 639 3-letter language code */
   uint16_t st_composition_id;
   uint16_t st_ancillary_id;