Exemplo n.º 1
0
class UmbrellaTracker:
    def __init__(self): 
        self.init = False

        #Are we paused
        self.pause = False
        #training mode enabled 
        self.train = False
        self.tracking = False
        self.useMog = True
        self.showBGDiff = False
        self.lastTimeWarningSent = 1
        self.alertInterval = 10000
        #seperate bg weight ratio's
        self.weightBG1 = 0.2
        self.weightBG2 = 0.6

        self.trained = {}
        self.counter = 0
        self.trainCouter = 0
        self.erodeIter = 1
        #gausion weight ratio
        self.blurKp = 1
        #cutoff threshold for BW image
        self.cutOffThresh=30;
        #what size to limit our bounding boxes too
        self.sizeL = 4000
        self.sizeM = 1500
        #kernal size for erode and dilate
        self.kernalH = 3
        self.kernalW = 3
        self.kernel = np.ones((self.kernalH,self.kernalW),'uint8')
        #kernal size
        self.rec = Recognizer()
        self.currFrame = None
        self.currFrameOpt = None

        self.contours = []
        #tracking will start after this many frames
        self.start = 10
        #track interval don't do tracking every frame
        self.track_interval = 20

        #Use MOG BG extractor
        self.bgs = cv2.BackgroundSubtractorMOG(24*60, 1
                                               , 0.8, 0.5)
        self.svmTracker = SVM()
        self.svmReady = True
        
        # Disabling twitter for the time being - not sure how its configured.
        #self.twtr = TwitterAlert()

    def onMouseClick(self,event, x, y, flags, param ):
        #select which roi to train for features
        if event == cv.CV_EVENT_LBUTTONDOWN:
            if len(self.contours) > 0:
                for cont in self.contours:
                    if self.isWithinBB(x,y,cont):
                        #x,y,w,h = cv2.boundingRect(cont)
                        box = cv2.boundingRect(cont)
                        if self.train:
                            print "Trainig for umbrella"
                            name = "./train/pos/%s.png"%self.trainCouter
                            self.trained[name] = box
                            cv2.imwrite("./train/pos/%s.png"%self.trainCouter,self.currFrame)
                            #self.rec.createTargetFeatures(self.getROI(self.currFrame,box),box)
                            #self.svmTracker.train(self.getROI(self.currFrame,box),np.float32)
                            #self.svmReady = True
                        else:
                            print "Trainig for None umbrella"
                            name = "./train/neg/%s.png"%self.trainCouter
                            self.trained[name] = box
                            cv2.imwrite("./train/neg/%s.png"%self.trainCouter,self.currFrame)

                        self.trainCouter = self.trainCouter + 1 
                               
            else:
                print "wait for contours to get initialized"

    def isWithinBB(self,x,y,cont):
        #check weather a point clicked on screen is within a bounding box defined by contours
        xb,yb,wb,hb = cv2.boundingRect(cont)
        xbC = xb + wb
        ybC = yb + hb
        if x > xb and x <xbC:
            if y > yb and y < ybC:
                return True
        return False

    def getROI(self,image,box):
        #cut out and return the section from an image extraceted from a contour
        if len(box) == 4:
            #make sure we have for courners seems sometimes we don'e #paranoid
            x,y,w,h = box
            return image[y:y+h,x:x+w]
        return None

    #get Video stream
    def addGausianBlur(self,f1,level):
        return cv2.GaussianBlur(f1,(level,level),1)

    #get diff 
    def subtractFrames(self,f1,f2):
        dif = cv2.absdiff(f1,f2)
        #dif = cv2.bitwise_and(f1,f2)
        return dif

    #get image threshold to remove some noise
    def getThreshold(self,f,cutOffThreshVal):
        return cv2.threshold(f,cutOffThreshVal,255,0)

    #find the contours in the image
    def findCountoursBW(self,f):
        conts,hier = cv2.findContours(f,cv2.RETR_LIST,cv2.CHAIN_APPROX_SIMPLE)
        #check size of countours ignore ones that are too small or large
        return conts

    def drawContours(self,image,conts,sizeL,sizeM,color):
        #draw contors on image
        for cnt in conts:   
            x,y,w,h = cv2.boundingRect(cnt)
            cv2.rectangle(image,(x,y),(x+w,y+h), color,2)
        return image

    def counterSizeFilter(self,contours,sizeL,sizeM):
        #filter out contours that fit a certain size
        filtered = []
        for cnt in contours:   
            x,y,w,h = cv2.boundingRect(cnt)
            area = float(w)*float(h)
            if area < sizeL and area > sizeM:
                filtered.append(cnt)
        return filtered

    def drawContour(self,image,cont,color):
        #TODO create a version that which works for single contours
        x,y,w,h = cv2.boundingRect(cont)
        cv2.rectangle(image,(x,y),(x+w,y+h), color,2)
        return image

    def drawContours2(self,image,conts,sizeL,sizeM,color):
        #draw contors on image given array od contours in contour
        for cnt in conts:   
            x,y,w,h = cnt[0],cnt[1],cnt[2],cnt[3]
            #area = float(w)*float(h)
            #if area < sizeL and area > sizeM:
            cv2.rectangle(image,(x,y),(x+w,y+h), color,2)
        return image

    def drawBox(self,image,x,y,size,color):
        cv2.rectangle(image,(x,y),(x+size,y+size),color)
        return image
    
    # TODO - does this comment go with cv2.rectangle or vc2.cvtColor ?    
    #this really slows things down!

    def makeBW(self,f):
        """
        Overview of cvtColor -  Converts an image from one color space to another.

            C++: void cvtColor(InputArray src, OutputArray dst, int code, int dstCn=0 )
            Python: cv2.cvtColor(src, code[, dst[, dstCn]]) -> dst
        """    
        # http://docs.opencv.org/modules/imgproc/doc/miscellaneous_transformations.html#cvtcolor
        bwImg = cv2.cvtColor(f,cv.CV_RGB2GRAY)
        return bwImg
    
    def createTrainigOutput(self):
        count = 1
        fpP = open("info.dat","w")
        fpN = open("bg.txt","w")
        for key in self.trained.keys():
                if key.find("neg") > 0:
                    output = "%s \n"%(key)
                    fpN.write(output)
                else:
                    output = "%s  %s  %s %s %s %s \n"%(key,count,self.trained[key][0],self.trained[key][1],self.trained[key][2],self.trained[key][3])
                    fpP.write(output)
        fpP.close()
        fpN.close()
    #Home Brew BG removal
    def bgSmoothing(self,frameBG,arr1, arr2):
 
        frameBG = cv2.equalizeHist(frameBG)
        #blurring helps to normalize and smooth out noise
        arr1 = self.addGausianBlur(arr1,self.blurKp)
        arr2 = self.addGausianBlur(arr2,self.blurKp)
        #add to wighted averages of the background in order to obtained weighterd average to remove noise
        cv2.accumulateWeighted(frameBG,arr1,self.weightBG1)
        cv2.accumulateWeighted(frameBG,arr2,self.weightBG2)
        #normalize
        res1 = cv2.convertScaleAbs(arr1)
        res2 = cv2.convertScaleAbs(arr2)
        
        #res1 = self.makeBW(res1)
        #res2 = self.makeBW(res2)
        return res1,res2

    def bgSmooothing2(self,frame):
        #OpenCV BG removal using MOG
        #frameBW = self.makeBW(frame)
        frame = self.addGausianBlur(frame,self.blurKp)
        frame = cv2.equalizeHist(frame)       
        fgMask = self.bgs.apply(frame)
        return fgMask

    def track(self, frame_raw, video_capture):
        """
        Primary function for rain detection - main loop to read the video input, 
        extract a frame, and process it.
        """
        frame_black_white = self.makeBW(frame_raw)
        #create numpy arrays from image frames
        avg1 = np.float32(frame_black_white)
        avg2 = np.float32(frame_black_white)
        objs = []
        while 1:
            if not self.pause:
                #get frame from video
                _, frame_raw = video_capture.read()
                
                # Check to see if we have a valid frame so we
                # don't get a strange error from opencv. 
                # http://stackoverflow.com/questions/16703345/how-can-i-use-opencv-python-to-read-a-video-file-without-looping-mac-os
                if (type(frame_raw) == type(None)):
                    break

                frame_black_white = self.makeBW(frame_raw)
                res2 = frame_black_white.copy()
                #make BW first

                if self.useMog:
                    mask = self.bgSmooothing2(res2)
                else:
                    res1,res2 = self.bgSmoothing(frame_black_white,avg1,avg2)
                    #get diff
                    mask = self.subtractFrames(res1,res2)
                    #lets threshold the image
                    _,mask = self.getThreshold(mask,self.cutOffThresh)

                #Dilate and erode
                #res3 = cv2.dilate(res3,kernel,iterations=3)
                #res3 = cv2.erode(res3,kernel,iterations=2)
                res3 = cv2.erode(mask,self.kernel,iterations=self.erodeIter)
            
                #set this for later use
                self.currFrame = res2
                self.currFrameOpt = res3

                cimage = np.copy(res3)
                #find countours
                if self.counter > self.start:
                    self.contours  = self.findCountoursBW(cimage)
                    self.contours = self.counterSizeFilter(self.contours,self.sizeL,self.sizeM)

                #make it color again
                #lets do some interesting stuff
                if len(self.contours) > 0 :
                    for cont in self.contours:
                        if self.counter % self.track_interval == 0:
                            box = self.getROI(res2,cv2.boundingRect(cont))
                            #self.rec.findUmbrellas(box)
                            #dont search every frame
                            if self.tracking and self.counter % self.track_interval == 0:
                                for thing in self.contours:
                                    objs = self.rec.detectUmbrellaCascade(box)
                                    umb = thing


                if len(objs) > 0:
                    if  self.lastTimeWarningSent == 1 or self.counter - self.lastTimeWarningSent >  self.alertInterval:
                        print "Umbrella Detected!!!! run for cover!!!!"
                        #self.twtr.sendAlert()
                        self.lastTimeWarningSent = self.counter
                    else:
                        print "Just sent an alert so won't send one again but just FYI there is an umbrella"

                res2 = cv2.cvtColor(res2,cv2.COLOR_GRAY2BGR)
                if len(objs) > 0:
                    res2 = self.drawContour(res2,umb,(255,0,0))

                res2 = self.drawContours(res2,self.contours,self.sizeL,self.sizeM,(255,255,0))
                #res3 = cv2.dilate(res3,kernel)
                self.counter = self.counter + 1
                #cv2.imwrite("./pngs/image-"+str(counter).zfill(5)+".png", mask)
                #self.rec.getFeatures(res2,2)
                #res1 = self.rec.drawFeatures(res2)

                if len(self.rec.toTrack) > 0 and self.svmReady:
                        for person in self.contours:
                            #masks = self.rec.flannMatching(cv2.boundingRect(person))
                            matches = self.rec.trackObjects(res2)
                            #res3 = self.drawBox()
            if self.showBGDiff:
                cv2.imshow('bg1',res3)
            cv2.imshow('bg2',res2)
            if not self.init:
                cv2.setMouseCallback('bg2',self.onMouseClick,None);
                init = True

            #break if esc is hit
            k = cv2.waitKey(20) 
            if k == ord('t'):
                if self.train:
                    print "training for Negative samples"
                    self.train = False
                else:
                    print "training for positive samples"
                    self.train = True
            if k == ord('d'):
                if len(self.trained.keys()):
                    print "creating training output file.."
                    self.createTrainigOutput()
                    print "output training files ready"
                    self.trained = {}
                    self.trainCouter = 0
            if k == ord('r'):
                self.rec.reset()
            if k == ord(' '):
                print "paused"
                self.pause = not self.pause
            if k == ord('g'):
                self.tracking = not self.tracking
                print "Tracking:%s"%self.tracking
            if k == ord('b'):
                self.useMog = not self.useMog
                print "MOG bg extraction %s"%self.useMog
            if k == ord('x'):
                self.showBGDiff = not self.showBGDiff
                print "enable bg diff:%s"%self.showBGDiff
            if k == ord('+'):
                self.erodeIter = self.erodeIter + 1
                print "erode : %s"%self.erodeIter
            if k == ord('-') and self.erodeIter > 0:
                self.erodeIter = self.erodeIter -1
                print "erode : %s"%self.erodeIter
            if k == 27:
                break