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