#!/usr/bin/perl

#
# Version 0.02
#
# ChangeLog:
#
#   Mon Apr 14
#     Added <image> tag, which allows one to generate frames using a still.
#     Fixed read behaviour when reading until end of stream, so it no longer reports an error.
#

#
# Run perldoc on me if you want to know how I work. :)
#

$SIG{PIPE} = sub {};

package Script;

use PDL;
use IO::File;
use YUV4MPEG ':all';
use UNIVERSAL 'isa';

my %rates = (0 => FPS_UNKNOWN,
             1 => FPS_NTSC_FILM,
             2 => FPS_FILM,
             3 => FPS_PAL,
             4 => FPS_NTSC,
             5 => FPS_30,
             6 => FPS_PAL_FIELD,
             7 => FPS_NTSC_FIELD,
             8 => FPS_60);

my %ilace = ("none" => Y4M_ILACE_NONE,
             "bf"   => Y4M_ILACE_BOTTOM_FIRST,
             "tf"   => Y4M_ILACE_TOP_FIRST);

my %aspect = (0 => DAR_UNKNOWN,
              1 => new YUV4MPEG::Ratio(1, 1),
              2 => DAR_4_3,
              3 => DAR_16_9,
              4 => DAR_221_100);

sub new
{
  my ($class, $params, $children) = @_;
  my $self = bless {}, $class;
  my $si;

  $self->{"children"} = $children;

  $si = $self->get_header();

  if (defined $params->{"fps"}) 
    {
      my $rate = $rates{$params->{"fps"}};
      
      die "Invalid framerate specified" if (! defined $rate);
      
      $si->set_framerate($rate);
    } 

  if (defined $params->{"width"}) 
    {
      $si->set_width($params->{"width"});
    }

  if (defined $params->{"height"}) 
    {
      $si->set_height($params->{"height"});
    } 

  if (defined $params->{"interlace"}) 
    {
      my $mode = $ilace{$params->{"interlace"}};
      
      die "Invalid interlacing mode specified" if (! defined $mode);
      
      $si->set_interlace($mode);
    } 

  if (defined $params->{"aspect"}) 
    {
      my $ratio = $aspect{$params->{"aspect"}};
      
      die "Invalid aspect ratio specified" if (! defined $ratio);
      
      $ratio = YUV4MPEG::Ratio::guess_sar($ratio, 
                                          $si->get_width(),
                                          $si->get_height());
      
      $si->set_sampleaspect($ratio);
    } 

  $self->{"target"} = $params->{"target"};
  $self->{"si"} = $si;
  $self->{"curr"} = shift @{ $children };
  $self->{"children"} = $children;  

  if ($self->{"target"} eq "-")
    {
      $self->{"handle"} = \*STDOUT;
    }
  else
    {
      $self->{"handle"} = new IO::File ">" . $self->{"target"};
    }

  if (! defined $self->{"handle"})
    {
      die "Unable to open $self->{target}: $!";
    }

  return $self;
}

sub get_header
{
  my $self = shift;

  foreach (@{ $self->{"children"} })
    {
      if (! isa($_, "Image"))
        {
          return $_->get_header();
        }
    }

  return new YUV4MPEG::StreamInfo;
}

sub write_header
{
  my $self = shift;

  if ($self->{"si"}->write($self->{"handle"}) != Y4M_OK)
    {
      die "Unable to write stream header to '$self->{target}'";
    }
}

sub execute
{
  my $self = shift;
  my $frame_pair;
  my $frame_num = 0;

  $self->write_header();

  while (defined ($frame_pair = $self->read()))
    {
      my ($fi, $frame) = @{ $frame_pair };
      my $out_frame = new YUV4MPEG::Frame($self->{"si"});

      #
      # This ensures that we use the correct interlacing mode
      #

      $out_frame->set_pdl_data($frame->get_pdl_data);

      print STDERR "   INFO: [yuvpipe] Frame $frame_num...\n";

      if ($out_frame->write($fi, $self->{"handle"}) != Y4M_OK)
        {
          die "Unable to write frame to '$self->{target}'";
        }

      $frame_num++;
    }
 }

sub read
{
  my $self = shift;
  my $frame_pair;

  if (! defined ($frame_pair = $self->{"curr"}->read()))
    {
      $self->{"curr"} = shift @{ $self->{"children"} };

      return undef if (! defined $self->{"curr"});

      $frame_pair = $self->{"curr"}->read();
    }

  return $frame_pair;
}

package Filter;

use IO::Pipe;
use YUV4MPEG ':all';
use UNIVERSAL 'isa';

