human-segmentation/unit_tests.py

#!/usr/local/bin/python

import unittest
from python_core import *
import argparse
import string

import inspect

def FunctionName(): 
    return inspect.stack()[1][3]

def DownloadSource(input_file_base_name, input_url):

	if DirExists(input_file_base_name) == False:
		log.info("downloading " + input_file_base_name)
		ExecuteCommand("wget " + input_url)
		ExecuteCommand("unzip -o " + input_file_base_name)
		ExecuteCommand("rm -rf __MACOSX")
	else:
		log.info("Skip downloading Adeel_Abbas_20210611T145405 - folder already exists ...")
	
class Tests(unittest.TestCase):

	def verify_matte_file(self, matte_file: string, num_frames: int, width: int, height: int):
		videoInfo = VideoInfo(matte_file)
		self.assertEqual(num_frames, videoInfo.frame_count)
		self.assertEqual(width, videoInfo.width)
		self.assertEqual(height, videoInfo.height)	

	def __init__(self, *args, **kwargs):
		super(Tests, self).__init__(*args, **kwargs)
		self._num_frames = 100
		self._input_file = "Adeel_Abbas_20210611T145405/Adeel_Abbas_20210611T145405-1X.MP4"

	@unittest.skipIf(int(os.getenv('RUN_SKIPPED_TESTS', 0)) == 0, "washington model will not be unit tested by default since it does not get used in production")
	def test_washington(self):
		log.info("Running test_washington")
		output_dir = FunctionName() + "/"
		MakeDir(output_dir)
		output_file = output_dir + FunctionName() + ".mp4"

		num_frames = 17

		command = args.human_segmentation + " -i " + self._input_file + " -o " + output_file 
		command += " --slate Adeel_Abbas_20210611T145405/Adeel_Abbas_20210611T145405-1X-BG.PNG -n " + str(num_frames)
		ExecuteCommand(command)
		self.verify_matte_file(output_file, num_frames, 1920, 1440)

	def test_temporal(self):
		log.info("Running test_temporal")
		output_dir = FunctionName() + "/"
		MakeDir(output_dir)
		output_file = output_dir + FunctionName() + ".mp4"

		num_frames = 39

		command = args.human_segmentation + " -i " + self._input_file + " -o " + output_file + " -n " + str(num_frames)
		ExecuteCommand(command)
		self.verify_matte_file(output_file, num_frames, 1920, 1440)
  
	def test_sam_temporal(self):
		log.info("Running test_sam_temporal")
		output_dir = FunctionName() + "/"
		MakeDir(output_dir)
		output_file = output_dir + FunctionName() + ".mp4"

		num_frames = 39

		command = args.human_segmentation + " -i " + self._input_file + " -o " + output_file + " -n " + str(num_frames)
		ExecuteCommand(command)
		self.verify_matte_file(output_file, num_frames, 1920, 1440)
  
	def test_sps_temporal(self):
		log.info("Running test_sps_temporal")
		output_dir = FunctionName() + "/"
		MakeDir(output_dir)
		output_file = output_dir + FunctionName() + ".mp4"

		num_frames = 39

		command = args.human_segmentation + " -i " + self._input_file + " -o " + output_file + " -n " + str(num_frames)
		ExecuteCommand(command)
		self.verify_matte_file(output_file, num_frames, 1920, 1440)

	def test_batch_sizes(self):
		log.info("Running test_batch_sizes")
		# Test for case when number of frames is indivisible by batch size.
		
		output_dir = FunctionName() + "/"
		MakeDir(output_dir)
		
		num_frames = 93

		batch_sizes = [1, 2, 3, 5]
		for batch_size in batch_sizes:

			output_file = output_dir + FunctionName() + "_" + str(batch_size) + ".mp4"
			
			command = args.human_segmentation + " -n " + str(num_frames) + " -i " + self._input_file + " -o " + output_file + " --batch-size " + str(batch_size)								
			ExecuteCommand(command)
			self.verify_matte_file(output_file, num_frames, 1920, 1440)

	def test_active_timecodes(self):
		log.info("Running test_active_timecodes")

		output_dir = FunctionName() + "/"
		MakeDir(output_dir)
		output_file = output_dir + FunctionName() + ".mp4"

		active_timecodes = "500-3000"
		num_frames = 24 + 72 + 24
		batch_size = 7

		command = args.human_segmentation + " -i " + self._input_file + " -o " + output_file + " --active-timecodes " + str(active_timecodes) + " --batch-size " + str(batch_size)
		ExecuteCommand(command)	
		self.verify_matte_file(output_file, num_frames, 1920, 1440)

	def test_multipass(self):
		log.info("Running test_multipass")
		
		output_dir = FunctionName() + "/"
		MakeDir(output_dir)
		output_file = output_dir + FunctionName() + ".mp4"

		num_frames = 93

		output_file_pass1 = output_dir + FunctionName() + "-pass1.mp4"
		command = args.human_segmentation + " -i " + self._input_file + " -o " + output_file_pass1 + " -n " + str(num_frames)		
		ExecuteCommand(command)
		self.verify_matte_file(output_file_pass1, num_frames, 1920, 1440)

		output_file_pass2 = output_dir + FunctionName() + "-pass2.mp4"
		command = args.human_segmentation + " -i " + self._input_file + " -o " + output_file_pass2 + " -n " + str(num_frames) + " --input-matte " + output_file_pass1
		ExecuteCommand(command)
		self.verify_matte_file(output_file_pass2, num_frames, 1920, 1440)
		
def sanitize_args(args):
	return args

if __name__ == '__main__':

	parser = argparse.ArgumentParser(description='Human segmentation unit tests')
	parser.add_argument('--log-timestamp', type=int, default=DefaultLogTimestamp(), help=LogTimestampHelp())
	parser.add_argument('--log-level', type=str, default=DefaultLogLevel(), help=LogLevelHelp())

	parser.add_argument('--human-segmentation', type=str, default=GetFullPath("human-segmentation"), help="human-segmentation executable path [human-segmentation]")

	parser.add_argument('--tests-dir', type=str, default="unit_tests")

	args = parser.parse_args()

	InitLog(args.log_level, args.log_timestamp)
	del sys.argv[1:]

	args = sanitize_args(args)
	
	MakeDir(args.tests_dir)
	ChangeDir(args.tests_dir)

	DownloadSource("Adeel_Abbas_20210611T145405", "https://www.dropbox.com/s/eir0i1gx4fswcpg/Adeel_Abbas_20210611T145405.ZIP")

	unittest.main()