

In [19]:

    
def nothing(x):
    pass

def hist_equalize_bgr(frame):
    '''
    Converts to YCrCb colorspace and applies histogram equalization to help with contrast. 
    Only applies to Y channel. Returns an RGB image.'''
    #Convert to     
    frame = cv2.cvtColor(frame,cv2.COLOR_RGB2YCrCb)
    frame[:,:,0] = cv2.equalizeHist(frame[:,:,0])
    frame = cv2.cvtColor(frame,cv2.COLOR_YCrCb2RGB)
    return frame

def smooth(frame, kern=(3,3)):
    '''
    Smooth an image using gaussian blur:
        kern:Shape of the gaussian kernel, larger numbers make the image more blurred. 
    '''
    return cv2.blur(frame, kern)

def resize_show(frame, params, window='Unmanipulated'):
    '''
    Resizes and image and displays it. 
    params: dictionary containing the resize ratio for the x direction, 'fx' and y direction 'fy'
    window: name of opencv2 window to display image. See cv2.imshow doc for more info. 
    '''
    frame = cv2.resize(frame,None,fx=params['fx'], fy=params['fy'], interpolation = cv2.INTER_CUBIC)
    cv2.imshow(window,frame)
    cv2.waitKey(1)

def convert_color(frame, conversion=cv2.COLOR_RGB2HSV):
    '''
    Convenience function to wrap converstion from RGB to HSV
    '''
    return cv2.cvtColor(frame, conversion) 

def apply_adaptive_thresh(frame,params):
    '''
    Equalizes the intensity values of a frame, converts to gray, uses adaptive thresholding and applies
    morphological opening to remove small bits of noise. 
    frame: image frame, assumed to be recently read from a video (bgr)
    params: dictionary containing
            kern: size of structuring element for mophological opening
            BS: Block size, defines the area by which 
    '''
    #Create a kernel
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(params['kern'],params['kern']))
    #Equalize histogram
    frame = hist_equalize_bgr(frame)
    #Convert to Grayscale
    frame = cv2.cvtColor(frame, cv2.COLOR_RGB2GRAY)
    #Threshold and remove small peices of noise. ADAPTIVE_THRESH_GAUSSIAN_C could be used... slightly better results 
    #but nontrivial cost in terms of processing speed. 
    frame = cv2.adaptiveThreshold(frame,255,cv2.ADAPTIVE_THRESH_MEAN_C,cv2.THRESH_BINARY_INV,params['BS'],params['C'])     
    frame = cv2.morphologyEx(frame, cv2.MORPH_OPEN, kernel)
    #Show Image, resized
    return frame 

def get_thresh_contours(frame, params):
    im2, contours, hierarchy = cv2.findContours(frame,cv2.RETR_TREE,cv2.CHAIN_APPROX_SIMPLE)
    
    goodCont = []
    for cnt in contours:
        area = cv2.contourArea(cnt)
        if  params['amin'] < area < params['amax']:
            goodCont.append(cnt)

    return goodCont
    
    
    