sub new
{
  my ($class, $params, $inputs) = @_;
  my $self = bless {}, $class;

  if (! defined $params->{command})
    {
      die "No command specified in filter";
    }

  $self->{"command"} = $params->{"command"};
  $self->{"inputs"} = $inputs;

  $self->spawn();

  $si = $self->get_header();

  $self->{"si"} = new YUV4MPEG::StreamInfo($si);

  if ($self->{"si"}->write($self->{"writer"}) != Y4M_OK)
    {
      die "Error writing stream header to '$self->{command}'";
    }

  if ($self->{"si"}->read($self->{"reader"}) != Y4M_OK)
    {
      die "Error reading stream header from '$self->{command}'";
    }

  return $self;
}

sub get_header
{
  my $self = shift;

  foreach (@{ $self->{"inputs"} })
    {
      if (! isa($_, "Image"))
        {
          return $_->get_header();
        }
    }
}

sub spawn
{
  my $self = shift;
  my $p1 = new IO::Pipe;
  my $p2 = new IO::Pipe;
  my $pid;

  if ($pid = fork())
    {
      $p1->reader();
      $p2->writer();

      $p2->autoflush(1);

      $self->{"pid"} = $pid;
      $self->{"reader"} = $p1;
      $self->{"writer"} = $p2;
    }
  elsif ($pid == 0)
    {
      $p1->writer();
      $p2->reader();

      $p1->autoflush(1);

      local *WRITE_PARENT = $p1;
      local *READ_PARENT = $p2;

      open(STDIN, "<&READ_PARENT");
      open(STDOUT, ">&WRITE_PARENT");

      exec $self->{"command"};

      die "Error executing '$self->{command}': $!\n";
    }
  else
    {
      die "Couldn't fork filter: $!\n";
    }
}

sub kill_spawn
{
  my $self = shift;

  kill 9, $self->{"pid"};
  waitpid $self->{"pid"}, 0;
}

sub read
{
  my $self = shift;
  my $input;
  my $cnt = 0;

  foreach $input (@{ $self->{"inputs"} })
    {
      my $frame_pair = $input->read();

      if (! defined $frame_pair)
        {
          $self->kill_spawn();
          return undef;
        }

      my ($fi, $frame) = @{ $frame_pair };

      if ($frame->write($fi, $self->{"writer"}) != Y4M_OK)
        {
          $self->kill_spawn();
          return undef;
        }

      $cnt++;
    }

  my $fi = new YUV4MPEG::FrameInfo;
  my $frame = new YUV4MPEG::Frame($self->{"si"});

  if ($frame->read($fi, $self->{"reader"}) != Y4M_OK)
    {
      $self->kill_spawn();

      die "Error reading frame from '$self->{command}'";
    }

  return [$fi, $frame];
}

package Input;

use YUV4MPEG;
use IO::File;

sub new
{
  my ($class, $params) = @_;
  my $self = bless {}, $class;
  my $si = new YUV4MPEG::StreamInfo;
  
  $self->{"source"} = $params->{"src"};
  $self->{"start"} = $params->{"start"};
  $self->{"end"} = $params->{"end"};  
  $self->{"pos"} = 0;

  if (! defined $self->{"start"})
    { $self->{"start"} = 0; }

  if (! defined $self->{"end"})
    { $self->{"end"} = -1; }

  if (defined $params->{"exclude"})
    {
      $self->{"exclude"} = parse_range($params->{"exclude"});
    }
  else
    {
      $self->{"exclude"} = [];
    }

  if ($self->{"source"} eq "-")
    {
      $self->{"handle"} = \*STDIN;
    }
  else
    {
      $self->{"handle"} = new IO::File $self->{"source"};      
    }
  
  if (! defined $self->{"handle"})
    {
      die "Unable to open $self->{source}: $!";
    }
  
  $self->{"si"} = new YUV4MPEG::StreamInfo;
  
  if ($self->{"si"}->read($self->{"handle"}) != Y4M_OK) 
    {
      die "Error reading stream header from $self->{source}";
    }

  return $self;
}

sub exclude_frame
{
  my $self = shift;
  my @ranges = @{ $self->{"exclude"} };
  my $num = $self->{"pos"};
  
  foreach (@ranges) {
    if (($_->{"low"} <= $num) &&
        ($_->{"high"} >= $num)) {
      return 1;
    }
  }
  
  return 0;
}

