#! /Library/Frameworks/Python.framework/Versions/Current/bin/python

"""
Names:  James G. Malcolm
        Shawn M. Lankton
        Jesus H. Christ

Project Number: 3

Notes:

  To view intermediate representation
  set the following variable to False
"""

FINAL = True

"""
Summary of Algorithm:

  :Background Modeling:
  :Motion Estimation:     
  :Corridor Detection:
  :Adaptive Parameters:
  :Blip Detection:         ("Holy Crap! What's That?!")
  :Propogate People:       ("Move it along, boys")
  :Associate Blips:        ("I want that one!")
  :Update People:          ("So *that's* where you were")
  :Peoplize Orphan Blips:  ("You're a *real* boy")
  :Drop Low Confidence:    ("Take no prisoners!")
  :Update L-R Counts:      ("How many of you were there?")

Description of Algorithm:

  :Background Modeling:
     Update bg model only when pixels aren't moving
     Use "Learn Rate" and "Forget Rate" to govern rate of change
     (Refer to Lankton P1)

  :Motion Estimation:     
     Subtract image from model
     Smooth with median filter
     Normalize remaining data and weight RGB information from image

  :Corridor Detection:
     Threshold motion estimate
     Keep cumulative representation of binary motion
     Logical OR each frame's binary motion
     Compute row sums over binary image
     Compute cumulative sum over row sums
     Detect corridor at 20% and 80% density
     Drop bottom half of corridor (to ignore feet)

  :Adaptive Parameters:
     Based on the size of the detected corridor:
     update parameters based on linear relationships
     (see report... we have graphs!)

  :Blip Detection:
     Determine which columns have significant motion
     Filter to remove tiny noise
     Find groups with no movement between
     For each blip, compute a 'feature profile'
     'feature profile' is a weighted probability density estimate
     we use a 16x16x16 joint RGB color space

  :Propogate People:
     Update each person's position based on their estimated velocity

  :Associate Blips and People:
     Each person finds the 'best match' blip
     'best match' is a combination of spatial proximity and color
     color comparison is performed with the bhattacharyya measaure
     a person with a nearby, color-matched blips 'own' that blip
     people who find matching blips have confidence increased.
     people who did not are impugned!

  :Update People with Measurement:
     Based found blips, update people's velocity and position
     position and velocity are controled with recursive average filters
     we assume constant velocity models
     
  :Peoplize Orphan Blips:
     Blips that were not associated to people become new people
     New people start with intermediate confidence
     
  :Drop Low Confidence:    ("Take no prisoners!")
     people with 0 confidence are removed
     people whose position moves them off the screen are removed

  :Update L-R Counts:      ("How many of you were there?")
     When a person's confidence reaches a max, they are counted
     The count used corresponds to the person's direction

Insights:
  Problem:  There are too many people to count on dynamics to carry you through 
            occlusions
  Solution: We realized that color information would be needed to help
            identify people.  Initially we tried to use the mean of
            the peroson in RGB space.  This provided decent results,
            but it was insufficient to identify people across the
            whole sequence.  Next we moved to weighted histograms
            (weighted by probability of foreground) as representations
            and used the Bhattacharyya measure similarity between
            these density estimates of blip and person (see Comaniciu
            2003)

  Problem:  Much of the image contains useless (and noisy) data that
            slows computation and can mess up detection.
  Solution: We detected the corridor along which people were walking
            and only tracking in this region.  We began by computing
            the total motion along each row.  The corridor was taken
            to be the rows comprising the middle 80% of that motion.
            We used the verticle cumulative distribution of motion to
            determine this region.  From then on, we stayed within
            this narrow region (much like Luke Skywalker when he
            attacked the Death Star by flying through the tight
            corridors).


  Problem:  The sequences provided were very different.  The relative
            scale and speed of people made various parameters
            impossible to set at a single value Solution: We noticed
            that many of the parameters (relating to the approximate
            speed of pedestrians) correlated well to the size of the
            corridor.  Furthermore, our corridor detection was very
            robust and worked well on all sequences without changing
            parameters. We based our other parameters on the corridor
            estimate by first finding parameter settings that worked
            well for several sequences, and then finding a linear
            mapping that related them to corridor size.
            

Interesting Facts:
  Turkeys have great hearing but no ears.
  Coca-Cola would be green if not for coloring.
  Google does not search for Chuck Norris.  Chuck searches for Google.

Bibliography:
  We do all our own stunts, but took inspiration from the following Jedi.

  @Article{Comaniciu2003,
    author = {D. Comaniciu and V. Ramesh and P. Meer},
    title = {Kernel-based object tracking},
    journal = {Trans. on Pattern Analysis and Machine Intelligence},
    volume = 25,
    number = 5,
    year = 2003
  }

NOTE:: Expecte more discussion in our *FULL FINAL PAPER***
        (prepare to be amazed!)


"""

