fits2png

#!/usr/bin/python
# -*- coding: utf-8 -*-
"""fits2png module.

Author: Jan Milík <milikjan@fit.cvut.cz>
"""


import sys
import argparse
import datetime

if sys.version_info < (2, 7):
    sys.exit("Python version 2.7 or higher required.")

try:
    import pyfits
except ImportError:
    sys.exit("pyfits library not found.")

try:
    import Image
    import ImageFont
    import ImageDraw
except ImportError:
    sys.exit("PIL (Python Imaging Library) not found.")

try:
    import numpy
except ImportError:
    sys.exit("numpy library not found.")


def get_image_info(image):
    return (
        image.header["NAXIS1"],
        image.header["NAXIS2"],
        image.data.min(),
        image.data.max(),
        image.data.mean(),
        image.data.var(),
    )


FN_IDENTITY = lambda image: image.data
FN_LOG      = lambda image: numpy.log(image.data[image.data != 0])


# A UTC class.
class UTC(datetime.tzinfo):
    """UTC"""
    ZERO = datetime.timedelta(0)
    
    def utcoffset(self, dt):
        return self.ZERO
    
    def tzname(self, dt):
        return "UTC"
    
    def dst(self, dt):
        return self.ZERO


class View(object):
    @property
    def requested_size(self):
        self.update()
        return self._requested_size
    
    def __init__(self, font):
        self.font = font
        self.foreground = 255
        self.background = 0
        
        self.dirty = True
        self._requested_size = None
    
    def invalidate(self):
        self.dirty = True
    
    def update(self):
        if not self.dirty:
            return
        self.dirty = False
        self._update()
    
    def _update(self):
        pass
    
    def draw(self, draw, rect):
        self.update()
        self._draw(draw, rect)
    
    def _draw(self, draw, rect):
        pass


class ImageAxis(object):
    def __init__(self, ref, value, delta, length, type = None):
        self.reference = ref - 1 # coordinate of the reference pixel (FITS counts from 1)
        self.value  = value      # logical value of the reference pixel
        self.delta  = delta      # difference of the logical value between two subsequent pixels along the axis
        self.length = length
        self.type   = type
        
        self.scale  = 1.0
    
    def __repr__(self):
        return "ImageAxis(%d, %f, %f)" % (self.reference + 1, self.value, self.delta)
    
    def __str__(self):
        offset_str = ""
        if self.value <> 0:
            offset_str = "%f + " % (self.value, )
        x_str = "x"
        if self.reference <> 0:
            x_str = "(x - %d)" % (self.reference, )
        scale_str = ""
        if self.delta <> 1:
            scale_str = " * %f" % (self.delta, )
        result = offset_str + x_str + scale_str
        if self.type is not None:
            result += " (%s)" % (self.type, )
        return result
    
    def to_logical(self, physical):
        return self.value + ((self.delta / self.scale) * float(physical - self.scale * self.reference))
    
    def to_physical(self, logical):
        return int(self.scale * self.reference + (self.scale / self.delta) * (logical - self.value))
    
    def to_logical_length(self, physical):
        return (self.delta / self.scale) * float(physical)
    
    def to_physical_length(self, logical):
        return (self.scale / self.delta) * float(logical)
    
    def draw(self, draw, rect, vertical = True):
        ticks = 1.0
        direction = 1.0
        if self.to_physical_length(ticks) < 0.0:
            ticks = -1.0
            direction = -1.0
        for i in range(10):
            ticks = direction * (10.0 ** i)
            if self.to_physical_length(ticks) > 20:
                break
        
        x, y, width, height = rect
        draw.rectangle((x, y, x + width, y + height), fill = 0)
        
        physical = 0
        logical = self.to_logical(physical)
        
        if vertical:
            draw.line((x + width, y, x + width, y + height), fill = 255)
            while physical < height:
                draw.line((x, y + physical, x + width, y + physical), fill = 255)
                draw.text((x, y + physical), "%d" % logical, fill = 255)
                logical += ticks
                physical = self.to_physical(logical)
        else:
            draw.line((x, y, x + width, y), fill = 255)
            while physical < width:
                draw.line((x + physical, y, x + physical, y + height), fill = 255)
                draw.text((x + physical, y), "%d" % logical, fill = 255)
                logical += ticks
                physical = self.to_physical(logical)
    
    @classmethod
    def from_image(cls, image, index):
        try:
            length = image.header["NAXIS%d" % (index + 1, )]
        except KeyError:
            raise Exception("Image does not seem to have axis %d." % (index, ))
        
        ref = image.header.get("CRPIX%d" % (index + 1, ), 1)
        value = image.header.get("CRVAL%d" % (index + 1, ), 1)
        delta = image.header.get("CDELT%d" % (index + 1, ), 1)
        atype = image.header.get("CTYPE%d" % (index + 1, ), None)
        
        return ImageAxis(ref, value, delta, length, atype)


