ktracker.cpp

从Beevers的个人网站上下载的

/*
  ktracker (c) 2006 Kris Beevers

  This file is part of ktracker.

  ktracker is free software; you can redistribute it and/or modify it
  under the terms of the GNU General Public License as published by
  the Free Software Foundation; either version 2 of the License, or
  (at your option) any later version.

  ktracker is distributed in the hope that it will be useful, but
  WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with ktracker; if not, write to the Free Software Foundation,
  Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA

  $Id: ktracker.cpp,v 1.1.1.1 2006/10/10 20:41:29 beevek Exp $
*/

#include <inttypes.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <SDL/SDL.h>
#include <fstream>
#include <signal.h>
#include <unistd.h>
#include "options.hpp"
#include "network.hpp"
#include "video_input.hpp"
#include "video_output.hpp"
#include "tracker.hpp"

bool gui = false, write_video = false;
video_input *input = 0;
video_output vidout;
SDL_Surface *screen = 0;
FILE *logfile = 0;
tracker::robot_t pose, origin;

// statistics
uint32_t zero, start, now, time_decode = 0, time_tracker = 0, frames = 0;

void cleanup()
{
  // print some statistics
  float total_secs = float(SDL_GetTicks() - zero) / 1000.0;
  fprintf(stderr, "Total frames: %d\n", frames);
  fprintf(stderr, "Total time: %f sec\n", total_secs);
  fprintf(stderr, "FPS: %f\n", float(frames) / total_secs);
  fprintf(stderr, "Decoding time: %f sec (%f/frame)\n",
	  float(time_decode) / 1000.0,
	  float(time_decode) / 1000.0 / float(frames));
  fprintf(stderr, "Tracking time: %f sec (%f/frame)\n",
	  float(time_tracker) / 1000.0,
	  float(time_tracker) / 1000.0 / float(frames));
  if(gui)
    SDL_Quit();
  if(input)
    input->close();
  if(write_video)
    vidout.close();
  tracker::cleanup();
  if(logfile)
    fclose(logfile);
}

void sigint(int s) { exit(0); }

void init_video_input()
{
  std::string in = options::quickget<std::string>("input");
  input = autoopen_video(in.c_str());
  if(!input) {
    fprintf(stderr, "Opening %s: ", in.c_str());
    perror("");
    exit(1);
  }
}

// initialize SDL
void init_gui()
{
  if(SDL_Init(SDL_INIT_VIDEO) < 0) {
    fprintf(stderr, "Initializing SDL: %s\n", SDL_GetError());
    exit(1);
  }

  screen = SDL_SetVideoMode(input->width, input->height,
			    input->bpp, SDL_SWSURFACE);
  if(!screen) {
    fprintf(stderr, "Can't set %dx%dx%d video: %s\n", input->width, input->height,
	    input->bpp, SDL_GetError());
    exit(1);
  }
}

bool handle_packet(int sock)
{
  uint8_t cmd;
  if(read(sock, &cmd, 1) < 0 || cmd == 0x00)
    return false;

  static tracker::robot_t xformed;

  switch(cmd) {
  case 0x00: // get current pose
    xformed.x = pose.x - origin.x;
    xformed.y = pose.y - origin.y;
    xformed.t = pose.t - origin.t;
    // ALERT: this assumes the client is running on the same kind of
    // platform as this server!  (i.e., we're not translating to any
    // kind of network-safe number representation)
    if(write(sock, &xformed, sizeof(tracker::robot_t)) != sizeof(tracker::robot_t))
      return false;
    break;
  case 0x01: // set current pose as origin
    origin = pose;
    break;
  }
  return true;
}

void screenshot(const char *prefix, const uint8_t *img)
{
  static uint32_t count = 0;
  static char fn[128];
  sprintf(fn, "%s%08d.ppm", prefix, count++);
  FILE *out = fopen(fn, "wb");
  if(!out) return;
  // don't use write_ppm because we need to swap blue and red
  fprintf(out, "P6\n%d\n%d\n255\n", input->width, input->height);
  for(uint32_t i = 0; i < input->imgsize; i += 3) { // need to convert BGR->RGB
    const uint8_t *p = img+i;
    fwrite(p+2, 1, 1, out);   // red
    fwrite(p+1, 1, 1, out);   // green
    fwrite(p, 1, 1, out);     // blue
  }
  fclose(out);
}

int parse_command_line(int argc, char **argv);