from FrameWork import *
from numpy import *
import ImageFilter
import numpy
import Image

#-- parameters for background subtraction
learn_rate = 1
forget_rate = 1
var_thresh = 50
smooth_size = 5
motion_thresh = 40

#-- parameter for corridor
corr_construction_frames = 150

#-- parameters for confidence (when to live and let die)
conf_max = 10;  #maximum value
conf_show = 5;  #value where display starts
conf_start = 3; #inital value
conf_kill = 0

#-- parameters for model comparison
bandwidth = 16
nearby_multiplier = 4

#-- paramters for motion model filter
gain_x  = 0.5; 
gain_dx = 0.5; # higher allows faster changes

#-- global variables (not paramaters!! OMG)
sz = 240,320
newID = 0
LeftCount = 0
RightCount = 0
nearby = 30
bhat_sum = .5
bhat_cnt = 1

class Model(object) :
    def __init__(self,img,weights) :
        self.density = self.create_density(img, weights)

    def create_density(self, data, w) :
        hist = zeros(bandwidth*ones(3))
        for c in range(0,data.shape[1]) :
            for r in range(0,data.shape[0]) :
                if w[r,c] == motion_thresh : continue
                bin = floor( data[r,c,:] / bandwidth )
                hist[bin[0],bin[1],bin[2]] += w[r,c]
        return hist / (sum(hist) + 1e-9) # normalize
            
    def compare(self, m) : #bhattacharyyayrraaaaya
        return sum(sqrt(self.density*m.density + 1e-9))

class Person(object) :
    def __init__(self, pos, model) :
        self.id = None
        self.x = pos # position
        self.dir = sign(sz[1]/2-pos)
        self.dx = 7*self.dir # velocity
        self.conf = conf_start # confidence [0,10]
        self.counted = False
        self.model = model

    def update(self, x) :
        #if (new_dx*self.dx>0) : #same sign
        self.dx = (1-gain_dx)*self.dx + gain_dx*(x - self.x)
        self.x  =  (1-gain_x)*self.x  +  gain_x*x

    def propagate(self) : #walk like a maaaaAAAAAAaaaaannnnnnnnn
        self.x += self.dx

    def compliment(self) :
        self.conf  = min(self.conf + 1, conf_max)
        if self.conf >= conf_max :
            self.conf = conf_max
            if self.counted == False :
                global LeftCount, RightCount
                if sign(self.dx) > 0 : RightCount += 1
                else :                 LeftCount  += 1
                self.counted = True

    def impugn(self) :
        self.conf = max(self.conf - 1, 0)

# add/remove from histogram
def enum(h, data, w, fn) :
    for i in range(0,data.shape[0]) : 	
        if w[i] == motion_thresh : continue
        bin = floor( data[i,:] / bandwidth )
        h[bin[0],bin[1],bin[2]] = fn(h[bin[0],bin[1],bin[2]], w[i])

def  yank(h, data, w) : 
    enum(h, data, w, lambda x,w : x - w)

def spank(h, data, w) : 
    enum(h, data, w, lambda x,w : x + w)