class AxisView(View):
    def __init__(self, font, axis, vertical = True):
        super(AxisView, self).__init__(font)
        self.axis = axis
        self.vertical = vertical
        
        #self.margin = 0.3
        
        #self.ticks_size = 0.5
        #self.width = 20
    
    def _update(self):
        self.ticks = 1.0
        direction = 1.0
        if self.axis.to_physical_length(self.ticks) < 0.0:
            self.ticks = -1.0
            direction = -1.0
        
        for i in range(10):
            self.ticks = direction * (10.0 ** i)
            t, max_size = self._get_ticks(self.axis.length)
            max_size = max_size[1] if self.vertical else max_size[0]
            physical = self.axis.to_physical_length(self.ticks)
            if (physical > 20) and (physical > max_size):
                break
            
            self.ticks = direction * (10.0 ** i) * 2
            t, max_size = self._get_ticks(self.axis.length)
            max_size = max_size[1] if self.vertical else max_size[0]
            physical = self.axis.to_physical_length(self.ticks)
            if (physical > 20) and (physical > max_size):
                break
        
        #self.ticks = ticks
        
        labels, tick_label_size = self._get_ticks(self.axis.length)
        
        self.margin = int(tick_label_size[1] * 0.3)
        if self.vertical:
            self.ticks_size = tick_label_size[0] + self.margin
        else:
            self.ticks_size = tick_label_size[1] + self.margin
        
        if (self.axis.type is not None) and len(self.axis.type) > 0:
            label_size = self.font.getsize(str(self.axis.type))
        else:
            label_size = (0, 0)
        
        if self.vertical:
            self._requested_size = (self.ticks_size + label_size[0] + self.margin, self.axis.length, )
        else:
            self._requested_size = (self.axis.length, int(self.ticks_size + label_size[1] + self.margin))
    
    def _get_ticks(self, length):
        physical = 0
        logical = self.axis.to_logical(physical)

        result = []
        max_width = 0
        max_height = 0
        
        if self.axis.type and self.axis.type.upper() == "TIME":
            print "TIME"
            text_fn = lambda x: datetime.datetime.fromtimestamp(x, UTC()).strftime("%H:%M:%S")
        else:
            print "OTHER"
            text_fn = lambda x: "%d" % (x, )
        
        while physical < length:
            #text = "%d" % (logical, )
            text = text_fn(logical)
            size = self.font.getsize(text)
            max_width = max(max_width, size[0])
            max_height = max(max_height, size[1])
            result.append((physical, text, size, ))
            logical += self.ticks
            physical = self.axis.to_physical(logical)
        
        return (result, (max_width, max_height), )
    
    def _draw(self, draw, rect):
        x, y, width, height = rect
        
        # Check if zero is visible
        physical = self.axis.to_physical(0.0)
        if (physical >= 0) and (physical < self.axis.length):
            pass
        
        physical = 0
        logical = self.axis.to_logical(physical)

        if self.vertical:
            tick_end   = x + width
            tick_start = tick_end - self.ticks_size
            
            draw.line((tick_end, y, tick_end, y + height), fill = self.foreground)
            #while physical < height:
            #    draw.line((tick_start, y + physical, tick_end, y + physical), fill = self.foreground)
            #    draw.text((tick_start, y + physical), "%d" % logical, fill = self.foreground)
            #    logical += self.ticks
            #    physical = self.axis.to_physical(logical)

            for physical, logical, lbl_size in self._get_ticks(height)[0]:
                draw.line((tick_start, y + physical, tick_end, y + physical), fill = self.foreground)
                draw.text((tick_start, y + physical), logical, fill = self.foreground)
            
            label = (" %s" % (self.axis.type, )) or ""
            if len(self.axis.type) > 0:
                label_size = self.font.getsize(label)
                #label_x = x + width / 2 - label_size[0] / 2
                label_x = x
                #label_y = y + height - self.ticks_size + self.margin
                label_y = y + height / 2 - label_size[1] / 2
                draw.text((label_x, label_y), label, font = self.font, fill = self.foreground)
            
            #label = self.axis.type or ""
            #if len(self.axis.type) > 0:
            #    label_size = self.font.getsize(label)
            #    label_img = Image.new("L", label_size)
            #    d = ImageDraw.Draw(label_img)
            #    d.text((0, 0), label, font = self.font, fill = self.foreground)
            #    label_img = label_img.rotate(90)
        else:
            tick_end = y
            tick_start = tick_end + self.ticks_size
            
            draw.line((x, tick_end, x + width, tick_end), fill = self.foreground)
            #while physical < width:
            #    draw.line((x + physical, tick_start, x + physical, tick_end), fill = self.foreground)
            #    draw.text((x + physical, tick_start), "%d" % logical, font = self.font, fill = self.foreground)
            #    logical += self.ticks
            #    physical = self.axis.to_physical(logical)
            
            for physical, logical, lbl_size in self._get_ticks(width)[0]:
                draw.line((x + physical, tick_start, x + physical, tick_end), fill = self.foreground)
                draw.text((x + physical, tick_start - lbl_size[1]), logical, font = self.font, fill = self.foreground)
            
            label = self.axis.type or ""
            if len(self.axis.type) > 0:
                label_size = self.font.getsize(label)
                label_x = x + width / 2 - label_size[0] / 2
                label_y = y + height - self.ticks_size + self.margin
                draw.text((label_x, label_y), label, font = self.font, fill = self.foreground)