class TrackVideo():
    
    def __init__(self,video, params):
        ''' 
        Starts an instance of a video tracking objects. 
        keyword arguments:
            video: The video file to be tracked
            params: a dictionary containing the following
                fx:resize ratio in the x direction for display
                fy:resize ratio in the y direction for display
        '''
        
        self.video = video
        self.params = params 
    
        self.init()
        '''
        for idx in range(100):
        
            frame = self.update()
        '''
    
    def test_settings_adaptivethresh(self, nframes=1000):
        '''
        Runs nframes of adaptive thresholding and allows you to play with the key paramters. 
        Use save parameters method to save paramters after the window closes. 
        '''
        cv2.namedWindow('TestSettings')
        
        cv2.createTrackbar('Block Size','TestSettings',3,255,nothing)
        cv2.createTrackbar('C','TestSettings',1,255,nothing)
        cv2.createTrackbar('KERN','TestSettings',1,25,nothing)
        
        #Start with some OK defaults
        cv2.setTrackbarPos('Block Size', 'TestSettings', self.params['BS'])
        cv2.setTrackbarPos('C','TestSettings',self.params['C'])
        cv2.setTrackbarPos('KERN', 'TestSettings',self.params['kern'])
        
        for idx in range(nframes):
            #Get Settings, clean them and store
            self.params['BS'] = cv2.getTrackbarPos('Block Size','TestSettings')
            if self.params['BS'] % 2 ==0:
                self.params['BS'] +=1 #Needs to be odd
                cv2.setTrackbarPos('Block Size', 'TestSettings', self.params['BS'])
            if self.params['BS'] < 3:
                self.params['BS'] == 3
                cv2.setTrackbarPos('Block Size', 'TestSettings', self.params['BS'])
                
            self.params['C']  = cv2.getTrackbarPos('C','TestSettings')
            
            self.params['kern'] = cv2.getTrackbarPos('KERN','TestSettings')
            if self.params['kern'] % 2 == 0:
                self.params['kern'] +=1
                cv2.setTrackbarPos('KERN', 'TestSettings',self.params['kern'])
            if self.params['kern'] < 3:
                self.params['kern'] == 3
                cv2.setTrackbarPos('KERN', 'TestSettings',self.params['kern'])

            
        
            #Proceses and display Frame
            ret, frame = self.cap.read()
            frame = apply_adaptive_thresh(frame, self.params)
            resizedFrame = resize_show(frame,self.params,'TestSettings')

        cv2.destroyWindow('TestSettings')
        cv2.waitKey(1)
    
    def test_settings_ath(self,segment, nframes=500, amaxPossible=500, aminPossible=0):
        '''
        Allows you to test area threshold 
        segment: function used to transform image into binary image
        nframes: how long to allow you to play with settings
        amaxPossible: largest value for the scroll bar to search for a good amax
        aminPossible: smallest value for the scroll bar to search for a good amin
        '''
        #Make a window
        cv2.namedWindow('TestSettings')
        
        #add track bars, set them based on params
        cv2.createTrackbar('A Min','TestSettings',params['amin'],amaxPossible,nothing)
        cv2.createTrackbar('A Max','TestSettings',params['amax'],amaxPossible,nothing)
        
        #Run segmentation, get values, and show results. 
        for idx in range(nframes):
            self.params['amin'] = cv2.getTrackbarPos('A Min', 'TestSettings')
            self.params['amax'] = cv2.getTrackbarPos('A Max', 'TestSettings')
            ret, unmanipulatedFrame = self.cap.read()
            frame = segment(unmanipulatedFrame,self.params)
            
            #Find contours 
            cnts = get_thresh_contours(frame, self.params)
            out = cv2.drawContours(unmanipulatedFrame, cnts, -1, (0,255,255), 2)
            
            resizedFrame = resize_show(out,self.params,'TestSettings')

            
        cv2.destroyWindow('TestSettings')
        cv2.waitKey(1)
        
    def process_adaptive_thresh(self,show=True,save='pickle',saveLoc='./test.p'):
        
        self.cap.release()
        self.cap = cv2.VideoCapture(self.video)
        
        allCnts = []
        
        #while(cap.isOpened())
        for ii in range(300):
            ret, unmanipulatedFrame = self.cap.read()
            
            frame = apply_adaptive_thresh(unmanipulatedFrame, self.params)
            cnts = get_thresh_contours(frame, self.params)
            
            allCnts.append(cnts)
            
            if show:
                out = cv2.drawContours(unmanipulatedFrame, cnts, -1, (0,255,255), 2)
                resizedFrame = resize_show(out,self.params,'Processing')
                
                
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
        
        
        self.cap.release()
        
        #Refresh the cap incase it needs to be used later. 
        self.cap = cv2.VideoCapture(self.video)
        
        save(self,cnt)
        
    def save(self,cnt):
        pass
     
    def save_params(self):
        pass
    
    def load_params(self):
        pass
    
    def init(self):
        #Clean Params
        if self.params['method'] == 'adaptiveThresh':
            if self.params['BS'] % 2 ==0:
                self.params['BS'] +=1 #Needs to be odd
            if self.params['BS'] < 3:
                self.params['BS'] == 3
            if self.params['kern'] % 2 == 0:
                self.params['kern'] +=1
            if self.params['kern'] < 3:
                self.params['kern'] == 3

        #Start a cap and grab a frame. 
        self.cap = cv2.VideoCapture(self.video)
        ret, self.testImg = self.cap.read()

    def kill_all(self):
        self.cap.release()
        cv2.destroyAllWindows()
        cv2.waitKey(1)