diff literal.py @ 0:5257ce9d9184

Initial literal.py tool
author Nick Stoler <nstoler@psu.edu>
date Sun, 02 Mar 2014 13:51:03 -0500
parents
children 58160195728e
line wrap: on
line diff
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/literal.py	Sun Mar 02 13:51:03 2014 -0500
@@ -0,0 +1,132 @@
+#!/usr/bin/env python
+from __future__ import division
+import os
+import sys
+import Image
+import argparse
+import fastareader
+
+OPT_DEFAULTS = {'size':'512x512', 'verbose':True,
+  'A':'0,255,0', 'T':'255,0,0', 'G':'255,255,255', 'C':'0,0,255'}
+USAGE = "%(prog)s [options] genome.fasta"
+DESCRIPTION = """Convert DNA sequence into a PNG image by representing each base
+  with one colored pixel."""
+EPILOG = """"""
+
+def main():
+
+  parser = argparse.ArgumentParser(
+    description=DESCRIPTION, usage=USAGE, epilog=EPILOG)
+  parser.set_defaults(**OPT_DEFAULTS)
+
+  parser.add_argument('fasta', metavar='genome.fasta',
+    help="""Input sequence. Can be in FASTA format or a plain text file
+      containing only the sequence. Any non-ATGC characters (case-insensitive)
+      will be skipped.""")
+  parser.add_argument('-s', '--size',
+    help="""The output image size, in pixels, in the format "widthxheight", e.g.
+      "640x480". If the sequence is larger than the number of pixels in the
+      image, it will be cut off. Default size: %(default)s""")
+  parser.add_argument('-o', '--outfile', metavar='image.png',
+    help="""Output filename. Overrides the default, which is
+      to use the input filename base plus .png.""")
+  parser.add_argument('-d', '--display', action='store_true',
+    help="""Display the image instead of saving it.""")
+  parser.add_argument('-c', '--clobber', action='store_true',
+    help="""If the output filename already exists, overwrite it instead of
+      throwing an error (the default).""")
+  parser.add_argument('-v', '--verbose', action='store_true',
+    help="""Verbose mode. On by default.""")
+  parser.add_argument('-q', '--quiet', action='store_false', dest='verbose',
+    help="""Quiet mode.""")
+  group = parser.add_argument_group('Color customization', """Use these options
+    to use custom colors for bases. Specify with a comma-delimited RGB value
+    like "100,150,10".""")
+  group.add_argument('-A', metavar='R,G,B',
+    help="""default: %(default)s""")
+  group.add_argument('-T', metavar='R,G,B',
+    help="""default: %(default)s""")
+  group.add_argument('-G', metavar='R,G,B',
+    help="""default: %(default)s""")
+  group.add_argument('-C', metavar='R,G,B',
+    help="""default: %(default)s""")
+
+  args = parser.parse_args()
+
+  try:
+    size = parse_size(args.size)
+  except ValueError:
+    parser.print_help()
+    fail('\nError: Invalid size string "%s".' % args.size)
+
+  fasta = fastareader.FastaLineGenerator(args.fasta)
+  bases = fasta.bases()
+
+  if not args.display:
+    outfile = args.outfile if args.outfile else outfile_name(args.fasta)
+    if os.path.exists(outfile) and not args.clobber:
+      fail('Error: Output filename already taken: "%s"' % outfile)
+
+  colors = {}
+  colors['A'] = parse_rgb(args.A)
+  colors['T'] = parse_rgb(args.T)
+  colors['G'] = parse_rgb(args.G)
+  colors['C'] = parse_rgb(args.C)
+
+  image = Image.new('RGB', size, 'white')
+  pixels = image.load()
+
+  done = False
+  for i in range(image.size[1]):
+    for j in range(image.size[0]):
+      try:
+        base = next(bases).upper()
+      except StopIteration:
+        done = True
+        break
+      if base in colors:
+        pixels[j,i] = colors[base]
+    if done:
+      break
+
+  if args.display:
+    image.show()
+  else:
+    image.save(outfile, 'PNG')
+
+
+
+def parse_size(size_str):
+  """Parse size string, return a tuple of (width, height).
+  Accepts size strings in the format "640x480".
+  If not valid, raises ValueError."""
+  size = map(int, size_str.split('x'))
+  if len(size) != 2:
+    raise ValueError
+  else:
+    return tuple(size)
+
+
+def parse_rgb(rgb_str):
+  """Parse RGB string, return a tuple of (R, G, B).
+  If not valid, raises ValueError."""
+  rgb = map(int, rgb_str.split(','))
+  if len(rgb) != 3:
+    raise ValueError
+  else:
+    return tuple(rgb)
+
+
+def outfile_name(infilename):
+  base = infilename.split('.')[0]
+  if not base:
+    base = infilename
+  return base+'.png'
+
+
+def fail(message):
+  sys.stderr.write(message+"\n")
+  sys.exit(1)
+
+if __name__ == '__main__':
+  main()