int main(int argc, char **argv)
{
  parse_command_line(argc, argv);

  atexit(cleanup);
  signal(SIGINT, sigint);

  if(init_network(handle_packet) < 0)
    exit(1);

  init_video_input();
  gui = options::quickget<bool>("gui");
  if(gui)
    init_gui();
  if(!tracker::init(*input)) {
    fprintf(stderr, "Can't initialize tracking\n");
    exit(1);
  }

  // use the calibrated reference frame unless we're told otherwise by
  // a client
  origin.x = origin.y = origin.t = 0;

  const char *ss_prefix = options::quickget<std::string>("ss-prefix").c_str();

  if(options::quickget<std::string>("output").size() > 1) {
    write_video = true;
    vidout.set_context(input->width, input->height);
    if(!vidout.open(options::quickget<std::string>("output").c_str())) {
      fprintf(stderr, "Can't open output video stream\n");
      exit(1);
    }
  }

  // write the pose from each frame to an "RPI ASCII format" log file
  if(options::quickget<std::string>("logfile").size() > 1) {
    logfile = fopen(options::quickget<std::string>("logfile").c_str(), "a");
    if(!logfile) {
      fprintf(stderr, "Can't open logfile %s\n", options::quickget<std::string>("logfile").c_str());
      exit(1);
    }
  }

  // move to the background if necessary
  if(options::quickget<bool>("daemon") && daemon(1, 0) < 0) {
    perror("daemon()");
    exit(1);
  }

  // main loop

  SDL_Event sdlevent;
  uint8_t *img;

  zero = start = SDL_GetTicks();
  while((img = input->get_next_frame()) != 0) { // decode a frame
    time_decode += SDL_GetTicks() - start;
    ++frames;

    // motion tracking
    start = SDL_GetTicks();
    tracker::track_robot(img, pose);
    now = SDL_GetTicks();
    time_tracker = now - start;

    if(logfile) {
      // write a timestamp and the pose to the log file
      struct timeval tv;
      gettimeofday(&tv, 0);
      fprintf(logfile, "T %d %g\n", now, double(tv.tv_sec)+double(tv.tv_usec)/1000000.0);
      fprintf(logfile, "P %g %g %g\n", pose.x, pose.y, pose.t * 180.0 / M_PI);
    }

    if(gui) {
      // draw to the screen
      memcpy(screen->pixels, img, input->imgsize);
      SDL_UpdateRect(screen, 0, 0, 0, 0);

      // check for input
      if(SDL_PollEvent(&sdlevent) == 1) {
	if(sdlevent.type == SDL_QUIT ||
	   (sdlevent.type == SDL_KEYDOWN && sdlevent.key.keysym.sym == SDLK_q))
	  break;
	else if(sdlevent.type == SDL_KEYDOWN && sdlevent.key.keysym.sym == SDLK_s)
	  screenshot(ss_prefix, img);
      }
    }

    start = SDL_GetTicks();

    if(write_video) // save the video with PTS=start
      vidout.encode_frame(img, start);

    start = SDL_GetTicks();
  }

  return 0;
}

int parse_command_line(int argc, char **argv)
{
  // set up commandline/configuration file options
  options::add<bool>("help", 0, "Print usage information", 0, false, options::nodump);
  options::set_cf_options("config", "c");
  options::add<std::string>("save-config", 0, "Save configuration file", 0, "", options::nodump);

  options::add<std::string>("input", "i", "Input source (def. /dev/video1394/0)", 0, "/dev/video1394/0");
  options::add<std::string>("output", "o", "Output MPEG movie filename (def. none)", 0, "");
  options::add<bool>("daemon", "d", "Run in background as daemon (def. false)", 0, false);
  options::add<std::string>("logfile", "l", "Log file of robot poses (def. none)", 0, "");
  options::add<bool>("gui", "g", "Show video in a window (def. false)", 0, false);
  options::add<bool>("no-overlays", 0, "Don't draw vision overlays (def. false)", 0, false);

  options::add<std::string>("ss-prefix", 0, "Prefix for screenshot filenames", 0, "frame");

  options::add<std::string>("calibration-file", "a", "Camera calibration file",
			    "Calibration", "");
  options::add<std::string>("front-template", 0, "Template image for front marker",
			    "Calibration", "");
  options::add<std::string>("back-template", 0, "Template image for back marker",
			    "Calibration", "");

  // parse commandline (and read config file if it is specified)
  int inpidx = options::parse_cmdline(argc, argv);

  if(inpidx < 0) // some kind of error
    exit(1);    // relevant info printed by getopt

  // print usage information
  if(inpidx > argc || options::quickget<bool>("help") == true) {
    std::cerr << "Usage: " << argv[0] << " [options]" << std::endl;
    options::print_options(std::cout);
    exit(1);
  }

  // save a config file based on these options?
  std::string cfname = options::quickget<std::string>("save-config");
  if(cfname.length() > 0) {
    std::ofstream conf(cfname.c_str());
    if(!conf)
      std::cerr << "Can't write configuration file " << cfname << std::endl;
    else
      options::dump(conf);
  }

  return 0;
}