doomcc/dcc/extract.py

import dcc.config
import omg
import numpy as np
import wand.color
import wand.image

class Extract(dcc.config.base):
	def get_parser(self, prog_name):
		parser = super().get_parser(prog_name)
		parser.add_argument('wad')
		parser.add_argument('lump')
		return parser

	def take_action(self, parsed_args):
		wads = sorted(self.pwads.joinpath(self.wad).glob('*.wad', case_sensitive=False), reverse=True)

		for w in wads:
			try:
				# TODO: handle anything other than graphics.
				wad = omg.WadIO(w)
				gl = omg.Graphic(wad.read(parsed_args.lump))
				# With no arguments, convert() changes a paletted image to an RGB one.
				with wand.image.Image.from_array(np.array(gl.to_Image().convert())) as img:
					img.transparent_color(wand.color.Color("#ff00ff"), 0.0)
					img.save(filename=self.output.joinpath(parsed_args.wad).joinpath(parsed_args.lump + ".png"))
				return
			except Exception as e:
				print("Wad {} likely has no lump {} (exception {}).".format(w, parsed_args.lump, e))

		print("Lump {} not found in any wad in {}".format(parsed_args.lump, parsed_args.wad))