scripts/NHMLFixup.lua

   1 #!/usr/bin/env lua
   2 ----------------------------------------------------------------------------------------------------------------------
   3 ----------------------------------------------------------------------------------------------------------------------
   4 -- NHMLFixup v5 by Ilari (2010-09-18).
   5 -- Update timecodes in NHML Audio/Video track timing to conform to given MKV v2 timecodes file.
   6 -- Syntax: NHMLFixup <video-nhml-file> <audio-nhml-file> <mkv-timecodes-file> [delay=<delay>] [tvaspect]
   7 -- <delay> is number of milliseconds to delay the video (in order to compensate for audio codec delay, reportedly
   8 -- does not work right with some demuxers).
   9 -- The 'tvaspect' option makes video track to be automatically adjusted to '4:3' aspect ratio.
  10 --
  11 -- Version v5 by Ilari (2010-09-18):
  12 --      - Move the files first out of way, since rename-over fails on Windows.
  13 --
  14 -- Version v4 by Ilari (2010-09-17):
  15 --      - Change audio track ID if it collides with video track..
  16 --
  17 -- Version v3 by Ilari (2010-09-17):
  18 --      - Support setting aspect ratio correction.
  19 --
  20 -- Version v2 by Ilari (2010-09-16):
  21 --      - If sound and video NHMLs are wrong way around, automatically swap them.
  22 --      - Check that one of the NHMLs is sound and other is video.
  23 --
  24 ----------------------------------------------------------------------------------------------------------------------
  25 ----------------------------------------------------------------------------------------------------------------------
  26
  27 ----------------------------------------------------------------------------------------------------------------------
  28 -- Function reduce_fraction(number numerator, number denumerator)
  29 -- Returns reduced fraction.
  30 ----------------------------------------------------------------------------------------------------------------------
  31 reduce_fraction = function(numerator, denumerator)
  32         local x, y = numerator, denumerator;
  33         while y > 0 do
  34                 x, y = y, x % y;
  35         end
  36         return numerator / x, denumerator / x;
  37 end
  38
  39 ----------------------------------------------------------------------------------------------------------------------
  40 -- Function load_timecode_file(FILE file, number timescale)
  41 -- Loads timecode data from file @file, using timescale of @timescale frames per second. Returns array of scaled
  42 -- timecodes.
  43 ----------------------------------------------------------------------------------------------------------------------
  44 load_timecode_file = function(file, timescale)
  45         local line, ret;
  46         line = file:read("*l");
  47         if line ~= "# timecode format v2" then
  48                 error("Timecode file is not in MKV timecodes v2 format");
  49         end
  50         ret = {};
  51         while true do
  52                 line = file:read("*l");
  53                 if not line then
  54                         break;
  55                 end
  56                 local timecode = tonumber(line);
  57                 if not timecode then
  58                         error("Can't parse timecode '" .. line .. "'.");
  59                 end
  60                 table.insert(ret, math.floor(0.5 + timecode / 1000 * timescale));
  61         end
  62         return ret;
  63 end
  64
  65 ----------------------------------------------------------------------------------------------------------------------
  66 -- Function make_reverse_index_table(Array array)
  67 -- Returns table, that has entry for each entry in given array @array with value being rank of value, 1 being smallest
  68 -- and #array the largest. If @lookup is non-nil, values are looked up from that array.
  69 ----------------------------------------------------------------------------------------------------------------------
  70 make_reverse_index_table = function(array, lookup)
  71         local sorted, ret;
  72         local i;
  73         sorted = {};
  74         for i = 1,#array do
  75                 sorted[i] = array[i];
  76         end
  77         table.sort(sorted);
  78         ret = {};
  79         for i = 1,#sorted do
  80                 ret[sorted[i]] = (lookup and lookup[i]) or i;
  81         end
  82         return ret;
  83 end
  84
  85 ----------------------------------------------------------------------------------------------------------------------
  86 -- Function max_causality_violaton(Array CTS, Array DTS)
  87 -- Return the maximum number of time units CTS and DTS values violate causality. #CTS must equal #DTS.
  88 ----------------------------------------------------------------------------------------------------------------------
  89 max_causality_violation = function(CTS, DTS)
  90         local max_cv = 0;
  91         local i;
  92         for i = 1,#CTS do
  93                 max_cv = math.max(max_cv, DTS[i] - CTS[i]);
  94         end
  95         return max_cv;
  96 end
  97
  98 ----------------------------------------------------------------------------------------------------------------------
  99 -- Function fixup_video_times(Array sampledata, Array timecodes, Number spec_delay)
 100 -- Fixes video timing of @sampledata (fields CTS and DTS) to be consistent with timecodes in @timecodes. Returns the
 101 -- CTS offset of first sample (for fixing audio). @spec_delay is special delay to add (to fix A/V sync).
 102 ----------------------------------------------------------------------------------------------------------------------
 103 fixup_video_times = function(sampledata, timecodes, spec_delay)
 104         local cts_tab = {};
 105         local dts_tab = {};
 106         local k, v, i;
 107         if #sampledata ~= #timecodes then
 108                 error("Number of samples (" .. #sampledata .. ") does not match number of timecodes (" .. #timecodes
 109                         .. ").");
 110         end
 111         for i = 1,#sampledata do
 112                 cts_tab[i] = sampledata[i].CTS;
 113                 dts_tab[i] = sampledata[i].DTS;
 114         end
 115         cts_lookup = make_reverse_index_table(cts_tab, timecodes);
 116         dts_lookup = make_reverse_index_table(dts_tab, timecodes);
 117
 118         -- Perform time translation and find max causality violation.
 119         local max_cv = 0;
 120         for i = 1,#sampledata do
 121                 sampledata[i].CTS = cts_lookup[sampledata[i].CTS];
 122                 sampledata[i].DTS = dts_lookup[sampledata[i].DTS];
 123                 max_cv = math.max(max_cv, sampledata[i].DTS - sampledata[i].CTS);
 124         end
 125         -- Add maximum causality violation to CTS to eliminate the causality violations.
 126         -- Also find the minimum CTS.
 127         local min_cts = 999999999999999999999;
 128         for i = 1,#sampledata do
 129                 sampledata[i].CTS = sampledata[i].CTS + max_cv + spec_delay;
 130                 --Spec_delay should not apply to audio.
 131                 min_cts = math.min(min_cts, sampledata[i].CTS - spec_delay);
 132         end
 133         return min_cts;
 134 end
 135
 136 ----------------------------------------------------------------------------------------------------------------------
 137 -- Function fixup_video_times(Array sampledata, Number min_video_cts, Number video_timescale, Number audio_timescale)
 138 -- Fixes video timing of @sampledata (field CTS) to be consistent with video minimum CTS of @cts. Video timescale
 139 -- is assumed to be @video_timescale and audio timescale @audio_timescale.
 140 ----------------------------------------------------------------------------------------------------------------------
 141 fixup_audio_times = function(sampledata, min_video_cts, video_timescale, audio_timescale)
 142         local fixup = math.floor(0.5 + min_video_cts * audio_timescale / video_timescale);
 143         for i = 1,#sampledata do
 144                 sampledata[i].CTS = sampledata[i].CTS + fixup;
 145         end
 146 end
 147
 148 ----------------------------------------------------------------------------------------------------------------------
 149 -- Function translate_NHML_TS_in(Array sampledata);
 150 -- Translate NHML CTSOffset fields in @sampledata into CTS fields.
 151 ----------------------------------------------------------------------------------------------------------------------
 152 translate_NHML_TS_in = function(sampledata, default_dDTS)
 153         local i;
 154         local dts = 0;
 155         for i = 1,#sampledata do
 156                 if not sampledata[i].DTS then
 157                         sampledata[i].DTS = dts + default_dDTS;
 158                 end
 159                 dts = sampledata[i].DTS;
 160                 if sampledata[i].CTSOffset then
 161                         sampledata[i].CTS = sampledata[i].CTSOffset + sampledata[i].DTS;
 162                 else
 163                         sampledata[i].CTS = sampledata[i].DTS;
 164                 end
 165         end
 166 end
 167
 168 ----------------------------------------------------------------------------------------------------------------------
 169 -- Function translate_NHML_TS_out(Array sampledata);
 170 -- Translate CTS fields in @sampledata into NHML CTSOffset fields.
 171 ----------------------------------------------------------------------------------------------------------------------
 172 translate_NHML_TS_out = function(sampledata)
 173         local i;
 174         for i = 1,#sampledata do
 175                 sampledata[i].CTSOffset = sampledata[i].CTS - sampledata[i].DTS;
 176                 if sampledata[i].CTSOffset < 0 then
 177                         error("INTERNAL ERROR: translate_NHML_TS_out: Causality violation: CTS=" .. tostring(
 178                                 sampledata[i].CTS) .. " DTS=" .. tostring(sampledata[i].DTS) .. ".");
 179                 end
 180                 sampledata[i].CTS = nil;
 181         end
 182 end
 183
 184 ----------------------------------------------------------------------------------------------------------------------
 185 -- Function map_table_to_number(Table tab);
 186 -- Translate all numeric fields in table @tab into numbers.
 187 ----------------------------------------------------------------------------------------------------------------------
 188 map_table_to_number = function(tab)
 189         local k, v;
 190         for k, v in pairs(tab) do
 191                 local n = tonumber(v);
 192                 if n then
 193                         tab[k] = n;
 194                 end
 195         end
 196 end
 197
 198 ----------------------------------------------------------------------------------------------------------------------
 199 -- Function map_fields_to_number(Array sampledata);
 200 -- Translate all numeric fields in array @sampledata into numbers.
 201 ----------------------------------------------------------------------------------------------------------------------
 202 map_fields_to_number = function(sampledata)
 203         local i;
 204         for i = 1,#sampledata do
 205                 map_table_to_number(sampledata[i]);
 206         end
 207 end
 208
 209 ----------------------------------------------------------------------------------------------------------------------
 210 -- Function escape_xml_text(String str)
 211 -- Return XML escaping of text str.
 212 ----------------------------------------------------------------------------------------------------------------------
 213 escape_xml_text = function(str)
 214         str = string.gsub(str, "&", "&amp;");
 215         str = string.gsub(str, "<", "&lt;");
 216         str = string.gsub(str, ">", "&gt;");
 217         str = string.gsub(str, "\"", "&quot;");
 218         str = string.gsub(str, "\'", "&apos;");
 219         return str;
 220 end
 221
 222 ----------------------------------------------------------------------------------------------------------------------
 223 -- Function escape_xml_text(String str)
 224 -- Return XML unescaping of text str.
 225 ----------------------------------------------------------------------------------------------------------------------
 226 unescape_xml_text = function(str)
 227         str = string.gsub(str, "&apos;", "\'");
 228         str = string.gsub(str, "&quot;", "\"");
 229         str = string.gsub(str, "&gt;", ">");
 230         str = string.gsub(str, "&lt;", "<");
 231         str = string.gsub(str, "&amp;", "&");
 232         return str;
 233 end
 234
 235 ----------------------------------------------------------------------------------------------------------------------
 236 -- Function serialize_table_to_xml_entity(File file, String tag, Table data, bool noclose);
 237 -- Write @data as XML start tag of type @tag into @file. If noclose is true, then tag will not be closed.
 238 ----------------------------------------------------------------------------------------------------------------------
 239 serialize_table_to_xml_entity = function(file, tag, data, noclose)
 240         local k, v;
 241         file:write("<" .. tag .. " ");
 242         for k, v in pairs(data) do
 243                 file:write(k .. "=\"" .. escape_xml_text(tostring(v)) .. "\" ");
 244         end
 245         if noclose then
 246                 file:write(">\n");
 247         else
 248                 file:write("/>\n");
 249         end
 250 end
 251
 252 ----------------------------------------------------------------------------------------------------------------------
 253 -- Function serialize_array_to_xml_entity(File file, String tag, Array data);
 254 -- Write each element of @data as empty XML tag of type @tag into @file.
 255 ----------------------------------------------------------------------------------------------------------------------
 256 serialize_array_to_xml_entity = function(file, tag, data)
 257         local i;
 258         for i = 1,#data do
 259                 serialize_table_to_xml_entity(file, tag, data[i]);
 260         end
 261 end
 262
 263 ----------------------------------------------------------------------------------------------------------------------
 264 -- Function write_NHML_data(File file, Table header, Table sampledata)
 265 -- Write entiere NHML file.
 266 ----------------------------------------------------------------------------------------------------------------------
 267 write_NHML_data = function(file, header, sampledata)
 268         file:write("<?xml version=\"1.0\" encoding=\"UTF-8\" standalone=\"no\"?>\n");
 269         serialize_table_to_xml_entity(file, "NHNTStream", header, true);
 270         serialize_array_to_xml_entity(file, "NHNTSample", sampledata);
 271         file:write("</NHNTStream>\n");
 272 end
 273
 274 ----------------------------------------------------------------------------------------------------------------------
 275 -- Function open_file_checked(String file, String mode, bool for_write)
 276 -- Return file handle to file (checking that open succeeds).
 277 ----------------------------------------------------------------------------------------------------------------------
 278 open_file_checked = function(file, mode, for_write)
 279         local a, b;
 280         a, b = io.open(file, mode);
 281         if not a then
 282                 error("Can't open '" .. file .. "': " .. b);
 283         end
 284         return a;
 285 end
 286
 287 ----------------------------------------------------------------------------------------------------------------------
 288 -- Function call_with_file(fun, string file, String mode, ...);
 289 -- Call fun with opened file handle to @file (in mode @mode) as first parameter.
 290 ----------------------------------------------------------------------------------------------------------------------
 291 call_with_file = function(fun, file, mode, ...)
 292         -- FIXME: Handle nils returned from function.
 293         local handle = open_file_checked(file, mode);
 294         local ret = {fun(handle, ...)};
 295         handle:close();
 296         return unpack(ret);
 297 end
 298
 299 ----------------------------------------------------------------------------------------------------------------------
 300 -- Function xml_parse_tag(String line);
 301 -- Returns the xml tag type for @line plus table of attributes.
 302 ----------------------------------------------------------------------------------------------------------------------
 303 xml_parse_tag = function(line)
 304         -- More regexping...
 305         local tagname;
 306         local attr = {};
 307         tagname = string.match(line, "<(%S+).*>");
 308         if not tagname then
 309                 error("'" .. line .. "': Parse error.");
 310         end
 311         local k, v;
 312         for k, v in string.gmatch(line, "([^ =]+)=\"([^\"]*)\"") do
 313                 attr[k] = unescape_xml_text(v);
 314         end
 315         return tagname, attr;
 316 end
 317
 318 ----------------------------------------------------------------------------------------------------------------------
 319 -- Function load_NHML(File file)
 320 -- Loads NHML file @file. Returns header table and samples array.
 321 ----------------------------------------------------------------------------------------------------------------------
 322 load_NHML = function(file)
 323         -- Let's regexp this shit...
 324         local header = {};
 325         local samples = {};
 326         while true do
 327                 local line = file:read();
 328                 if not line then
 329                         error("Unexpected end of NHML file.");
 330                 end
 331                 local xtag, attributes;
 332                 xtag, attributes = xml_parse_tag(line);
 333                 if xtag == "NHNTStream" then
 334                         header = attributes;
 335                 elseif xtag == "NHNTSample" then
 336                         table.insert(samples, attributes);
 337                 elseif xtag == "/NHNTStream" then
 338                         break;
 339                 elseif xtag == "?xml" then
 340                 else
 341                         print("WARNING: Unrecognized tag '" .. xtag .. "'.");
 342                 end
 343         end
 344         return header, samples;
 345 end
 346
 347 ----------------------------------------------------------------------------------------------------------------------
 348 -- Function reame_errcheck(String old, String new)
 349 -- Rename old to new. With error checking.
 350 ----------------------------------------------------------------------------------------------------------------------
 351 rename_errcheck = function(old, new, backup)
 352         local a, b;
 353         os.remove(backup);
 354         a, b = os.rename(new, backup);
 355         if not a then
 356                 error("Can't rename '" .. new .. "' -> '" .. backup .. "': " .. b);
 357         end
 358         a, b = os.rename(old, new);
 359         if not a then
 360                 error("Can't rename '" .. old .. "' -> '" .. new .. "': " .. b);
 361         end
 362 end
 363
 364
 365 if #arg < 3 then
 366         error("Syntax: NHMLFixup.lua <video.nhml> <audio.nhml> <timecodes.txt> [delay=<delay>] [tvaspect]");
 367 end
 368
 369 -- Load the NHML files.
 370 io.stdout:write("Loading '" .. arg[1] .. "'..."); io.stdout:flush();
 371 video_header, video_samples = call_with_file(load_NHML, arg[1], "r");
 372 io.stdout:write("Done.\n");
 373 io.stdout:write("Loading '" .. arg[2] .. "'..."); io.stdout:flush();
 374 audio_header, audio_samples = call_with_file(load_NHML, arg[2], "r");
 375 io.stdout:write("Done.\n");
 376 io.stdout:write("String to number conversion on video header..."); io.stdout:flush();
 377 map_table_to_number(video_header);
 378 io.stdout:write("Done.\n");
 379 io.stdout:write("String to number conversion on video samples..."); io.stdout:flush();
 380 map_fields_to_number(video_samples);
 381 io.stdout:write("Done.\n");
 382 io.stdout:write("String to number conversion on audio header..."); io.stdout:flush();
 383 map_table_to_number(audio_header);
 384 io.stdout:write("Done.\n");
 385 io.stdout:write("String to number conversion on audio samples..."); io.stdout:flush();
 386 map_fields_to_number(audio_samples);
 387 io.stdout:write("Done.\n");
 388 if video_header.streamType == 4 and audio_header.streamType == 5 then
 389         -- Ok.
 390 elseif video_header.streamType == 5 and audio_header.streamType == 4 then
 391         print("WARNING: You got audio and video wrong way around. Swapping them for you...");
 392         audio_header,audio_samples,arg[2],video_header,video_samples,arg[1] =
 393                 video_header,video_samples,arg[1],audio_header,audio_samples,arg[2];
 394 else
 395         error("Expected one video track and one audio track");
 396 end
 397
 398 if video_header.trackID == audio_header.trackID then
 399         print("WARNING: Audio and video have the same track id. Assigning new track id to audio track...");
 400         audio_header.trackID = audio_header.trackID + 1;
 401 end
 402
 403 io.stdout:write("Computing CTS for video samples..."); io.stdout:flush();
 404 translate_NHML_TS_in(video_samples, video_header.DTS_increment or 0);
 405 io.stdout:write("Done.\n");
 406 io.stdout:write("Computing CTS for audio samples..."); io.stdout:flush();
 407 translate_NHML_TS_in(audio_samples, audio_header.DTS_increment or 0);
 408 io.stdout:write("Done.\n");
 409
 410 -- Alter timescale if needed and load the timecode data.
 411 delay = 0;
 412 rdelay = 0;
 413 for i = 4,#arg do
 414         if arg[i] == "tvaspect" then
 415                 do_aspect_fixup = true;
 416         elseif string.sub(arg[i], 1, 6) == "delay=" then
 417                 local n = tonumber(string.sub(arg[i], 7, #(arg[i])));
 418                 if not n then
 419                         error("Bad delay.");
 420                 end
 421                 rdelay = n;
 422                 delay = math.floor(0.5 + rdelay / 1000 * video_header.timeScale);
 423         end
 424 end
 425
 426
 427 timecode_data = call_with_file(load_timecode_file, arg[3], "r", video_header.timeScale);
 428 MAX_MP4BOX_TIMECODE = 0x7FFFFFF;
 429 if timecode_data[#timecode_data] > MAX_MP4BOX_TIMECODE then
 430         -- Workaround MP4Box bug.
 431         divider = math.ceil(timecode_data[#timecode_data] / MAX_MP4BOX_TIMECODE);
 432         print("Notice: Dividing timecodes by " .. divider .. " to workaround MP4Box timecode bug.");
 433         io.stdout:write("Performing division..."); io.stdout:flush();
 434         video_header.timeScale = video_header.timeScale / divider;
 435         for i = 1,#timecode_data do
 436                 timecode_data[i] = timecode_data[i] / divider;
 437         end
 438         --Recompute delay.
 439         delay = math.floor(0.5 + rdelay / 1000 * video_header.timeScale);
 440         io.stdout:write("Done.\n");
 441 end
 442
 443 -- Do the actual fixup.
 444 io.stdout:write("Fixing up video timecodes..."); io.stdout:flush();
 445 audio_fixup = fixup_video_times(video_samples, timecode_data, delay);
 446 io.stdout:write("Done.\n");
 447 io.stdout:write("Fixing up audio timecodes..."); io.stdout:flush();
 448 fixup_audio_times(audio_samples, audio_fixup, video_header.timeScale, audio_header.timeScale);
 449 io.stdout:write("Done.\n");
 450
 451 if do_aspect_fixup then
 452         video_header.parNum, video_header.parDen = reduce_fraction(4 * video_header.height, 3 * video_header.width);
 453 end
 454
 455 -- Save the NHML files.
 456 io.stdout:write("Computing CTSOffset for video samples..."); io.stdout:flush();
 457 translate_NHML_TS_out(video_samples);
 458 io.stdout:write("Done.\n");
 459 io.stdout:write("Computing CTSOffset for audio samples..."); io.stdout:flush();
 460 translate_NHML_TS_out(audio_samples);
 461 io.stdout:write("Done.\n");
 462 io.stdout:write("Saving '" .. arg[1] .. ".tmp'..."); io.stdout:flush();
 463 call_with_file(write_NHML_data, arg[1] .. ".tmp", "w", video_header, video_samples);
 464 io.stdout:write("Done.\n");
 465 io.stdout:write("Saving '" .. arg[2] .. ".tmp'..."); io.stdout:flush();
 466 call_with_file(write_NHML_data, arg[2] .. ".tmp", "w", audio_header, audio_samples);
 467 io.stdout:write("Done.\n");
 468 io.stdout:write("Renaming '" .. arg[1] .. ".tmp' -> '" .. arg[1] .. "'..."); io.stdout:flush();
 469 rename_errcheck(arg[1] .. ".tmp", arg[1], arg[1] .. ".bak");
 470 io.stdout:write("Done.\n");
 471 io.stdout:write("Renaming '" .. arg[2] .. ".tmp' -> '" .. arg[2] .. "'..."); io.stdout:flush();
 472 rename_errcheck(arg[2] .. ".tmp", arg[2], arg[2] .. ".bak");
 473 io.stdout:write("Done.\n");
 474 io.stdout:write("All done.\n");