tgautil.py (7540B)
1 ''' 2 Copyright 2020 Electronic Arts Inc. 3 4 This program is is free software: you can redistribute it and/or modify it under the terms of 5 the GNU General Public License as published by the Free Software Foundation, 6 either version 3 of the License, or (at your option) any later version. 7 8 This program is is distributed in the hope that it will be useful, but with permitted additional restrictions 9 under Section 7 of the GPL. See the GNU General Public License in LICENSE.TXT 10 distributed with this program. You should have received a copy of the 11 GNU General Public License along with permitted additional restrictions 12 with this program. If not, see https://github.com/electronicarts/CnC_Remastered_Collection 13 ''' 14 import argparse 15 import io 16 import json 17 from PIL import Image 18 import os 19 import StringIO 20 import sys 21 import zipfile 22 23 def overwrite_prompt(question, default=False): 24 prompt = " [Y/n] " if default else " [y/N] " 25 while True: 26 sys.stdout.write(question + prompt) 27 choice = raw_input().lower() 28 if choice == '': 29 return default 30 elif choice == 'y': 31 return True 32 elif choice == 'n': 33 return False 34 else: 35 sys.stdout.write("\n") 36 37 def crop(tga_file): 38 with Image.open(tga_file) as image: 39 image = image.convert('RGBA') 40 alpha = image.split()[-1] 41 left, top, right, bottom = 0, 0, image.width, image.height 42 found_left, found_top, found_right, found_bottom = False, False, False, False 43 for y in range(0, image.height): 44 if found_top and found_bottom: 45 break 46 for x in range(0, image.width): 47 if found_top and found_bottom: 48 break 49 if not found_top and alpha.getpixel((x, y)) != 0: 50 top = y 51 found_top = True 52 if not found_bottom and alpha.getpixel((x, image.height - y - 1)) != 0: 53 bottom = image.height - y 54 found_bottom = True 55 for x in range(0, image.width): 56 if found_left and found_right: 57 break 58 for y in range(top, bottom): 59 if found_left and found_right: 60 break 61 if not found_left and alpha.getpixel((x, y)) != 0: 62 left = x 63 found_left = True 64 if not found_right and alpha.getpixel((image.width - x - 1, y)) != 0: 65 right = image.width - x 66 found_right = True 67 tga_data = StringIO.StringIO() 68 meta = None 69 if left == 0 and top == 0 and right == image.width and bottom == image.height: 70 image.save(tga_data, 'TGA') 71 else: 72 image.crop((left, top, right, bottom)).save(tga_data, 'TGA') 73 meta = json.dumps({ 74 'size': [image.width, image.height], 75 'crop': [left, top, right, bottom] 76 }, separators=(',',':')) 77 return (tga_data.getvalue(), meta) 78 79 def expand(tga_data, meta, tga_file): 80 with Image.open(io.BytesIO(tga_data)) as image: 81 if meta: 82 crop = meta['crop'] 83 image_size = (crop[2] - crop[0], crop[3] - crop[1]) 84 image = image.resize(image_size) 85 expanded_crop = (crop[0], crop[1], crop[2], crop[3]) 86 expanded_size = (meta['size'][0], meta['size'][1]) 87 with Image.new('RGBA', expanded_size, (0, 0, 0, 0)) as expanded: 88 expanded.paste(image, expanded_crop) 89 expanded.save(tga_file) 90 else: 91 image.save(tga_file) 92 93 def zip(args): 94 if not os.path.isdir(args.directory): 95 print >> sys.stderr, '\'{}\' does not exist or is not a directory\n'.format(args.directory) 96 sys.exit(1) 97 tga_files = [f for f in os.listdir(args.directory) if os.path.isfile(os.path.join(args.directory, f)) and os.path.splitext(f)[1].lower() == '.tga'] 98 if not tga_files: 99 print >> sys.stderr, '\'{}\' does not contain any TGA files\n'.format(args.directory) 100 sys.exit(1) 101 out_file = os.path.basename(os.path.normpath(args.directory)).upper() + '.ZIP' 102 if os.path.exists(out_file): 103 if not os.path.isfile(out_file): 104 print >> sys.stderr, '\'{}\' already exists and is not a file\n'.format(out_file) 105 sys.exit(1) 106 if not args.yes and not overwrite_prompt('\'{}\' already exists, overwrite?'.format(out_file)): 107 sys.exit(0) 108 with zipfile.ZipFile(out_file, 'w', zipfile.ZIP_DEFLATED) as zip: 109 for tga_file in tga_files: 110 tga_data, meta = crop(os.path.join(args.directory, tga_file)) 111 zip.writestr(tga_file, tga_data) 112 if meta: 113 zip.writestr(os.path.splitext(tga_file)[0] + '.meta', meta) 114 print 'Wrote ZIP archive \'{}\''.format(out_file) 115 116 def unzip(args): 117 if not os.path.isfile(args.archive): 118 print >> sys.stderr, '\'{}\' does not exist or is not a file\n'.format(args.archive) 119 sys.exit(1) 120 out_dir = os.path.normpath(os.path.splitext(args.archive)[0]) 121 if os.path.exists(out_dir): 122 if not os.path.isdir(out_dir): 123 print >> sys.stderr, '\'{}\' already exists and is not a directory\n'.format(out_dir) 124 sys.exit(1) 125 if len(os.listdir(out_dir)) > 0: 126 if not args.yes and not overwrite_prompt('\'{}\' is not empty, overwrite?'.format(out_dir)): 127 sys.exit(0) 128 else: 129 os.mkdir(out_dir) 130 files = {} 131 with zipfile.ZipFile(args.archive, 'r', zipfile.ZIP_DEFLATED) as zip: 132 for filename in zip.namelist(): 133 fileparts = os.path.splitext(filename) 134 name, ext = fileparts[0].lower(), fileparts[1].lower() 135 data = files.setdefault(name, {'tga': None, 'meta': None}) 136 if data['tga'] is None and ext == '.tga': 137 data['tga'] = zip.read(filename) 138 elif data['meta'] is None and ext == '.meta': 139 data['meta'] = json.loads(zip.read(filename).decode('ascii')) 140 if data['tga'] is not None and data['meta'] is not None: 141 expand(data['tga'], data['meta'], os.path.join(out_dir, name) + '.tga') 142 del files[name] 143 for name, data in files.items(): 144 expand(data['tga'], None, os.path.join(out_dir, name) + '.tga') 145 print 'Extracted files to \'{}\''.format(out_dir) 146 147 parser = argparse.ArgumentParser(description='TGA archive utility.') 148 subparsers = parser.add_subparsers() 149 150 parser_zip = subparsers.add_parser('z', help='Build a ZIP archive from a directory of TGA files.') 151 parser_zip.add_argument('directory', help='Directory of TGA files.') 152 parser_zip.add_argument('-o', '--out', nargs='?', help='Output archive path (defaults to input directory name with ZIP extension in the current path).') 153 parser_zip.add_argument('-y', '--yes', action='store_true', help='Confirm overwrite of existing ZIP archives.') 154 parser_zip.set_defaults(func=zip) 155 156 parser_unzip = subparsers.add_parser('u', help='Extract a ZIP archive of TGA files to a directory.') 157 parser_unzip.add_argument('archive', help='ZIP archive of TGA files.') 158 parser_unzip.add_argument('-o', '--out', nargs='?', help='Output directory (defaults to directory with name of the ZIP archive in the current path).') 159 parser_unzip.add_argument('-y', '--yes', action='store_true', help='Confirm overwrite of files in output directory.') 160 parser_unzip.set_defaults(func=unzip) 161 162 args = parser.parse_args() 163 args.func(args)