class Lines(View):
    @property
    def size(self):
        self.update()
        return (self._width, self._height, self._margin, self._positions)
    
    @property
    def lines(self):
        return self._lines
    
    @lines.setter
    def lines(self, value):
        self._lines = value
        self.invalidate()
    
    def __init__(self, font, lines = None):
        super(Lines, self).__init__(font)
        
        if lines is None:
            lines = []
        self._lines = lines
    
    def _update(self):
        width = 0
        height = 0
        heights = []
        size = (0, 0)
        
        for line in self.lines:
            size = self.font.getsize(line)
            width = max(width, size[0])
            height += size[1]
            heights.append(size[1])
        
        margin = size[1] / 4
        width += margin * 2
        height += margin + margin * len(self.lines)
        
        positions = []
        y = margin
        for h in heights:
            positions.append(y)
            y += h + margin
        
        self._width = width
        self._height = height
        self._positions = positions
        self._margin = margin
        
        self._requested_size = (width, height)
    
    def draw(self, draw, rect):
        self.update()
        
        position = (rect[0], rect[1], )
        
        x = position[0]
        for line, y in zip(self.lines, self._positions):
            draw.text((x, position[1] + y), line, font = self.font, fill = self.foreground)


class Legend(View):
    def __init__(self, source, font):
        super(Legend, self).__init__(font)
        self.source = source
        self.lines = Lines(font)
    
    def _update(self):
        origin = self.source.header.get("ORIGIN", "-")
        lines = [" Origin: %s " % (origin, ), ]
        
        if "DATE" in self.source.header:
            date = self.source.header["DATE"]
            lines += [" Date: %s " % (date, ), ]
        if "DATE-OBS" in self.source.header:
            date = self.source.header["DATE-OBS"]
            lines += [" Observation date: %s " % (date, ), ]
        
        self.lines.lines = lines
        self.lines.update()

        self._requested_size = self.lines.requested_size
    
    def _draw(self, draw, rect):
        self.lines.draw(draw, rect)


