#! /usr/bin/env python
#
# Copyright (C) Oct 2012 Angelos Tzotsos <tzotsos@gmail.com>		  
#
# gdal_segm_accuracy.py is an implementation of the methods described in:
# Liu et al, 2012 "Discrepancy measures for selecting optimal combination of 
# parameter values in object-based image analysis", ISPRS Journal of 
# Photogrammetry and Remote Sensing 68 pp. 144-156
# 
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

from __future__ import division
try:
    from osgeo import gdal
    from osgeo.gdalconst import *
except ImportError:
    import gdal
    from gdalconst import *

try:
    import numpy
except ImportError:
    import Numeric as numpy

import sys
import math

# =============================================================================
def Usage():
    print('Usage: gdal_segm_accuracy.py segmentation_file reference_file')
    print('')
    sys.exit( 1 )

# =============================================================================
seg_file = None
ref_file = None

# Parse command line arguments.
i = 1
while i < len(sys.argv):
    arg = sys.argv[i]

    if seg_file is None:
        seg_file = arg

    elif ref_file is None:
        ref_file = arg

    else:
        Usage()

    i = i + 1

if seg_file is None:
    Usage()
if  ref_file is None:
    Usage()

#Open segmentation and reference files
segdataset = gdal.Open( seg_file, GA_ReadOnly )
refdataset = gdal.Open( ref_file, GA_ReadOnly )

if (segdataset.RasterCount != 1):
    print('Error: Segmentation file has more than 1 band')
    sys.exit( 1 )
if (refdataset.RasterCount != 1):
    print('Error: Reference file has more than 1 band')
    sys.exit( 1 )
if ((segdataset.RasterXSize != refdataset.RasterXSize) or
(segdataset.RasterYSize != refdataset.RasterYSize)):
    print('Error: Segmentation and Reference file have different dimensions')
    sys.exit( 1 )
    
#Get bands
segband = segdataset.GetRasterBand(1)
refband = refdataset.GetRasterBand(1)

#Define classes
class robject:
    def __init__(self,area=0,under_area=0,over_area=0,overlap_area=0):
        self.area = area
        self.under_area = under_area
        self.over_area = over_area
        self.overlap_area = overlap_area
        self.corobj_area = 0
        self.corobjects = {}

#class segment:
#    area=0
#    id=0
#    
#class corobject:
#    common_area=0
#    id=0

#Using dictionaries instead of segment classes
segments={}
robjects={}

#Populate dictionaries by running both images once
for l in range(segband.YSize-1, -1, -1):
    scanline_seg = segband.ReadAsArray(0, l, segband.XSize, 1, segband.XSize, 1)
    scanline_ref = refband.ReadAsArray(0, l, refband.XSize, 1, refband.XSize, 1)
    print(' Processing images at %2.2f %%  \r' % 
        ((segband.YSize-1-l)*100/(segband.YSize-1) )),
    for c in range(0, segband.XSize):
        #read the position values
        val_seg = scanline_seg[0][c]
        val_ref = scanline_ref[0][c]
        
        #Find the segment in dictionary and increment area by one
        if val_seg in segments:
            segments[val_seg] = segments[val_seg] + 1
        else:
            segments[val_seg] = 1
                
        #Check if there is a reference object and add to area and common area
        if val_ref != 0:
            if val_ref in robjects:
                robjects[val_ref].area = robjects[val_ref].area + 1
                if val_seg in robjects[val_ref].corobjects:
                    robjects[val_ref].corobjects[val_seg] = \
                        robjects[val_ref].corobjects[val_seg] + 1
                else:
                    robjects[val_ref].corobjects[val_seg] = 1
            else:
                robjects[val_ref] = robject()
                robjects[val_ref].area = 1
                robjects[val_ref].corobjects[val_seg] = 1
        
#move output to the next line
print

#debugging print-out
print("Number of reference objects : %d" % len (robjects))
print("Number of segments : %d" % len (segments))
print("")
     
#For each reference object find final coresponding segments and statistics
for rid,robj in robjects.iteritems():
    #debugging print-out
    print("Number of initial corresponding objects in reference %d : %d" % 
        (rid,len (robj.corobjects)))
    to_delete = []
    #For each initial coresponding segment    
    for cid,cobj in robj.corobjects.iteritems():
        print("segment %d : common area %d , total area %d, reference area %d" % 
            (cid, cobj, segments[cid], robj.area))
        #Check robjects for corresponding segments
        if ((cobj/segments[cid]) < 0.5) and ((cobj/robj.area) < 0.5):
            #If not corresponding object is found, add common area to 
            #over-segmentation and delete object
            print("Decision: Not corresponding segment")
            robj.over_area = robj.over_area + cobj
            to_delete.append(cid)
        else:
            #If corresponding object compare common area with segment area and  
            #add to under-segmentation
            print("Decision: Corresponding segment")
            robj.under_area = robj.under_area + (segments[cid]-cobj)
            robj.overlap_area = robj.overlap_area + cobj
            robj.corobj_area = robj.corobj_area + segments[cid]
    for i in to_delete:
        del robj.corobjects[i]
    #debugging print-out
    print("Number of final corresponding objects in reference %d : %d" % 
        (rid,len (robj.corobjects)))
    print("Total area of reference object %d : %d" % (rid, robj.area))
    print("Total overlap area of reference object %d : %d" % 
        (rid, robj.overlap_area))
    print("Total over-segments area of reference object %d : %d" % 
        (rid, robj.over_area))
    print("Total under-segments area of reference object %d : %d" % 
        (rid, robj.under_area))
    print('-------------------------------------------------------------')

#Calculate global metrics
total_area = 0
total_overlap_area = 0
total_over_area = 0
total_under_area = 0
total_corobjects = 0
total_corobj_area = 0

for rid,robj in robjects.iteritems():
    total_area = total_area + robj.area
    total_overlap_area = total_overlap_area + robj.overlap_area
    total_over_area = total_over_area + robj.over_area
    total_under_area = total_under_area + robj.under_area
    total_corobjects = total_corobjects + len(robj.corobjects)
    total_corobj_area = total_corobj_area + robj.corobj_area

print("Reference Area: %d" % (total_area))
print("Total Overlap Area: %d" % (total_overlap_area))
print("Total Oversegmentation Area: %d" % (total_over_area))
print("Total Undersegmentation Area: %d" % (total_under_area))
print("Total Coresponding Segments: %d" % (total_corobjects))
print("Total Coresponding Segments Area: %d" % (total_corobj_area))

QR = 1 - (total_overlap_area/(total_overlap_area + total_over_area +  
    total_under_area))
OR = total_over_area/total_area
UR = total_under_area/total_corobj_area
ED1 = math.sqrt(((OR*OR)+(UR*UR))/2)
PSE = total_under_area/total_area
num_ref = len(robjects)
NSR = (math.fabs(num_ref-total_corobjects))/num_ref
ED2 = math.sqrt((PSE*PSE)+(NSR*NSR))

#Print Segmentation Accuracy

print('-------------------------------------------------------------')
print('Segmentation Accuracy')
print('-------------------------------------------------------------')
print('  QR  |  OR  |  UR  |  ED1  |  PSE  |  NSR |  ED2')
print(' %1.3f| %1.3f| %1.3f| %1.3f | %1.3f | %1.3f| %1.3f' %
          (QR, OR, UR, ED1, PSE, NSR, ED2))
print('-------------------------------------------------------------')