#              Lane Tracking Annotation
#
#   (c) Scott & Linda Wills         2 November 2009

import os, Image, ImageDraw

SeqDir = "./seqs/"
RptDir = "./reports/"

Colors = ['red', 'green', 'blue', 'orange', 'purple', 'skyblue', 'yellow', 'pink']

def LaneTracker() :
    """ This routine walks through each sequence in the seqs directory, parses
    the report.txt file in that sequence. It annotates a copy of each frame with
    lane lines and a lane change symbol (when appropriate) using information
    contained in the report.txt file. """
    print "Lane Tracker: version 1.0"
    DirList = os.listdir(".")
    if SeqDir[2:-1] not in DirList :
        print "sequence directory %s is missing" % (SeqDir)
    elif RptDir[2:-1] not in DirList :
        print "report directory %s is missing" % (RptDir)
    else :
        Reports = os.listdir(RptDir)
        for Report in Reports :
            if Report[-4:] == '.txt' :
                print "annotating report %s" % (Report)
                (Seq, Start, End, FrameData) = Parse_Report(RptDir, Report)
                if Seq and Start and End :
                    Annotate_Seq(Seq, Start, End, SeqDir, RptDir, FrameData)

def Parse_Report(RptDir, Report) :
    """ This routine parses a report file and returns
    a frame dictionary of report data. """
    ReportFile = open(RptDir + Report, 'r')
    FrameData = {}
    Seq = Start = End = None
    for Line in ReportFile.readlines() :
        Line = Line.strip()
        Tokens = Line.split(':')
        Keyword = Tokens[0].strip()
        if Line == '' :
            continue
        elif Keyword == 'Seq' : # sequence name
            Seq = Tokens[1].strip()
        elif Keyword == 'Start' : # starting seq num
            Start = int(Tokens[1].strip())
        elif Keyword == 'End' : # ending seq num
            End = int(Tokens[1].strip())
        elif Keyword == 'FN' :  # frame entry
            FN = int(Tokens[1].strip())
            FrameData[FN] = []
        elif Keyword == 'LE' :  # lane entry
            FrameData[FN].append(eval(Tokens[1].strip()))
        elif Keyword == 'DE' :
            FrameData[FN] = [int(Tokens[1].strip())] + FrameData[FN]
        else :                  # do nothing
            print "can't parse report line '%s'" % (Line.strip())
    ReportFile.close()
    if not Seq :
        print "ERROR: Sequence name not specified"
    if not Start :
        print "ERROR: Starting frame number not specified"
    if not End :
        print "ERROR: Ending frame number not specified"
    return (Seq, Start, End, FrameData)

def Annotate_Seq(Seq, Start, End, SeqDir, RptDir, FrameData) :
    """ This routine copies and annotates each frame in the
    specified sequences with report information. The output
    annotated sequence is stored in the report directory. """
    if Seq not in os.listdir(SeqDir) :
        print "ERROR: sequence %s not in sequence directory" % (Seq)
        return
    if Seq not in os.listdir(RptDir) :
        print "   creating annotation directory %s" % (RptDir + Seq)
        os.mkdir(RptDir + Seq)
    for FN in range(Start, End+1) :
        Path = "%s%s/%05d.jpg" % (SeqDir, Seq, FN)
        Frame = Image.open(Path)
        Draw = ImageDraw.Draw(Frame)
        if FN in FrameData :
            DepartureEvent = FrameData[FN][0]
            Lanes = FrameData[FN][1:]
            Index = 0
            for Lane in Lanes :
                Draw_Lane(Lane, FN, Draw, Colors[Index%len(Colors)])
                # Draw.text((X+3, Y+3), '%02d' % (V), fill='white')
                Index += 1
            Draw_Departure_Event(DepartureEvent, FN, Draw)
        else :
            print "WARNING: Missing data for frame %d" % (FN)
        Path = "%s/%05d.jpg" % (RptDir + Seq, FN)
        Frame.save(Path, 'JPEG')

def Draw_Departure_Event(Data, FN, Draw, Color="orange", Xmax=1279, Ymax=719) :
    """ This routine draws a triangle pointing upward, right, or left for
    Data = 0, 1, -1, respectively."""
    Xb = (Xmax / 2) - 20
    Yb = Ymax - 100
    if Data == 0 :
        Draw.polygon([(Xb,Yb+20), (Xb+20,Yb), (Xb+40,Yb+20), (Xb+30,Yb+20), (Xb+30, Yb+40), \
                      (Xb+10,Yb+40), (Xb+10,Yb+20)], fill=Color)
    elif Data == 1 :
        Draw.polygon([(Xb,Yb+10), (Xb+20,Yb+10), (Xb+20,Yb), (Xb+40,Yb+20), (Xb+20, Yb+40), \
                      (Xb+20,Yb+30), (Xb,Yb+30)], fill=Color)
    elif Data == -1 :
        Draw.polygon([(Xb,Yb+20), (Xb+20,Yb), (Xb+20,Yb+10), (Xb+40,Yb+10), (Xb+40, Yb+30), \
                      (Xb+20,Yb+30), (Xb+20,Yb+40)], fill=Color)

def Draw_Lane(Data, FN, Draw, Color="red", Xmax=1279, Ymax=719) :
    """ This routine draws a multi-segment line on the frame. """
    Points = []
    for (X,Y) in Data :
        if X < 0 or X > Xmax or Y < 0 or Y > Ymax :
            print "   Warning: FN %d marker out of frame, adjusted" % (FN)
            if X < 0 :
                X = 0
            if X > Xmax :
                X = Xmax
            if Y < 0 :
                Y = 0
            if Y > Ymax :
                Y = Ymax
        Points.append(X)
        Points.append(Y)
    Draw.line(Points, fill=Color, width=8)

if __name__ == '__main__' :
    LaneTracker()