frameshot.c

   1 /*****************************************************************************
   2 * frameshot: a frame-accurate screenshot generator.
   3 *****************************************************************************
   4 * Copyright (C) 2009
   5 *
   6 * Authors: Nathan Caldwell <saintdev@gmail.com>
   7 *          Various parts taken from x264.
   8 *
   9 * This program is free software; you can redistribute it and/or modify
  10 * it under the terms of the GNU General Public License as published by
  11 * the Free Software Foundation; either version 2 of the License, or
  12 * (at your option) any later version.
  13 *
  14 * This program is distributed in the hope that it will be useful,
  15 * but WITHOUT ANY WARRANTY; without even the implied warranty of
  16 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  17 * GNU General Public License for more details.
  18 *
  19 * You should have received a copy of the GNU General Public License
  20 * along with this program; if not, write to the Free Software
  21 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02111, USA.
  22 *****************************************************************************/
  23
  24 #include <stdint.h>
  25 #include <stdlib.h>
  26 #include <stdio.h>
  27 #include <string.h>
  28 #include <unistd.h>
  29 #include <sys/stat.h>
  30 #include <sys/types.h>
  31 #include <limits.h>
  32
  33 #include <zlib.h>
  34
  35
  36 #define _GNU_SOURCE
  37 #include <getopt.h>
  38
  39 #include "common.h"
  40 #include "utils.h"
  41 #include "output.h"
  42 #include "input.h"
  43
  44 enum
  45 {
  46         FORMAT_UNKNOWN,
  47         FORMAT_H264,
  48         FORMAT_DIRAC,
  49         FORMAT_OGG,
  50         FORMAT_M4V
  51 };
  52
  53 typedef struct {
  54         char *psz_outdir;
  55         int i_zlevel;
  56         hnd_t hin;
  57 } cli_opt_t;
  58
  59 /* input file function pointers */
  60 int (*p_open_infile)( char *psz_filename, hnd_t *p_handle, config_t *p_config );
  61 int (*p_read_frame)( hnd_t handle, picture_t *p_pic, int i_frame );
  62 int (*p_close_infile)( hnd_t handle );
  63
  64 /* output file function pointers */
  65 static int (*p_open_outfile)( char *psz_filename, hnd_t *p_handle, int i_compression );
  66 // static int (*p_set_outfile_param)( hnd_t handle, config_t *p_config );
  67 static int (*p_write_image)( hnd_t handle, picture_t *p_pic, config_t *p_config );
  68 static int (*p_close_outfile)( hnd_t handle );
  69
  70 static int parse_options( int argc, char **argv, config_t *config, cli_opt_t *opt );
  71 static int grab_frames( config_t *config, cli_opt_t *opt );
  72
  73 int main(int argc, char **argv)
  74 {
  75         config_t config;
  76         cli_opt_t opt;
  77         int ret = 0;
  78
  79         parse_options(argc, argv, &config, &opt);
  80
  81         ret = grab_frames( &config, &opt );
  82
  83         return ret;
  84 }
  85
  86 static void show_help(void)
  87 {
  88         #define HELP printf
  89         HELP( "Syntax: frameshot [options] infile\n"
  90         "\n"
  91         "Infile is a raw bitstream of one of the following codecs:\n"
  92         "  YUV4MPEG(*.y4m), Dirac(*.drc)\n"
  93         "\n"
  94         "Options:\n"
  95         "\n"
  96         "  -h, --help                  Displays this message.\n"
  97         );
  98         HELP( "  -f, --frames <int,int,...>  Frames numbers to grab.\n" );
  99         HELP( "  -o, --outdir <string>       Output directory for images.\n" );
 100         HELP( "  -z, --compression <integer> Ammount of compression to use.\n" );
 101         HELP( "  -1, --fast                  Use fastest compression.\n" );
 102         HELP( "  -9, --best                  Use best (slowest) compression.\n" );
 103         HELP( "\n" );
 104 }
 105
 106 static int parse_options( int argc, char **argv, config_t *config, cli_opt_t *opt )
 107 {
 108         char *psz_filename = NULL;
 109         int i_zlevel = Z_DEFAULT_COMPRESSION;
 110         char *psz, *psz_token;
 111         int b_y4m = 0;
 112         int b_dirac = 0;
 113         struct stat sb;
 114
 115         memset( opt, 0, sizeof(*opt) );
 116
 117         /* Default input driver */
 118         p_open_infile = open_file_y4m;
 119         p_read_frame = read_frame_y4m;
 120         p_close_infile = close_file_y4m;
 121
 122         /* Default output driver */
 123         p_open_outfile = open_file_png;
 124         p_write_image = write_image_png;
 125         p_close_outfile = close_file_png;
 126
 127         for( ;; )
 128         {
 129                 int long_options_index = -1;
 130                 static struct option long_options[] =
 131                 {
 132                         { "fast", no_argument, NULL, '1' },
 133                         { "best", no_argument, NULL, '9' },
 134                         { "frames", required_argument, NULL, 'f' },
 135                         { "help", no_argument, NULL, 'h' },
 136                         { "outdir", required_argument, NULL, 'o' },
 137                         { "compression", required_argument, NULL, 'z' },
 138                         {0, 0, 0, 0}
 139                 };
 140
 141                 int c = getopt_long( argc, argv, "19f:ho:z:",
 142                                                          long_options, &long_options_index);
 143
 144                 if( c == -1 )
 145                 {
 146                         break;
 147                 }
 148
 149                 switch( c )
 150                 {
 151                         case '1':
 152                                 i_zlevel = Z_BEST_SPEED;
 153                                 break;
 154                         case '9':
 155                                 i_zlevel = Z_BEST_COMPRESSION;
 156                                 break;
 157                         case 'f':
 158                                 for( config->i_frames = 0; config->i_frames < MAX_FRAMES; config->i_frames++, optarg = NULL )
 159                                 {
 160                                         psz_token = strtok(optarg, ",");
 161                                         if( psz_token == NULL )
 162                                                 break;
 163                                         config->frames[config->i_frames] = atoi(psz_token);
 164                                 }
 165                                 qsort(config->frames, config->i_frames, sizeof(*config->frames), intcmp);
 166                                 break;
 167                         case 'o':
 168                                 opt->psz_outdir = strdup(optarg);
 169                                 if( stat( opt->psz_outdir, &sb ) < 0 )
 170                                 {
 171                                         if( mkdir( opt->psz_outdir, S_IRWXU ) < 0 )
 172                                         {
 173                                                 fprintf( stderr, "ERROR: Unable to create output directory.\n" );
 174                                                 perror("mkdir");
 175                                                 return -1;
 176                                         }
 177                                 } else {
 178                                         if( !S_ISDIR(sb.st_mode) )
 179                                         {
 180                                                 fprintf(stderr, "ERROR: Outdir exists, and is not a directory." );
 181                                                 return -1;
 182                                         }
 183                                 }
 184                                 break;
 185                         case 'z':
 186                                 if( optarg == NULL || optarg[0] < '0' || optarg[0] > '9' )
 187                                 {
 188                                         opt->i_zlevel = Z_DEFAULT_COMPRESSION;
 189                                         break;
 190                                 }
 191                                 opt->i_zlevel = atoi(optarg);
 192                                 break;
 193                         case 'h':
 194                         default:
 195                                 show_help();
 196                                 exit(0);
 197                 }
 198         }
 199
 200         /* Get the input file name */
 201         if( optind > argc - 1 )
 202         {
 203                 fprintf( stderr, "ERROR: No input file.\n" );
 204                 show_help();
 205                 return -1;
 206         }
 207         psz_filename = argv[optind++];
 208
 209         psz = strrchr( psz_filename, '.' );
 210         if( !strncasecmp( psz, ".y4m", 4 ) )
 211                 b_y4m = 1;
 212         else if( !strncasecmp( psz, ".drc", 4 ) )
 213                 b_dirac = 1;
 214
 215         if( !opt->psz_outdir )
 216                 opt->psz_outdir = getcwd(NULL, 0);
 217
 218         if( b_y4m )
 219         {
 220                 p_open_infile = open_file_y4m;
 221                 p_read_frame = read_frame_y4m;
 222                 p_close_infile = close_file_y4m;
 223         }
 224
 225         if( b_dirac )
 226         {
 227                 p_open_infile = open_file_dirac;
 228                 p_read_frame = read_frame_dirac;
 229                 p_close_infile = close_file_dirac;
 230         }
 231
 232         if( p_open_infile( psz_filename, &opt->hin, config ) )
 233         {
 234                 fprintf( stderr, "ERROR: could not open input file '%s'\n", psz_filename );
 235                 return -1;
 236         }
 237
 238         return 0;
 239 }
 240
 241 static int grab_frames( config_t *config, cli_opt_t *opt )
 242 {
 243         hnd_t hout;
 244         picture_t pic;
 245         int i;
 246         char tmp[PATH_MAX];
 247
 248         pic.img.plane[0] = calloc(1, 3 * config->i_width * config->i_height / 2);
 249         pic.img.plane[1] = pic.img.plane[0] + config->i_width * config->i_height;
 250         pic.img.plane[2] = pic.img.plane[1] + config->i_width * config->i_height / 4;
 251         pic.img.plane[3] = NULL;
 252
 253         pic.img.i_stride[0] = config->i_width;
 254         pic.img.i_stride[1] = pic.img.i_stride[2] = config->i_width / 2;
 255         pic.img.i_stride[3] = 0;
 256
 257         for(i = 0; i < config->i_frames; i++ )
 258         {
 259                 p_read_frame( opt->hin, &pic, config->frames[i] );
 260
 261                 snprintf(tmp, PATH_MAX, "%s/%05d.png", opt->psz_outdir, config->frames[i]);
 262                 p_open_outfile( tmp, &hout, opt->i_zlevel );
 263                 p_write_image( hout, &pic, config );
 264                 p_close_outfile( hout );
 265         }
 266
 267         p_close_infile( opt->hin );
 268
 269         if( opt->psz_outdir )
 270                 free(opt->psz_outdir);
 271
 272         return 0;
 273 }