#!/usr/bin/env python

import cv2
import os
import help_parser as hp
import lxml.etree as ET
from xml.dom import minidom
import numpy as np
import warnings

fw_orig = warnings.formatwarning
warnings.formatwarning = lambda msg, categ, fname, lineno, line=None: fw_orig(msg, categ, fname, lineno, '')

def parse_args():
	parser = hp.HelpParser(description='Create bounding box for masking data')
	parser.add_argument('dir', help='Folder to process')
	parser.add_argument('datadir', help='Folder with rgb data files')
	parser.add_argument('xml', help='XML file to save as an output')
	parser.add_argument('-cls', help='String to use as class in resulting XML file [person]', default='person')
	parser.add_argument('-name', help='Dataset name to fill in XML file [victims]', default='Victims')
	args = parser.parse_args()
	return args
	

def run_folder(args):
	files = [f for f in os.listdir(args.dir) if os.path.isfile(os.path.join(args.dir, f))]
	files.sort()
	xml = ET.Element('folder')
	xml.set('path', args.datadir)
	for f in files:
		ret = bbox_im(args, f)
		if ret is not None:
			xml.append(ret)
	return xml
		

def bbox_im(args, f):
	im = cv2.imread(os.path.join(args.dir, f))
	if im is None:
		return None
	xml = ET.Element('file')
	xml.set('path', os.path.join(args.datadir,f.replace('_mask','_rgb')))
	gray = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY)
	ret,thresh = cv2.threshold(gray,127,255,0)
	contours, hierarchy = cv2.findContours(thresh,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)
	for i, c in enumerate(contours):
		if hierarchy[0][i][3] == -1:
			elem = ET.Element('object')
			elem.set('class', args.cls)
			elem.set('prob', 'gt')
			x,y,w,h = cv2.boundingRect(c)
			bb = (x, y, x+w, y+h)
			elem.set('bbox', str(bb))
			xml.append(elem)
	return xml

if __name__ == '__main__':
	args = parse_args()
	if os.path.isfile(args.xml):
		warning.warn('File ' + args.xml + ' will be overwritten')
		
	xml = run_folder(args)
	root = ET.Element('Ground_truth')
	root.set('dataset', args.name)
	root.append(xml)
	xml = root
	rough = ET.tostring(xml)
	with open(args.xml, mode='w') as f:
		minidom.parseString(rough).writexml(f, addindent='\t', newl='\n')
