yuv4mpeg.cpp

   1 #include "yuv4mpeg.hpp"
   2 #include <sstream>
   3 #include <stdexcept>
   4 #include <limits>
   5 #include <fcntl.h>
   6
   7 namespace
   8 {
   9         std::string extract_tok(const std::string& str, size_t& ptr)
  10         {
  11                 size_t eptr = ptr;
  12                 size_t strl = str.length();
  13                 while(eptr < strl) {
  14                         char ch = str[eptr];
  15                         if(ch == '\t' || ch == '\v' || ch == '\r' || ch == ' ')
  16                                 break;
  17                         eptr++;
  18                 }
  19                 std::string tmp = str;
  20                 tmp = tmp.substr(ptr + 1, eptr - ptr - 1);
  21                 ptr = eptr;
  22                 return tmp;
  23         }
  24
  25         template<typename T>
  26         T tonumber(const std::string& str)
  27         {
  28                 T r = 0;
  29                 size_t strl = str.length();
  30                 if(str == "")
  31                         throw std::runtime_error("Invalid number '" + str + "'");
  32                 for(size_t i = 0; i < strl; i++) {
  33                         if(str[i] < '0' || str[i] > '9')
  34                                 throw std::runtime_error("Invalid number '" + str + "'");
  35                         if(r > std::numeric_limits<T>::max() / 10)
  36                                 throw std::runtime_error("Invalid number '" + str + "'");
  37                         r *= 10;
  38                         if(r + (str[i] - '0') < r)
  39                                 throw std::runtime_error("Invalid number '" + str + "'");
  40                         r += (str[i] - '0');
  41                 }
  42                 return r;
  43         }
  44
  45         void parse_ratio(const std::string& str, uint32_t& n, uint32_t& d)
  46         {
  47                 std::string _str = str;
  48                 size_t p = str.find_first_of(':');
  49                 if(p > str.length())
  50                         throw std::runtime_error("Invalid ratio '" + str + "'");
  51                 n = tonumber<uint32_t>(_str.substr(0, p));
  52                 d = tonumber<uint32_t>(_str.substr(p + 1));
  53         }
  54
  55         std::string read_line(FILE* pipe)
  56         {
  57                 std::ostringstream s;
  58                 int ch;
  59                 while(true) {
  60                         ch = fgetc(pipe);
  61                         if(ch == '\n' || ch < 0)
  62                                 return s.str();
  63                         s << static_cast<char>(ch);
  64                 }
  65         }
  66 }
  67
  68 yuv4mpeg_stream_header::yuv4mpeg_stream_header()
  69 {
  70         width = 0;
  71         height = 0;
  72         chroma = "420jpeg";
  73         interlace = '?';
  74         fps_n = 0;
  75         fps_d = 0;
  76         sar_n = 0;
  77         sar_d = 0;
  78 }
  79
  80 yuv4mpeg_stream_header::yuv4mpeg_stream_header(const std::string& str)
  81 {
  82         bool seen_width = false;
  83         bool seen_height = false;
  84         width = 0;
  85         height = 0;
  86         chroma = "420jpeg";
  87         interlace = '?';
  88         fps_n = 0;
  89         fps_d = 0;
  90         sar_n = 0;
  91         sar_d = 0;
  92         if(str.length() < 9 || str.substr(0, 9) != "YUV4MPEG2")
  93                 throw std::runtime_error("Bad YUV4MPEG2 stream magic");
  94         size_t ptr = 9;
  95         size_t strl = str.length();
  96         std::string tok;
  97         while(ptr < strl) {
  98                 char ch = str[ptr];
  99                 switch(ch) {
 100                 case '\t':
 101                 case '\v':
 102                 case '\r':
 103                 case ' ':
 104                         //Skip.
 105                         ptr++;
 106                         break;
 107                 case 'A':
 108                         tok = extract_tok(str, ptr);
 109                         parse_ratio(tok, sar_n, sar_d);
 110                         break;
 111                 case 'C':
 112                         chroma = extract_tok(str, ptr);
 113                         break;
 114                 case 'F':
 115                         tok = extract_tok(str, ptr);
 116                         parse_ratio(tok, fps_n, fps_d);
 117                         break;
 118                 case 'H':
 119                         tok = extract_tok(str, ptr);
 120                         height = tonumber<size_t>(tok);
 121                         seen_height = true;
 122                         break;
 123                 case 'I':
 124                         tok = extract_tok(str, ptr);
 125                         if(tok.length() != 1)
 126                                 throw std::runtime_error("Invalid stream I-token");
 127                         interlace = tok[0];
 128                         break;
 129                 case 'W':
 130                         tok = extract_tok(str, ptr);
 131                         width = tonumber<size_t>(tok);
 132                         seen_width = true;
 133                         break;
 134                 case 'X':
 135                         tok = extract_tok(str, ptr);
 136                         if(tok == "")
 137                                 throw std::runtime_error("Empty extension token not allowed");
 138                         extension.push_back(tok);
 139                         break;
 140                 default:
 141                         throw std::runtime_error("Unknown token type '" + std::string(&ch, 1) + "'");
 142                 }
 143         }
 144         if(!seen_width || !seen_height)
 145                 throw std::runtime_error("W and H tokens are required");
 146 }
 147
 148 yuv4mpeg_stream_header::yuv4mpeg_stream_header(FILE* pipe)
 149 {
 150         *this = yuv4mpeg_stream_header(read_line(pipe));
 151 }
 152
 153 yuv4mpeg_stream_header::operator std::string() const
 154 {
 155         std::ostringstream h;
 156         h << "YUV4MPEG2 W" << width << " H" << height;
 157         if(chroma != "420jpeg")
 158                 h << " C" << chroma;
 159         if(interlace != '?')
 160                 h << " I" << interlace;
 161         if(fps_n || fps_d)
 162                 h << " F" << fps_n << ":" << fps_d;
 163         if(sar_n || sar_d)
 164                 h << " A" << sar_n << ":" << sar_d;
 165         for(auto i : extension)
 166                 h << " X" << i;
 167         return h.str();
 168 }
 169
 170 yuv4mpeg_frame_header::yuv4mpeg_frame_header()
 171 {
 172         representation = 0;
 173         temporal = 0;
 174         chroma = 0;
 175         duration = 1;
 176 }
 177
 178 yuv4mpeg_frame_header::yuv4mpeg_frame_header(const std::string& str)
 179 {
 180         representation = 0;
 181         temporal = 0;
 182         chroma = 0;
 183         duration = 1;
 184         if(str.length() < 5 || str.substr(0, 5) != "FRAME")
 185                 throw std::runtime_error("Bad YUV4MPEG2 frame magic");
 186         size_t ptr = 5;
 187         size_t strl = str.length();
 188         std::string tok;
 189         while(ptr < strl) {
 190                 char ch = str[ptr];
 191                 switch(ch) {
 192                 case '\t':
 193                 case '\v':
 194                 case '\r':
 195                 case ' ':
 196                         //Skip.
 197                         ptr++;
 198                         break;
 199                 case 'I':
 200                         tok = extract_tok(str, ptr);
 201                         if(tok.length() != 3)
 202                                 throw std::runtime_error("Invalid frame I-token");
 203                         representation = tok[0];
 204                         temporal = tok[1];
 205                         chroma = tok[2];
 206                         break;
 207                 case 'X':
 208                         tok = extract_tok(str, ptr);
 209                         if(tok == "")
 210                                 throw std::runtime_error("Empty extension token not allowed");
 211                         if(tok.length() > 9 && tok.substr(0, 9) == "duration=")
 212                                 duration = tonumber<uint32_t>(tok.substr(9));
 213                         else
 214                                 extension.push_back(tok);
 215                         break;
 216                 default:
 217                         throw std::runtime_error("Unknown token type '" + std::string(&ch, 1) + "'");
 218                 }
 219         }
 220 }
 221
 222 yuv4mpeg_frame_header::yuv4mpeg_frame_header(FILE* pipe)
 223 {
 224         *this = yuv4mpeg_frame_header(read_line(pipe));
 225 }
 226
 227 yuv4mpeg_frame_header::operator std::string() const
 228 {
 229         std::ostringstream h;
 230         h << "FRAME";
 231         if(representation || temporal || chroma)
 232                 h << " I" << representation << temporal << chroma;
 233         if(duration > 1)
 234                 h << " Xduration=" << duration;
 235         for(auto i : extension)
 236                 if(!regex_match("duration=.*", i))
 237                         h << " X" << i;
 238         return h.str();
 239 }
 240
 241 void mark_pipe_as_binary(FILE* pipe)
 242 {
 243 #if defined(_WIN32) || defined(_WIN64)
 244         setmode(fileno(pipe), O_BINARY);
 245 #endif
 246 }
 247
 248 void read_or_die(FILE* pipe, void* data, size_t datasize)
 249 {
 250         char* _data = reinterpret_cast<char*>(data);
 251         while(datasize > 0) {
 252                 size_t r;
 253                 size_t s = (datasize > 4096) ? 4096 : datasize;
 254                 r = fread(_data, 1, s, pipe);
 255                 if(r < s)
 256                         throw std::runtime_error("Error reading from file");
 257                 datasize -= r;
 258                 _data += r;
 259         }
 260 }
 261
 262 bool read_or_die2(FILE* pipe, void* data, size_t datasize)
 263 {
 264         bool first = true;
 265         char* _data = reinterpret_cast<char*>(data);
 266         while(datasize > 0) {
 267                 size_t r;
 268                 size_t s = (datasize > 4096) ? 4096 : datasize;
 269                 r = fread(_data, 1, s, pipe);
 270                 if(r == 0 && first)
 271                         return false;
 272                 if(r < s)
 273                         throw std::runtime_error("Error reading from file");
 274                 datasize -= r;
 275                 _data += r;
 276                 first = false;
 277         }
 278         return true;
 279 }
 280
 281 void write_or_die(FILE* pipe, const void* data, size_t datasize)
 282 {
 283         const char* _data = reinterpret_cast<const char*>(data);
 284         while(datasize > 0) {
 285                 size_t r;
 286                 size_t s = (datasize > 4096) ? 4096 : datasize;
 287                 r = fwrite(_data, 1, s, pipe);
 288                 if(r < s)
 289                         throw std::runtime_error("Error writing to file");
 290                 datasize -= r;
 291                 _data += r;
 292         }
 293 }
 294
 295 void write_or_die(FILE* pipe, const std::string& data)
 296 {
 297         std::vector<char> v;
 298         v.resize(data.size() + 1);
 299         std::copy(data.begin(), data.end(), v.begin());
 300         v[v.size() - 1] = '\n';
 301         write_or_die(pipe, &v[0], v.size());
 302 }
 303
 304 void get_512_scaling(size_t w, size_t h, uint32_t sar_n, uint32_t sar_d, size_t& w512, size_t& h512)
 305 {
 306         if(!sar_n || !sar_d) {
 307                 w512 = w;
 308                 h512 = h;
 309                 return;
 310         }
 311         if(sar_d < sar_n) {
 312                 //w / h < sar_n / sar_d, widen the screen.
 313                 w512 = static_cast<uint64_t>(w) * sar_n / sar_d;
 314                 h512 = h;
 315                 if(w512 & 1)
 316                         w512++;
 317         } else if(sar_d > sar_n) {
 318                 //w / h > sar_n / sar_d, tallen the screen.
 319                 w512 = w;
 320                 h512 = static_cast<uint64_t>(h) * sar_d / sar_n;
 321                 if(h512 & 1)
 322                         h512++;
 323         } else {
 324                 w512 = w;
 325                 h512 = h;
 326         }
 327 }
 328
 329 bool read_line2(FILE* pipe, std::string& line)
 330 {
 331         std::ostringstream s;
 332         int ch;
 333         bool first = true;
 334         while(true) {
 335                 ch = fgetc(pipe);
 336                 if(ch == '\n' || ch < 0)  {
 337                         line = s.str();
 338                         return (ch >= 0 || !first);
 339                 }
 340                 s << static_cast<char>(ch);
 341                 first = false;
 342         }
 343 }
 344
 345 regex_results yuv4mpeg_find_extension(const std::vector<std::string>& exts, const std::string& regexp)
 346 {
 347         for(auto i : exts) {
 348                 regex_results r;
 349                 if(r = regex(regexp, i))
 350                         return r;
 351         }
 352         return regex_results();
 353 }
 354
 355 void yuv4mpeg_delete_extensions(std::vector<std::string>& exts, const std::string& regexp)
 356 {
 357         bool redo = false;
 358         do {
 359                 redo = false;
 360                 for(auto i = exts.begin(); i != exts.end(); i++)
 361                         if(regex_match(regexp, *i)) {
 362                                 exts.erase(i);
 363                                 redo = true;
 364                                 break;
 365                         }
 366         } while(redo);
 367 }