sub parse_range
{
  my $spec = shift;
  my @ranges = split(/,/, $spec);
  my @result;
  
  foreach (@ranges)
    {
      my @parts = split(/-/, $_);
      
      die "Invalid range specified: $_" if (($#parts > 1) || ($#parts < 0));
      
      $parts[1] = $parts[0] if (! defined $parts[1]);
      
      push @result, { "low" => $parts[0],
                      "high" => $parts[1] };
    }
 
  return \@result;
}

sub get_header
{
  my $self = shift;

  return $self->{"si"};
}

sub read_frame
{
  my $self = shift;
  my $frame = shift;
  my $fi = new YUV4MPEG::FrameInfo;

  my $ret = $frame->read($fi, $self->{"handle"});

  if (($ret != Y4M_OK) && ($self->{"pos"} < $self->{"end"}))
    {
      die "Couldn't read frame $self->{pos} from $self->{source}";
    }
  elsif ($ret != Y4M_OK)
    {
      return undef;
    }
  
  $self->{"pos"}++;

  return $fi;
}

sub read
{
  my $self = shift;

  return undef if (($self->{"pos"} > $self->{"end"}) &&
                   ($self->{"end"} != -1));

  my $fi;
  my $frame = new YUV4MPEG::Frame($self->{"si"});

  while ($self->{"pos"} < $self->{"start"})
    {
      $fi = $self->read_frame($frame);
    }

  while ($self->exclude_frame())
    {
      $fi = $self->read_frame($frame);
    }

  $fi = $self->read_frame($frame);

  return undef if (! defined $fi);

  return [$fi, $frame];
}

package Image;

use YUV4MPEG ':all';
use IO::File;
use PDL;
use PDL::IO::Pic;

sub new
{
  my ($class, $params) = @_;
  my $self = bless {}, $class;

  die "No count specified in image tag" if (! defined $params->{"count"});

  $self->{"count"} = $params->{"count"};

  my $pdl = PDL->rpic($params->{"src"});

  die "Unable to open image: $params->{src}" if (! defined $pdl);

  my @dims = $pdl->dims();
  my $si = new YUV4MPEG::StreamInfo;

  $si->set_width($dims[1]);
  $si->set_height($dims[2]);
  $si->set_interlace(Y4M_ILACE_NONE);

  my $frame = new YUV4MPEG::Frame($si);
  my $fi = new YUV4MPEG::FrameInfo;

  $frame->set_pdl_data($pdl);
  $frame->to_yuv();

  $self->{"si"} = $si;
  $self->{"read_val"} = [$fi, $frame];
  $self->{"curr"} = 0;
  $self->{"count"} = $params->{"count"};

  return $self;
}

sub read
{
  my $self = shift;

  return undef if ($self->{"curr"} > $self->{"count"});

  $self->{"curr"}++;

  return $self->{"read_val"};
}

package Segment;

use UNIVERSAL 'isa';

sub new
{
  my ($class, $params, $children) = @_;
  my $self = bless {}, $class;

  $self->{"curr"} = shift @{ $children };
  $self->{"children"} = $children;

  return $self;
}

sub get_header
{
  my $self = shift;

  foreach (@{ $self->{"children"} })
    {
      if (! isa($_, "Image"))
        {
          return $_->get_header();
        }
    }
}

sub read
{
  my $self = shift;
  my $frame_pair;

  if (! defined ($frame_pair = $self->{"curr"}->read()))
    {
      $self->{"curr"} = shift @{ $self->{"children"} };

      return undef if (! defined $self->{"curr"});

      $frame_pair = $self->{"curr"}->read();
    }

  return $frame_pair;
}

package main;

use XML::Parser;

my %id_hash;
my @obj_list;

sub parse_script
{
  my $parser = new XML::Parser(Style => 'Tree');
  my $parse_tree = $parser->parsefile(shift);

  my $elem = $parse_tree->[0];
  my $contents = $parse_tree->[1];
  my $params = shift @$contents;  

  my $root = parse_element($elem, $params, $contents);

  return $root;
}

sub parse_element
{
  my ($elem, $params, $contents) = @_;
  my @children;
  my $i;
  my $obj;

  for ($i = 0; $i < $#{ $contents }; $i += 2)
  {
    next if ($contents->[$i] eq "0");

    my $child = $contents->[$i];
    my $child_contents = $contents->[$i + 1];
    my $child_params = shift @$child_contents;

    push @children, parse_element($child, $child_params, $child_contents);
  }

  if ($elem eq "script") {
    $obj = new Script($params, \@children);
  } elsif ($elem eq "filter") {
    $obj = new Filter($params, \@children);
  } elsif ($elem eq "input") {
    die "Input should not contain elements" if ($#children >= 0);

    $obj = new Input($params);
  } elsif ($elem eq "segment") {
    $obj = new Segment($params, \@children);
  } elsif ($elem eq "image") {
    $obj = new Image($params);
  } else {
    die "Unknown tag $elem";
  }

  if (defined $params->{"id"})
    {
      $id_hash{$params->{"id"}} = $obj;
    }

  push @obj_list, $obj;

  return $obj;
}

die "No script provided" if (! defined $ARGV[0]);

my $script = parse_script($ARGV[0]);

$script->execute();

exit(0);

__END__

=head1 NAME

yuvpipe - An editing tool for manipulating YUV4MPEG data.

=head1 SYNOPSIS

yuvpipe script.xml

=head1 DESCRIPTION

yuvpipe is a tool inspired by the basic ideas of lavpipe, a tool for
editing/filtering YUV4MPEG data.  lavpipe, part of mjpegtools, allows
one to specify how a series of inputs and filters should be combined
to form a single YUV4MPEG output stream.  While very useful, yuvpipe
does have some limitations, and it's "pipe script" format is, at least
in this authors opinion :), rather difficult to work with.

In yuvpipe, the concept of a "pipe script" is replaced with a more
powerful XML representation.  This XML is used to define how a series
of input YUV streams and filters should be combined to form a single
YUV stream as output.

=head1 OPTIONS

Just specify the XML file on the command line. :)

=head1 XML FORMAT

The XML format for yuvpipe is fairly simple.  The following is a list
of the allowed tags and their various parameters:

=over 8

=item <script>

The root node.  Defines the output target, and allows one to specify
various parameters for the output stream.  The script tag takes the
following parameters:

=over 8

=item target

The output target for the stream.  Currently, this must be a valid
file name, or "-", which indicates standard output.

=item fps

The frame rate for the output stream, specified using the regular
integer constants (see the -F parameter of mpeg2enc).

=item width

The width of the output stream.

=item height

The height of the output stream.

=item interlace

The interlacing mode of the stream.  Valid values are none, 
tf, and bf, which represent none/progressive, top first, and bottom
first interlacing, respectively.

=item aspect

The aspect ratio for the stream.  Uses the standard integer values found
in mjpeg2enc (see the -a option).

=back

Note, the default values are taken from the first child element in the
script tag.

The contents of this tag can be anything but another script tag.  The
output is a concatenation of the result of all the child elements of 
the tag.

=item <segment>

The segment tag, which has no parameters, simply concatenates all it's
child inputs together, much like the <script> tag.

=item <filter>

This tag reads a frame from each of its child elements and pipes them
through a user-defined command.  Its output is the result of this
command.  Thus, the filter tag takes a series of inputs and combines
them into to one output.  Valid parameters for the filter tag are:

=over 8

=item command

The command to execute and pipe frame data through.

=back

=item <input>

The input tag reads frame data from some source and returns it.
It's parameters are:

=over 8

=item src

The source for the stream.  This could be a file, or any other valid
Perl file specifier.  eg "lav2yuv scene1.avi|".

=item start

The starting frame number to output.  Default is 0.

=item end

The ending frame number to output.  A value of -1 indicates "read until
end of stream".  Default is -1.

=item exclude

Causes the input to exclude the specified set of frames when reading.
This parameter should be specified as a comma-delimited set of ranges.
eg: "1-5,10-15,20,25-100".

=back

=item <image>

Loads a still image and outputs the image as a frame the number of times
specified.  Note, an image tag alone isn't sufficient to define the
parameters for the YUV stream.  So, if you intend to just create YUV
still stream, you must specify all the stream parameters manually
(including width and height!).  Parameters:

=over 8

=item src

The image file to load.  Uses the PDL to load the image.

=item count

The number of frames this image tag should generated.

=back

=back

=head1 EXAMPLES

The following is a replacement for the transition example in the
lavpipe manpage, with the additional twist of denoising the output:

<script target="-">
  <filter command="yuvdenoise">
    <segment>
      <input src="lav2yuv scene1.avi|" end="50"/>
  
      <filter command="transist.flt -o 0 -O 255 -d 50">
        <input src="lav2yuv scene1.avi|" start="50" end="75"/>
        <input src="lav2yuv scene2.avi|" end="25"/>
      </filter>

      <input src="lav2yuv scene2.avi|" start="25" end="75"/>
    </segment>
  </filter>
</script>

=head1 BUGS

Probably many. :)  The script should really be checked more thoroughly
to see if it's malformed.

=head1 NOTES

None yet.  Give me time. :)

=head1 AUTHOR

brettk@frodo.dyn.gno.org

=head1 SEE ALSO

L<lavpipe>, L<yuv2lav>, L<lav2yuv>, L<yuvscaler>, L<mpeg2enc>

=end
