review_data.py

import summaryScriptAlignment
import xml.etree.ElementTree as ET
import random

movie_indices = {'butterfly': 0, 'shrek': 1, 'anastasia': 2}
movie_directories = ['The Butterfly Effect', 'Shrek', 'Anastasia (1997 film)']
class review_data:
	def __init__(self):
		self.true_pairs = {}
		self.false_pairs = {}
		for movie_code in movie_indices:
			folder = 'Aligned/' + movie_directories[movie_indices[movie_code]] + '/processed/'
			alignment_file_path = folder + 'standard ' + movie_code + ' alignment.txt'

			file = open(alignment_file_path)
			alignment = summaryScriptAlignment.SummaryScriptAlignment(file)
			alignment.getAlignment()
			for sentence in alignment.summaryScriptAlignment:
				for scene in alignment.summaryScriptAlignment[sentence]:
					if movie_code not in self.true_pairs:
						self.true_pairs[movie_code] = {}
					if sentence not in self.true_pairs[movie_code]:
						self.true_pairs[movie_code][sentence] = set()
					self.true_pairs[movie_code][sentence].add(scene)
			

			xml_file_path = folder + 'script.xml'
			tree = ET.parse(xml_file_path)
			root = tree.getroot()
			num_scene = 0
			for scene in root.iter('scene'):
				if num_scene < int(scene.get('count')):
					num_scene = int(scene.get('count'))

			for sentence in self.true_pairs[movie_code]:
				for scene in self.true_pairs[movie_code][sentence]:
					if movie_code not in self.false_pairs:
						self.false_pairs[movie_code] = {}
					if sentence not in self.false_pairs[movie_code]:
						self.false_pairs[movie_code][sentence] = set()
					while True:
						rand = random.randint(1, num_scene)
						if rand not in self.true_pairs[movie_code][sentence]:
							self.false_pairs[movie_code][sentence].add(rand)
							break

rd = review_data()

print('true_pairs:')
print(rd.true_pairs)

print('false_pairs:')
print(rd.false_pairs)