image 2 FIB and FIB 2 image - requires openCV, wxPython, PIL, numpy

convert_fib_str.py

#!/usr/bin/python
""" image 2 FIB and FIB 2 image
Usage:
  convert_fib_str.py fromstr <str> [--video] [--out=<out_file_path>]
  convert_fib_str.py tostr <image> [--out=<out_file_path>] [--stretch] [--invert]
"""
# from __future__ import print_function
import traceback
import sys
import math
import os
import struct
import copy
from PIL import Image, ImageDraw
import PIL.ImageOps
import cv2
import numpy as np

#import gui_related
import wx
app = wx.App(False)
#print wx.GetDisplaySize()

bmp=[]
pixes=[]
pix=[]


class fib_pattern(object):
    def __init__(self, filename):
        self.size = (4096, 4096)             # size of the image to create
        self.maxX=0
        self.maxY=0
        self.minY=self.size[1]
        self.minX=self.size[0]

        self.lines=[]
        print('loading image file: {}'.format(filename))
        print(os.path.isfile(filename))
        self.filename = filename
        self.filename_prefix = os.path.splitext(filename)[0]
        self.file_opened = False
        #with open(os.path.join(os.path.dirname(__file__), filename), 'rb') as fl:
    def open_str(self):
        with open(self.filename, 'rb') as fl:
            self.lines=fl.read()
        print 'A5 at: ', self.lines.index('\xA5')
        
        for i, x, y, dwell in self._iterate_pixels():
            pass
        self.file_opened = True

    def _xy_to_3bytes(self, x, y):
        #something to do with how nybble order gets packed by C compilers
        nib1 = y & 0xF
        nib2 = x >> 8
        b1 = x & 0xFF
        b3 = y >> 4
        b2 = (nib1<<4) | nib2
        return b1, b2, b3

    def _3bytes_to_xy(self, b1, b2, b3):
        #something to do with how nybble order gets packed by C compilers
        nib1 = b2 >> 4       # move data right     10010110 -> 00001001
        nib2 = b2 & 0xF      # keep rightside data 10010110 -> 00000110
        x = (nib2<<8) | b1   #             00000000bbbbbbbb -> 00000110bbbbbbbb
        y = (b3<<4) | nib1   #             00000000cccccccc -> 0000cccccccc1001
        return x, y

    def _iterate_pixels(self, _yield=False):
        #i = self.lines.index('\xd7') - 3
        num_pix = 0
        i = 20
        while True:
            try:    
                b1 = ord(self.lines[i])
                b2 = ord(self.lines[i+1])
                b3 = ord(self.lines[i+2])
                dwell = b4 = ord(self.lines[i+3])
                num_pix += 1

                x, y = self._3bytes_to_xy(b1, b2, b3)
                assert (b1, b2, b3) == self._xy_to_3bytes(x, y), 'x {} y {} orig {} return {}'.format(x, y, (b1, b2, b3), self._xy_to_3bytes(x, y))

                if _yield:
                    yield(i, x, y, dwell)
                else:
                    #keep some statistics
                    if x>self.maxX:
                        self.maxX=x
                    if y>self.maxY:
                        self.maxY=y
                    if x<self.minX:
                        self.minX=x
                    if y<self.minY:
                        self.minY=y
                #advance by 4 bytes
                i+=4
            except IndexError:
                break
        if not _yield:
            print "maxX %d, maxY %d" % (self.maxX, self.maxY)
            print "minX %d, minY %d" % (self.minX, self.minY)
            print 'num pixels: %d calced %d' % (num_pix, (i/4))
            self.numPixels = num_pix

    def write_12_bit_to_image(self, out_file_path):
        if not self.file_opened:
            self.open_str()
        lines = self.lines
        assert (self.size[0] > self.maxX) and (self.size[1] > self.maxY), 'size {} maxX {} maxY {}\n(self.size[0] == self.maxX) ({})\n(self.size[1] == self.maxY) ({})'.format(self.size, self.maxX, self.maxY, (self.size[0] > self.maxX), (self.size[1] > self.maxY))
        #fl1=open('out.txt1', 'wb')
        #i=lines.index('\xd7') - 3
        
        #size = (4096, 4096)             # size of the image to create
        im = Image.new('L', self.size) # create the image, 'L' for greyscale
        for i, x, y, dwell in self._iterate_pixels(True):
            #write the pixel to the image
            im.putpixel((x,y), dwell)
        im.save('{}.png'.format(self.filename_prefix) if not out_file_path else out_file_path)

    def write_12_bit_to_screen_as_video(self):
        if not self.file_opened:
            self.open_str()
        lines = self.lines
        #fl1=open('out.txt1', 'wb')
        # i=lines.index('\xd7') - 3
        # print 'i', i 
        
        # size = (4096, 4096)             # size of the image to create
        #im = Image.new('L', size) # create the image, 'L' for greyscale
        
        width = (self.maxX - self.minX) + 1
        height = (self.maxY - self.minY) + 1

        print 'width %d height %d maxX %d minX %d maxY %d minY %d' % (width, height, self.maxX, self.minX, self.maxY, self.minY)

        #out = cv2.VideoWriter('output.avi',cv2.cv.CV_FOURCC('D', 'I', 'V', 'X'), 20.0, (width,height))
        #fourcc = cv2.cv.CV_FOURCC('m', 'p', '4', 'v') # note the lower case
        #fourcc =  cv2.cv.CV_FOURCC('M','J','P','G')
        #fourcc = cv2.cv.CV_FOURCC('m', 'p', '4', '2')
        #fourcc = cv2.cv.CV_FOURCC('i','Y','U', 'V')#'D','I','V','X')  #CAP_PROP
        #fourcc = cv2.cv.CV_FOURCC(*'XVID')
        #print fourcc
        #out = cv2.VideoWriter()
        #success = out.open('output.avi',0,30,( width, height),True)
        m = max(width, height)
        #out = cv2.VideoWriter("test.avi", cv2.cv.CV_FOURCC('i','Y','u','v'), #'F', 'M', 'P', '4'),'H', '2', '6', '4'
        #                                    30, (width, height), False)  #import cv syntax CreateVideoWriter
        #out = cv2.VideoWriter("testmp42.mp4", cv2.cv.CV_FOURCC('M','P','4','2'), #'F', 'M', 'P', '4'),'H', '2', '6', '4'     'D','I','V','3' (works oks)
        #                                    90, (width, height), False) #'X','2','6','4'    width, height   
        #out = cv2.VideoWriter("testmp4.avi", cv2.cv.CV_FOURCC('D','I','V','3'), 30, (width, height), True)
        #out = cv2.VideoWriter("test_x264.avi", cv2.cv.CV_FOURCC('X','2','6','4'), 30, (width, height), True)
        #print 'video opened? %s' % out.isOpened()
        #print success
        blank_image = np.zeros((width, height,1), np.uint8)
        #blank_image = cv2.cv.CreateImage((m, m), 8, 1)
        #cv_image = cv2.cv.CreateImage((height, width), cv2.IPL_DEPTH_8U,1)
        cv2.namedWindow('FIB pattern raster sequence')#, cv2.CV_WINDOW_NORMAL|cv2.CV_WINDOW_KEEPRATIO | cv2.CV_GUI_EXPANDED)
        lastKey=0 
        #available_screen_size = gui_related.get_available_screen_size()
        
        topBarHeight, sideBarWidth, widthAvailable, heightAvailable = wx.ClientDisplayRect()
        heightAvailable -= topBarHeight # account for the highGui window's title bar

        if width > widthAvailable:
            scaleWidth=float(widthAvailable)/width
        else:
            scaleWidth = 1
        if height > heightAvailable:
            scaleHeight = float(heightAvailable)/height
        else:
            scaleHeight = 1
        scale = 1# min(scaleWidth, scaleHeight)
        scaled_frame_size = (int(width*scale), int(height*scale))

        for i, x, y, dwell in self._iterate_pixels(True):

            #write the pixel to the image
            blank_image[(y-self.minY, x-self.minX)] = dwell
            #advance by 4 bytes
            i+=4
            if i%256==0:
                pass
            if (i/4%(self.numPixels/100)==1):
                print 'percent done: %d' % (float(i/4) / self.numPixels *100)
                if not lastKey == 27:
                    cv2.imshow('FIB pattern raster sequence',cv2.resize(blank_image, scaled_frame_size) )
                lastKey = cv2.waitKey(30)                 # Waits forever for user to press any key
                print i/4, self.numPixels

        print "maxX %d, maxY %d" % (self.maxX, self.maxY)
        print "minX %d, minY %d" % (self.minX, self.minY)
        print 'num pixels: %d' % (i/4)
        print('press escape when done')
        k = 0
        while k != 27 and int(unicode(k&255)) != 27:
            k = cv2.waitKey(0)
            cv2.imshow('FIB pattern raster sequence',cv2.resize(blank_image, scaled_frame_size) )        
        cv2.destroyAllWindows()

    def image_to_str(self, _stretch=False, _invert=False, out_file_path = False):
        final_w, final_h = self.size
        if self.filename.endswith('.svg'):
            # import cairo
            # import rsvg
            # print('WARNING: converting from SVG to PNG first!!!')
            # handle = rsvg.Handle(self.filename)
            # w, h, ww, hh =  handle.get_dimension_data()

            # #img = cairo.ImageSurface(cairo.FORMAT_ARGB32, w, h)
            # img = cairo.ImageSurface(cairo.FORMAT_ARGB32, h, w)
            # ctx = cairo.Context(img)
            
            # ## or, for in memory SVG data:
            # # handle= rsvg.Handle(None, str(<svg data>))
            # handle.render_cairo(ctx)
            new_name = '{}_converted_from_svg.png'.format(self.filename_prefix)
            # img.write_to_png(new_name)
            # self.filename = new_name
            from wand.api import library
            import wand.color
            import wand.image

            svg_file = open(self.filename)
            with wand.image.Image() as image:
                # with wand.color.Color('transparent') as background_color:
                #     library.MagickSetBackgroundColor(image.wand, 
                #                                      background_color.resource) 
                # DPI???
                image.read(blob=svg_file.read(), resolution=300)
                w =  image.width
                h =  image.height
                print ('svg width {} height {}'.format(w, h))
                png_image = image.make_blob("png32")
            svg_file.close()

            # if max((w, h))< max(self.size):
            #     svg_file = open(self.filename)
            #     _scale = int(round(max(self.size)/max((w, h))))
                
            svg_file.close()

            with open(new_name, "wb") as out:
                out.write(png_image)
            self.filename = new_name

        im = Image.open(self.filename)
        im_w, im_h = im.size
        im_to_paste = None

        if _invert:
            inverted_image = PIL.ImageOps.invert(im.convert('L'))
            inverted_image.save('{}_inverted.png'.format(self.filename_prefix))
            im = inverted_image

        # print(im.format, im.size, im.mode)
        if im.size != self.size and max(im.size)<=max(self.size) and not _stretch:
            print('WARNING, image size is {} but target STR size is: {}  image will be centered in STR'
                  .format(im.size, self.size))
            offset = ((final_w - im_w) / 2, (final_h - im_h) / 2)
            im_to_paste = im
        elif im.size == self.size:
            if im.mode != 'L':
                ready_to_convert = im.convert('L')
            else:
             ready_to_convert = im
        elif _stretch:
            max_edge = max(im.size)
            min_edge = min(im.size)
            scaling = max(self.size) / float(max_edge)
            resized_w = int(round(scaling * im_w))
            resized_h = int(round(scaling * im_h))
            offset = ((final_w - resized_w) / 2, (final_h - resized_h) / 2)
            im_to_paste = im.resize((resized_w, resized_h))
        else:
            raise NotImplementedError('unsupported: image size {} probably needs stretched to {}'
                                      .format(im.size, self.size))
        
        if im_to_paste:
            blank = Image.new('L', self.size, (0))
            blank.paste(im_to_paste, offset)
            ready_to_convert = blank
            ready_to_convert.save('{}_ready_to_convert.png'.format(self.filename_prefix))
        
        if not out_file_path:
            out_file_path = '{}_ready_to_use.str'.format(self.filename_prefix)

        bin_file = open(out_file_path, 'wb')
        for i in range(20):
            bin_file.write(struct.pack('B', 0))
        # PIL 0,0 is in top-left of image
        for x in range(im_w):
            for y in range(im_h):
                try:
                    dwell = ready_to_convert.getpixel((x,y))
                    if not dwell:
                        continue
                except IndexError as e:
                    print('x {} y {} im_w {} im_h {}'.format(x, y, im_w, im_h))
                    import traceback
                    traceback.print_exc()
                    raise e
                b1, b2, b3 = self._xy_to_3bytes(x, y)
                # for b in [b1, b2, b3, dwell]:
                #     bin_file.write(struct.pack('b', b))
                bin_file.write(struct.pack('B', b1))
                bin_file.write(struct.pack('B', b2))
                bin_file.write(struct.pack('B', b3))
                bin_file.write(struct.pack('B', dwell))
        bin_file.close()


