#!/usr/bin/env python
# -*- coding: utf-8 -*-

# Copyright 2016 Massachusetts Institute of Technology

"""Extract images from a rosbag.
"""

import os
import argparse

import cv2
import numpy as np

import rosbag
from sensor_msgs.msg import Image
from cv_bridge import CvBridge

def main():
    """Extract a folder of images from a rosbag.
    """
    parser = argparse.ArgumentParser(description="Extract images from a ROS bag.")
    parser.add_argument("bag_file", help="Input ROS bag.")
    parser.add_argument("output_dir", help="Output directory.")
    parser.add_argument("image_topic", help="Image topic.")
    parser.add_argument("image_prefix", help="Image prefix.")
    parser.add_argument("normalize", help="True->perform normalization False->no normalization.")

    args = parser.parse_args()

    print "Extract images from %s on topic %s into %s [prefix %s]; normalization %s" % (args.bag_file, args.image_topic, args.output_dir, args.image_prefix, bool(args.normalize))

    bag = rosbag.Bag(args.bag_file, "r")
    bridge = CvBridge()
    count = 0
    for topic, msg, t in bag.read_messages(topics=[args.image_topic]):
        cv_img = bridge.imgmsg_to_cv2(msg, desired_encoding="mono8")
        if args.normalize:
        		cv_img_norm = np.zeros(shape=(cv_img.shape[0], cv_img.shape[1]), dtype=cv_img.dtype)
        		cv2.normalize (cv_img, cv_img_norm, 0, 255, cv2.NORM_MINMAX)
				
#        cv2.imwrite(os.path.join(args.output_dir, "frame%06i.png" % count), cv_img)
	if bool(args.normalize):
		print "Normalization bool is true"
		cv2.imwrite(os.path.join(args.output_dir, args.image_prefix + ("_%04i.png" % count)), cv_img_norm)
	else:
		cv2.imwrite(os.path.join(args.output_dir, args.image_prefix + ("_%04i.png" % count)), cv_img)

	print os.path.join(args.output_dir, args.image_prefix + ("_%04i.png" % count))
        print "Wrote image %i" % count

        count += 1

    bag.close()

    return

if __name__ == '__main__':
    main()