Mercurial > repos > nick > dna_visualizer
comparison 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 |
comparison
equal
deleted
inserted
replaced
-1:000000000000 | 0:5257ce9d9184 |
---|---|
1 #!/usr/bin/env python | |
2 from __future__ import division | |
3 import os | |
4 import sys | |
5 import Image | |
6 import argparse | |
7 import fastareader | |
8 | |
9 OPT_DEFAULTS = {'size':'512x512', 'verbose':True, | |
10 'A':'0,255,0', 'T':'255,0,0', 'G':'255,255,255', 'C':'0,0,255'} | |
11 USAGE = "%(prog)s [options] genome.fasta" | |
12 DESCRIPTION = """Convert DNA sequence into a PNG image by representing each base | |
13 with one colored pixel.""" | |
14 EPILOG = """""" | |
15 | |
16 def main(): | |
17 | |
18 parser = argparse.ArgumentParser( | |
19 description=DESCRIPTION, usage=USAGE, epilog=EPILOG) | |
20 parser.set_defaults(**OPT_DEFAULTS) | |
21 | |
22 parser.add_argument('fasta', metavar='genome.fasta', | |
23 help="""Input sequence. Can be in FASTA format or a plain text file | |
24 containing only the sequence. Any non-ATGC characters (case-insensitive) | |
25 will be skipped.""") | |
26 parser.add_argument('-s', '--size', | |
27 help="""The output image size, in pixels, in the format "widthxheight", e.g. | |
28 "640x480". If the sequence is larger than the number of pixels in the | |
29 image, it will be cut off. Default size: %(default)s""") | |
30 parser.add_argument('-o', '--outfile', metavar='image.png', | |
31 help="""Output filename. Overrides the default, which is | |
32 to use the input filename base plus .png.""") | |
33 parser.add_argument('-d', '--display', action='store_true', | |
34 help="""Display the image instead of saving it.""") | |
35 parser.add_argument('-c', '--clobber', action='store_true', | |
36 help="""If the output filename already exists, overwrite it instead of | |
37 throwing an error (the default).""") | |
38 parser.add_argument('-v', '--verbose', action='store_true', | |
39 help="""Verbose mode. On by default.""") | |
40 parser.add_argument('-q', '--quiet', action='store_false', dest='verbose', | |
41 help="""Quiet mode.""") | |
42 group = parser.add_argument_group('Color customization', """Use these options | |
43 to use custom colors for bases. Specify with a comma-delimited RGB value | |
44 like "100,150,10".""") | |
45 group.add_argument('-A', metavar='R,G,B', | |
46 help="""default: %(default)s""") | |
47 group.add_argument('-T', metavar='R,G,B', | |
48 help="""default: %(default)s""") | |
49 group.add_argument('-G', metavar='R,G,B', | |
50 help="""default: %(default)s""") | |
51 group.add_argument('-C', metavar='R,G,B', | |
52 help="""default: %(default)s""") | |
53 | |
54 args = parser.parse_args() | |
55 | |
56 try: | |
57 size = parse_size(args.size) | |
58 except ValueError: | |
59 parser.print_help() | |
60 fail('\nError: Invalid size string "%s".' % args.size) | |
61 | |
62 fasta = fastareader.FastaLineGenerator(args.fasta) | |
63 bases = fasta.bases() | |
64 | |
65 if not args.display: | |
66 outfile = args.outfile if args.outfile else outfile_name(args.fasta) | |
67 if os.path.exists(outfile) and not args.clobber: | |
68 fail('Error: Output filename already taken: "%s"' % outfile) | |
69 | |
70 colors = {} | |
71 colors['A'] = parse_rgb(args.A) | |
72 colors['T'] = parse_rgb(args.T) | |
73 colors['G'] = parse_rgb(args.G) | |
74 colors['C'] = parse_rgb(args.C) | |
75 | |
76 image = Image.new('RGB', size, 'white') | |
77 pixels = image.load() | |
78 | |
79 done = False | |
80 for i in range(image.size[1]): | |
81 for j in range(image.size[0]): | |
82 try: | |
83 base = next(bases).upper() | |
84 except StopIteration: | |
85 done = True | |
86 break | |
87 if base in colors: | |
88 pixels[j,i] = colors[base] | |
89 if done: | |
90 break | |
91 | |
92 if args.display: | |
93 image.show() | |
94 else: | |
95 image.save(outfile, 'PNG') | |
96 | |
97 | |
98 | |
99 def parse_size(size_str): | |
100 """Parse size string, return a tuple of (width, height). | |
101 Accepts size strings in the format "640x480". | |
102 If not valid, raises ValueError.""" | |
103 size = map(int, size_str.split('x')) | |
104 if len(size) != 2: | |
105 raise ValueError | |
106 else: | |
107 return tuple(size) | |
108 | |
109 | |
110 def parse_rgb(rgb_str): | |
111 """Parse RGB string, return a tuple of (R, G, B). | |
112 If not valid, raises ValueError.""" | |
113 rgb = map(int, rgb_str.split(',')) | |
114 if len(rgb) != 3: | |
115 raise ValueError | |
116 else: | |
117 return tuple(rgb) | |
118 | |
119 | |
120 def outfile_name(infilename): | |
121 base = infilename.split('.')[0] | |
122 if not base: | |
123 base = infilename | |
124 return base+'.png' | |
125 | |
126 | |
127 def fail(message): | |
128 sys.stderr.write(message+"\n") | |
129 sys.exit(1) | |
130 | |
131 if __name__ == '__main__': | |
132 main() |