def Main(SeqName='20061129-01', start=1850, end=1859) :
    MyTrial = Trial(Update=False)
    sin = Seq(MyTrial, SeqName)
    sout = PedSeq(MyTrial, 'fg')

    #--initialize global vars
    global newID, LeftCount, RightCount, nearby, bhat_sum, bhat_cnt
    newID = 0
    LeftCount = 0
    RightCount = 0
    nearby = 30
    bhat_sum = .5
    bhat_cnt = 1
    global conf_max, conf_show, conf_start
    conf_max = 10 
    conf_show = 5
    conf_start = 3
    img = downsample(numpy.asarray(sin.Load_Frame(start)).astype('float32'))

    bg = img.copy()             # data model
    var = zeros((sz[0],sz[1]))    # data variance
    motion_acc = tile(False, var.shape)  # for corridor detection
    corridor = 0                         # initial corridor is null
    people = list()

    for fn in range(start, end+1) :
        #-- get current and previous frame
        pre = img
        img_orig = numpy.asarray(sin.Load_Frame(fn)).astype('float32')
        img = downsample(img_orig)
        bg, var = update_bg(bg, var, img, pre);

        #-- get motion corridor
        motion=sum(abs(bg-img),2)/3         # get motion map
        motion=medfilt(motion,smooth_size)
        corridor,motion_acc = get_corridor(fn-start, motion, motion_acc, corridor)

        motion[0:corridor[0],  :]=0         # clean motion map
        motion[corridor[-1]:-1,:]=0
        img_ = img[corridor,:,:]
        w_   = motion[corridor,:]

        #-- get 'blips' (where motion is present)
        blip_strip = 255*(sum(motion[corridor,:] > motion_thresh,0) > 5)
        blip_strip = filter_blips(blip_strip)
        blips = blipstrip2blips(blip_strip, img_, w_)

        #-- propagate people
        for p in people: p.propagate()

        #-- associate people to blips (don't hate)
        unassoc_blips = associate_blips(people, blips)

        #-- peoplize remaining blips
        for b in unassoc_blips : people.append(Person(b[0],b[1]))

        #-- kill the impugned (and also the insecure)
        people = filter(is_alive, people)

        ##############  FINAL OUTPUT E########################
        if(FINAL) : 
            pedlist = people2pedlist_final(people)
            result = Image.fromarray(img_orig.astype('uint8'))
        else : 
            pedlist = people2pedlist(people)
            motion = visualizations(corridor, img, motion, blip_strip, people)
            result = Image.fromarray(motion.astype('uint8'))
        sout.Mark_Frame(fn, result, pedlist, (LeftCount,RightCount))
        sout.Store_Frame(fn, result)
        ######################################################

    #print SeqName + '  corridor: ' + str(corridor[-1] - corridor[0])
    MyTrial.Print('L-R: ' + str(RightCount) + ' R-L: ' + str(LeftCount))
    MyTrial.Done()


#-- converts people 2 pedlist
def people2pedlist(people) :
    global newID
    people_ = filter(lambda p: p.conf > conf_show, people)
    for p in people_ : 
        if(p.id<0) : 
            p.id = newID
            newID += 1
    pedlist = map(lambda p: (p.conf, p.x, sign(p.dx)), people_)
    pedlist.append((conf_max,10,1))
    pedlist.append((conf_show,sz[1]-10,1))
    return pedlist

#-- converts people 2 pedlist
def people2pedlist_final(people) :
    global newID
    people_ = filter(lambda p: p.conf > conf_show, people)
    for p in people_ : 
        if(p.id<0) : 
            p.id = newID
            newID += 1
    pedlist = map(lambda p: (p.id, p.x*2, sign(p.dx)), people_)
    return pedlist

#-- let live if on screen or secure in identity
def is_alive(p) :
    x = p.x # + p.dx
    return (0 < x and x < sz[1]) and p.conf > conf_kill

#-- update people based on blips
def associate_blips(people,blips) :
    for p in people :
        global bhat_sum, bhat_cnt
        if len(blips) == 0 : return [] # done, it's over, we're finished, so stop

        bhat = map(lambda b : cmp_blip(p,b), blips)
        bhat_best = max(bhat)

        if(bhat_best) :
            bhat_cnt+=1
            bhat_sum+=bhat_best

        bhat_target = bhat_sum/bhat_cnt-.2
        if bhat_best < bhat_target :
            p.impugn()
        else :
            blips = associate(p, blips, bhat, bhat_best)
            p.compliment()
    return blips

#-- compare blip to person
def cmp_blip(p, (pos,m)) :
    if abs(p.x - pos) > nearby : return 0
    return m.compare(p.model)

#-- update person and drop from blips
def associate(p, blips, bhat, bhat_best) :
    for i in range(0, len(blips)) :
        if bhat[i] != bhat_best : continue
        b = blips.pop(i)
        p.update(b[0])
        return blips

#-- convert blip strip to positions
def blipstrip2blips(strip, img, w) :
    blips = list()
    cnt = 0
    for i in range(0,len(strip)) :
        if cnt and strip[i] : # increment
            cnt += 1
        elif cnt and strip[i] == 0 : # end
            xx = range(i-cnt,i-1)
            model = Model(img[:,xx,:],w[:,xx])
            blips.append((int(i - cnt/2), model))
            cnt = 0
        elif strip[i] :
            cnt = 1 # start
    return blips
        

