#!/usr/bin/env python

from PIL import Image
import numpy as np
import argparse
import sys

p = argparse.ArgumentParser()
p.add_argument('picture', type=str, help='Input picture')
p.add_argument('vtkfile', type=str, help='Output VTK file')
args = p.parse_args()

# Load image data
img = Image.open(args.picture).convert('LA')

# Flip vertically (figure origin is top left instead of bottom left)
img = img.transpose(Image.FLIP_TOP_BOTTOM)

# Determine shape of data
picture_shape = img.size

# Get scalar greyscale data
scalar_data = img.getdata(band=0)

# Generate header
vtk_header="""# vtk DataFile Version 2.0
File generated by picture_to_vtk.py script
ASCII
DATASET STRUCTURED_POINTS
DIMENSIONS {} {} 1
SPACING 1.0 1.0 1.0
ORIGIN 0 0 0
POINT_DATA {}
SCALARS bc_data float 1
LOOKUP_TABLE default""".format(picture_shape[0],
                               picture_shape[1],
                               picture_shape[0] * picture_shape[1])

# Save output
np.savetxt(args.vtkfile, scalar_data, header=vtk_header, comments='')