class Converter(object):
    def __init__(self):
        self.data_fn = FN_LOG
        self.color_fn = self.default_color_fn
        self.file_name_fn = self.default_file_name_fn
        
        self.minv = 0
        self.maxv = 1
        self.mean = 0
        self.var  = 1
        
        self.minv_override = None
        self.maxv_override = None
        self.output_width  = None
        self.output_jpeg   = False
        self.verbose       = False
        
        self.font = ImageFont.load_default()
    
    def _write(self, message, *args):
        if not self.verbose:
            return
        if len(args) > 0:
            sys.stdout.write(message % tuple(args))
        else:
            sys.stdout.write(message)
    
    def default_color_fn(self, value, minv, maxv, mean, var):
        return (value - minv) / (maxv - minv)
    
    def default_file_name_fn(self, fn, frmt):
        return fn + "." + frmt
    
    def convert(self, image, file_name):
        if "NAXIS" not in image.header:
            raise "Image doesn't have \"NAXIS\" header."
        if image.header["NAXIS"] <> 2:
            raise "Expected 2 dimensional image."
        
        axes = [
            ImageAxis.from_image(image, 0),
            ImageAxis.from_image(image, 1),
        ]
        
        width, height, minv, maxv, mean, var = get_image_info(image)
        
        write = self._write
        write("Raw data:\n")
        write("size = %d, %d\n", width, height)
        write("x axis = %s\n", axes[0])
        write("y axis = %s\n", axes[1])
        write("min  = %f\n", minv)
        write("max  = %f\n", maxv)
        write("mean = %f\n", mean)
        write("var  = %f\n", var)
        
        minv = self.minv_override or minv
        maxv = self.maxv_override or maxv
        
        data = self.data_fn(image)
        
        minv = data.min()
        maxv = data.max()
        mean = data.mean()
        var = data.var()
        
        write("\nTransformed data:\n")
        write("min  = %f\n", minv)
        write("max  = %f\n", maxv)
        write("mean = %f\n", mean)
        write("var  = %f\n", var)
        
        result = Image.new("L", (width, height))
        MAX_PIXEL = 255
        
        i = 0
        pixels = numpy.zeros(width * height, numpy.uint8)
        
        progress = TextProgressBar(end_text = "Done.", enabled = self.verbose)
        progress.start()
        
        for y in range(height):
            row = self.color_fn(data[y], minv, maxv, mean, var) * MAX_PIXEL
            
            pixels[i:(i + len(row))] = row
            i += len(row)
            
            progress.update(float(y) / float(height))
        
        progress.end()
        
        result.putdata(pixels)
        
        if ((self.output_width is not None) and
            (self.output_width < width)):
            write("Resizing image...\n")
            ratio = float(self.output_width) / float(width)
            size = (self.output_width, int(float(height) * ratio))
            result = result.resize(size, Image.ANTIALIAS)
            for axis in axes:
                axis.scale = ratio
        
        #size = result.size
        #result = result.crop((0, -100, size[0], size[1]))
        
        #draw = ImageDraw.Draw(result)
        #self.draw_legend(image, draw, (10, 10), axes)
        result = self.draw_legend(image, result, (0, 0), axes)
        
        if self.output_jpeg:
            result.save(self.file_name_fn(file_name, "jpeg"), "JPEG")
        else:
            result.save(self.file_name_fn(file_name, "png"), "PNG")
    
    def draw_legend(self, source, dest, position, axes):
        size = dest.size
        
        legend = Legend(source, self.font)
        legend_size = legend.requested_size
        
        x_axis, y_axis = axes
        
        y_view = AxisView(self.font, y_axis, True)
        y_size = y_view.requested_size
        
        x_view = AxisView(self.font, x_axis, False)
        x_size = x_view.requested_size
        
        print "legend_size = %r" % (legend_size, )
        dest = dest.crop((-y_size[0], -legend_size[1], size[0], size[1] + x_size[1]))
        
        draw = ImageDraw.Draw(dest)
        legend.draw(draw, (y_size[0], 0, y_size[0] + size[0], legend_size[1]))
        #y_view.draw(draw, (0, legend_size[1], y_size[0], legend_size[1] + y_size[1], ))
        y_view.draw(draw, (0, legend_size[1], y_size[0], y_size[1], ))
        x_view.draw(draw, (y_size[0], legend_size[1] + size[1], x_size[0], x_size[1], ))
        
        return dest