#-- draw visualizations
def visualizations(corridor, img, motion, blips, people) :
    # blips
    motion[30:60] = tile(blips, [30,1])

    # grab corridor
    motion = g2rgb(motion)
    img_    = img[corridor,:,:]
    motion_ = motion[corridor,:]

    # motion weighted image
    motion[corridor] = img_ * motion_ / 100
    motion[motion>255] = 255
    motion[corridor,0:5,:] = 255

    # original image
    off = len(corridor) + 10
    corridor_ = map(lambda x : min(x+off, motion.shape[0]-1), corridor)
    motion[corridor_,:,:] = img[corridor,:,:]

    # people
    for p in people :
        motion[0:30,int(p.x),:] = 128
        if p.conf > conf_show : motion[0:30,int(p.x),1] = 255

    return motion

#-- run medfilt on blips (must make 2D first... fu pil)
def filter_blips(blips) :
    b = medfilt(tile(blips,[2,1]),5)
    return b[0,:]

def set_speed_unit(top,bot) :
    global nearby, conf_max, conf_show, conf_start
    h = bot - top
    speed_unit = 0.2744 * h + 5.2343
    nearby = speed_unit*nearby_multiplier

    conf_unit = -0.138 * h  + 13.396
    conf_max = max(int(conf_unit),3)
    conf_show = max(int(conf_unit/2),2)
    conf_start = max(int(conf_unit/3),1)
    #print 'nearby: ' + str(nearby) + ' conf_max: ' + str(conf_max)

#-- called every iteration to get corridor
def get_corridor(fn, motion, motion_acc, corridor) :
    if fn < corr_construction_frames :
        motion_acc = logical_or(motion_acc, motion>motion_thresh)
        corridor_top,corridor_bot = get_corridor_topbot(motion_acc)
        corridor_bot -= int( (corridor_bot - corridor_top)/2 )
        set_speed_unit(corridor_top, corridor_bot)
        if corridor_top==corridor_bot : return array([0]),motion_acc
        else: return range(corridor_top,corridor_bot), motion_acc
    else : return corridor, motion_acc

#-- determines top and bottom of corridor
def get_corridor_topbot(is_active) :
    corridor = cumsum(sum(is_active,1))
    max = corridor.max()
    if max < is_active.shape[1] : return 0,0  # activity threshold
    top = where(corridor > .2 * max)[0][0]
    bot = where(corridor > .8 * max)[0][0]
    height = float(bot-top);
    top = int(top-0.2*height)
    bot = int(bot+0.2*height)
    return top,bot

#-- update the background model
def update_bg(bg, var, img, pre) :
    #compute average frame-to-frame difference (recursive average)
    var = (var+sum(abs(img-pre),2))/(forget_rate*2)
    #get the ones that aren't changing
    mask = var<var_thresh*1
    maskrgb=g2rgb(mask)
    #put them into the model (recursive average)
    bg[maskrgb>0]=(bg[maskrgb>0]+learn_rate*img[maskrgb>0])/(1+learn_rate)
    return (bg,var)

#-- replace each pixel with the median of its neighbors
def medfilt(a,smooth_size):
  img = Image.fromarray(a.astype('uint8')).filter(ImageFilter.MedianFilter(smooth_size))
  return numpy.asarray(img).astype('float32')

#-- each pixel is replaced by the max of its neighbors
def maxfilt(u,win) :
    y = zeros(u.shape)
    for i in range(0,len(u)) :
        xx = range(max(0,i-win), min(i+win,len(u)))
        y[i] = min(255, max(u[xx]))
    return y

#-- reduce the size of the image
def downsample(a):
    return a[0:-1:2,0:-1:2]

#-- norm
def norm(x) : return sqrt(sum(x**2))

#-- converts a grayscale image to RGB
def g2rgb(a):
    r = zeros((sz[0],sz[1],3))
    r[:,:,0]=a
    r[:,:,1]=a
    r[:,:,2]=a
    return r

# these calls make the program run in one-shot!
# Main('20070403-09',1,750) #demo
# Main('20080411-01',1,750) #ours
# Main('20080411-02',957,1462) # CRC
# Main('20080411-03',1,750) #stucen
# Main('20080411-04',0,749) #mark&j
# Main('20080413-01',27,749) #fence (bigpeople)
Main('20080414-01',464,1000) # physics
# Main('20080414-02',0,599) #library
# Main('20080414-03',1,658) #mall
# Main('20080414-04',83,583) #van leer
# Main('20080414-05',5,580) #bful warner robbins
# Main('20080415-01',0,499) #cherry tree
# Main('20080415-02',0,749) #too fast
# Main('20080415-03',0,599) #skiles
# Main('20080416-01',1,1064) #too slow