def write_image(img):
    w = max([p[0] for p in img])
    h = max([p[1] for p in img])

    size = (4096, 4096)             # size of the image to create
    im = Image.new('L', size) # create the image
    draw = ImageDraw.Draw(im)   # create a drawing object that is
                                # used to draw on the new image
    red = (255,0,0)    # color of our text
    table = sum(zip(*img), ())
    #draw.point(img)
    im.putdata(table)

    """
    #fff7fb
    #ece7f2
    #d0d1e6
    #a6bddb
    #74a9cf
    #3690c0
    #0570b0
    #034e7b
    """

    del draw # I'm done drawing so I don't need this anymore
    im.save('out_test1.png', 'PNG')

if __name__ == '__main__':
    from docopt import docopt
    args = docopt(__doc__)
    if args['fromstr']:
        pat = fib_pattern(args['<str>'])
        if args['--video']:
            pat.write_12_bit_to_screen_as_video()
        else:
            #pat.write_as_12_bit(copy.deepcopy(lines))
            pat.write_12_bit_to_image(args['--out'])
    elif args['tostr']:
        pat = fib_pattern(args['<image>'])
        pat.image_to_str(args['--stretch'], args['--invert'], args['--out'])
    #write_as_byte_w_offset(lines)
    #write_image(pixes)

    sys.exit(1)