class TextProgressBar(object):
    def __init__(self, output = sys.stdout, end_text = "", enabled = True):
        self.output = output
        self.value = 0
        self.end_text = end_text
        self.enabled = enabled
    
    def start(self):
        if not self.enabled:
            return
        self.output.write("\r%d%%" % (self.value * 100, ))
        self.output.flush()
    
    def update(self, value):
        self.value = value
        if not self.enabled:
            return
        self.output.write("\r%d%%" % (self.value * 100, ))
        self.output.flush()
    
    def end(self):
        if not self.enabled:
            return
        if self.end_text and len(self.end_text) > 0:
            self.output.write("\r%s\n" % (self.end_text, ))
        else:
            self.output.write("\r")
        self.output.flush()


def main():
    parser = argparse.ArgumentParser(description = "FITS to PNG converter.")
    parser.add_argument("files", metavar = "FILE", nargs = "+",
                        help = "a FITS file to convert")
    parser.add_argument("--headers", action = "store_true",
                        help = "show FITS headers")
    parser.add_argument("--jpeg", action = "store_true",
                        help = "output JPEG instead of PNG")
    parser.add_argument("--fn", metavar = "DATA_FUNCTION",
                        nargs = "?", default = None,
                        help = "data transformation function, identity by default")
    parser.add_argument("--min", metavar = "MIN_VALUE",
                        type = float, nargs = "?", default = None,
                        help = "minimal value cut-off")
    parser.add_argument("--max", metavar = "MAX_VALUE",
                        type = float, nargs = "?", default = None,
                        help = "maximal value cut-off")
    parser.add_argument("--width", metavar = "WIDTH",
                        type = int, nargs = "?", default = None,
                        help = "output image width in pixels")
    parser.add_argument("--verbose", action = "store_true",
                        help = "be verbose")
    args = parser.parse_args()
    
    converter = Converter()
    converter.minv_override = args.min
    converter.maxv_override = args.max
    converter.output_width  = args.width
    converter.output_jpeg   = args.jpeg
    converter.verbose       = args.verbose
    
    converter.data_fn = {
        "id":  FN_IDENTITY,
        "log": FN_LOG,
    }.get(args.fn, FN_LOG)
    
    for file_name in args.files:
        f = None
        if args.verbose:
            print "%s %s %s" % (10 * "=", file_name, 10 * "=", )
        try:
            f = pyfits.open(file_name)
            if args.headers:
                print f[0].header
                continue
            converter.convert(f[0], file_name)
            f.close()
        #except Exception as e:
        #    print "\nERROR: %s" % (e, )
        finally:
            if f:
                f.close()


if __name__ == "__main__":
    main()