genome to image

genome2img.py

#!/usr/bin/env python3
# written by bruno02468 in 2020-06-18
# quick 'n dirty script to turn genomes into square images, pads with 0
# it's a mess 'cause it was written in an hour or so
# no rights reserved
# ref: https://twitter.com/cib0rges/status/1273634584220049409

try:
  from PIL import Image
except:
  print("Install PIL with pip3 install pillow.")
  raise

import sys, math

def nuc2num(nuc):
  try:
    return "acgt".index(nuc)
  except:
    return None

def quad2byte(quad):
  return sum([nuc2num(quad[i]) << 2*(3-i) for i in range(4)])

def bytes2img(lb):
  nucpixels = int(len(lb)/3)
  dim = math.ceil(math.sqrt(nucpixels))
  if nucpixels != dim*dim:
    print("Warning: pixels aren't a square number (image truncated).")
    lb += [0] * (dim*dim - nucpixels) * 3
  return (Image.frombytes("RGB", (dim, dim), bytes(lb)), dim)

def main(filename, outname, pixel_side):
  nuc = "catg"
  basebytes = []
  cbyte = []
  with open(filename, "r") as fd:
    c = fd.read(1).lower()
    while c:
      if c not in nuc:
        print("Warning: skipped invalid char 0x%02X." % (ord(c),))
      else:
        cbyte.append(c)
      if len(cbyte) == 4:
        basebytes.append(quad2byte(cbyte))
        cbyte = []
      c = fd.read(1).lower()
  print("Done reading: genome encodes %d bytes." % (len(basebytes),))
  missingnucs = (4 - len(cbyte)) % 4
  if missingnucs:
    print("Warning: 4 doesn't divide # of nucleotides (bytes truncated).")
    cbyte += ["a"] * missingnucs
    basebytes.append(quad2byte(cbyte))
    cbyte = []
  missingbytes = len(basebytes) % 3
  if missingbytes:
    print("Warning: 3 doesn't divide # of bytes (colors truncated).")
    basebytes += [0] * missingbytes * 3
  img, dim = bytes2img(basebytes)
  dim *= pixel_side
  img = img.resize((dim, dim), Image.NEAREST)
  print("Saved to %s (%dx%d)." % (outname, dim, dim))
  img.save(outname)
  return 0


if __name__ == "__main__":
  if len(sys.argv) < 3:
    print("Specify a text file and an image file name.")
    sys.exit(-1)
  else:
    pixel_side = int(sys.argv[3]) if len(sys.argv) > 3 else 5
    sys.exit(main(sys.argv[1], sys.argv[2], pixel_side))