WP3: Hybrid Model v3 qp

# Author: Inigo Sedano and Patrik Hummelbrunner.
# Version 3.2, 17-12-2011
# Usage: python hybrid_v3_qp.py [XML file or XML file in .gz to parse]
# Python 3.2
# Python 2.x change __next()__ to next()

import xml.etree.cElementTree as ET
import time
import sys, gzip


if len(sys.argv) < 2:
        print("No file name specified as argument")
        sys.exit()
else:
        ## using gzip.open to be able to open gzip
        if sys.argv[1].endswith("gz"):
                fileToParse = gzip.open(sys.argv[1])
        elif sys.argv[1].endswith("HIMX"):
                fileToParse = sys.argv[1]
        else: 
                print("Fileformat needs to be .gz or HIMX")
                sys.exit()
start = time.time()

# get an iterable
context = ET.iterparse(fileToParse, events=("start","end"))

# turn it into an iterator
context = iter(context)

# get the root element
### if use python 2.x need to use .next()
event, root = context.__next__()

#  output file
outputFile = open("hybridQP_"+fileToParse.name+".csv","w")

# QP module ----
numPictures = 0
totalMacros = 0
maxQP=0
minQP=100
sumQP=0.0
avgQP=0.0
# ---- QP module

for event, elem in context:
    if event == "end" and elem.tag == "Picture":

        # QP module ----
        # This version is a modified version of the paper:
        # Analysis of Freely Available Subjective Dataset for HDTV including Coding and Transmission Distortions, Fifth International Workshop on Video         # Processing and Quality Metrics for Consumer Electronics (VPQM-10), Scottsdale, Arizona, January 13-15 2010.
        # Authors: Marcus Barkowsky, Margaret Pinson, Romuald Pepion, Patrick Le Callet.
        # The modification calculates the average QP when the QP is not constant.
        qpList = elem.findall("./SubPicture/Slice/MacroBlock/QP_Y")
        for macroQP in qpList:
                if int(macroQP.text) > maxQP:
                        maxQP = int(macroQP.text)
                if int(macroQP.text) < minQP:
                        minQP = int(macroQP.text)
                sumQP = sumQP + int(macroQP.text)
        numPictures = numPictures + 1
        totalMacros = totalMacros + len(qpList)
        # ---- QP module

        root.clear()

# QP module ----
print("Number of pictures: "+str(numPictures))
print("Total number of macroblocks: "+str(totalMacros))
if maxQP != minQP:
        print("Max QP macroblock value: %i" % maxQP)
        print("Min QP macroblock value: %i" % minQP)
        print("The QP (Quantization Parameter) is not constant, using the average QP (sum of the QP values of the macroblocks divided by the number of macroblocks)")
avgQP = sumQP / totalMacros
print("QP value: %.2f" %avgQP)
if avgQP < 28:
        print("The MOS (Mean Opinion Score) estimated is higher than 4.4 (range 1-5)")
		# in .csv should be only a digit output
        outputFile.write(fileToParse.name+";4.7")
elif avgQP > 44:
        print("The MOS (Mean Opinion Score) estimated is lower than 1.7 (range 1-5)")
		# in .csv should be only a digit output
        outputFile.write(fileToParse.name+";1.3")
else:
        estimatedMOS = -0.172 * avgQP + 9.249
        print("Estimated MOS (Mean Opinion Score): %.2f (range 1-5)" % estimatedMOS)
        outputFIle.write(fileToParse.name+";"+estimatedMOS)
# ---- QP module

elapsed = (time.time() - start)
print("Time elapsed (approx): %i seconds\n" % elapsed)