ok

src/utils/datagen.py

@@ -0,0 +1,159 @@
+import cv2
+import numpy as np
+
+from matplotlib.colors import rgb_to_hsv, hsv_to_rgb
+
+
+def rand(a=0, b=1):
+    return np.random.rand()*(b-a) + a
+
+
+def get_random_data (
+    annotation_line, 
+    input_shape, 
+    max_boxes=25, 
+    scale=.3, 
+    hue=.1, 
+    sat=1.5, 
+    val=1.5, 
+    random=True
+    ):
+
+    '''
+    random preprocessing for real-time data augmentation
+    '''
+    
+    line = annotation_line.split('\t')
+    h, w = input_shape
+    box = np.array([np.array(list(map(int,box.split(',')))) for box in line[1:]])
+
+    image = cv2.imread(line[0])
+    ih, iw, ic = image.shape
+    
+    if not random:
+        resize_scale = min(h/ih, w/iw)
+
+        nw = int(iw * resize_scale)
+        nh = int(ih * resize_scale)
+        
+        max_offx = w - nw
+        max_offy = h - nh
+        
+        dx = max_offx//2
+        dy = max_offy//2
+
+        to_x0, to_y0 = max(0, dx),    max(0, dy)
+        from_x0, from_y0 = max(0, -dx), max(0, -dy)
+        wx, hy = min(w, dx+nw) - to_x0, min(h, dy+nh) - to_y0
+
+        # place image
+        image_data = np.zeros((*input_shape,ic), dtype='uint8') + 128
+        image_data[to_y0:to_y0+hy, to_x0:to_x0+wx, :] = cv2.resize(image, (nw, nh))[from_y0:from_y0+hy, from_x0:from_x0+wx, :]
+        
+        flip = False
+        image_data = image_data/255.
+    else:
+        if np.random.uniform() >= 0.5:
+            # scale Up
+            resize_scale = 1. + scale * np.random.uniform()
+            resize_scale = max( h*resize_scale/ih, w*resize_scale/iw)
+
+            nw = int(iw * resize_scale)
+            nh = int(ih * resize_scale)
+
+            max_offx = nw - w
+            max_offy = nh - h
+
+            dx = int(np.random.uniform() * max_offx)
+            dy = int(np.random.uniform() * max_offy)
+
+            # resize and crop
+            image = cv2.resize(image, (nw, nh))
+            image_data = image[dy : (dy + h), dx : (dx + w), :]
+
+            dx, dy = (-dx, -dy)
+        else:
+            # scale down
+            mul = 1 if np.random.uniform() >= 0.5 else -1
+
+            resize_scale = 1. + mul * scale * np.random.uniform()
+            resize_scale = min( h*resize_scale/ih, w*resize_scale/iw)
+
+            nw = int(iw * resize_scale)
+            nh = int(ih * resize_scale)
+
+            max_offx = w - nw
+            max_offy = h - nh
+
+            dx = int(np.random.uniform() * max_offx)
+            dy = int(np.random.uniform() * max_offy)
+
+            to_x0, to_y0 = max(0, dx),    max(0, dy)
+            from_x0, from_y0 = max(0, -dx), max(0, -dy)
+            wx, hy = min(w, dx+nw) - to_x0, min(h, dy+nh) - to_y0
+
+            # place image
+            image_data = np.zeros((*input_shape,ic), dtype='uint8') + 128
+            image_data[to_y0:to_y0+hy, to_x0:to_x0+wx, :] = cv2.resize(image, (nw, nh))[from_y0:from_y0+hy, from_x0:from_x0+wx, :]
+    
+        flip = np.random.uniform() >= 0.5
+        if flip: image_data = image_data[:,::-1,:]
+
+        # distort color of the image
+        hue = rand(-hue, hue)
+        sat = rand(1, sat) if rand()<.5 else 1/rand(1, sat)
+        val = rand(1, val) if rand()<.5 else 1/rand(1, val)
+        x = rgb_to_hsv(np.array(image_data)/255.)
+        x[..., 0] += hue
+        x[..., 0][x[..., 0]>1] -= 1
+        x[..., 0][x[..., 0]<0] += 1
+        x[..., 1] *= sat
+        x[..., 2] *= val
+        x[x>1] = 1
+        x[x<0] = 0
+        image_data = hsv_to_rgb(x) # numpy array, 0 to 1
+
+    # correct boxes
+    box_data = np.zeros((max_boxes,5))
+    if len(box)>0:
+        np.random.shuffle(box)
+        box[:, [0,2]] = box[:, [0,2]]*nw/iw + dx
+        box[:, [1,3]] = box[:, [1,3]]*nh/ih + dy
+        if flip: box[:, [0,2]] = w - box[:, [2,0]]
+        box[:, 0:2][box[:, 0:2]<0] = 0
+        box[:, 2][box[:, 2]>w] = w
+        box[:, 3][box[:, 3]>h] = h
+        box_w = box[:, 2] - box[:, 0]
+        box_h = box[:, 3] - box[:, 1]
+        box = box[np.logical_and(box_w>1, box_h>1)] # discard invalid box
+        if len(box)>max_boxes: box = box[:max_boxes]
+        box_data[:len(box)] = box
+
+    return image_data, box_data
+
+
+def data_generator(annotation_lines, batch_size, input_shape, random):
+    '''
+    data generator for fit_generator
+    '''
+    n = len(annotation_lines)
+    i = 0
+    while True:
+        image_data = []
+        box_data = []
+        for _ in range(batch_size):
+            image, box = get_random_data(annotation_lines[i], input_shape, max_boxes=50, random=random)
+            image_data.append(image)
+            box = box[np.sum(box, axis=1) != 0, :]
+            box_data.append(box)
+            i = (i+1) % n
+        image_data = np.array(image_data)
+        box_data = np.array(box_data)
+        
+        yield image_data, box_data
+        
+        
+def data_generator_wrapper(annotation_lines, batch_size, input_shape, random):
+    n = len(annotation_lines)
+    if n==0 or batch_size<=0: return None
+    return data_generator(annotation_lines, batch_size, input_shape, random)

src/utils/fixes.py

@@ -0,0 +1,19 @@
+import tensorflow as tf
+
+
+def fix_tf_gpu():
+    '''
+    Fix for the following error message:
+    UnknownError: Failed to get convolution algorithm. 
+    This is probably because cuDNN failed to initialize...
+
+    More:
+    https://www.tensorflow.org/api_docs/python/tf/config/experimental/set_memory_growth
+    '''
+
+    physical_devices = tf.config.experimental.list_physical_devices('GPU')
+    
+    try:
+        tf.config.experimental.set_memory_growth(physical_devices[0], True)
+    except:
+        pass

src/utils/image.py

@@ -0,0 +1,82 @@
+import cv2
+import numpy as np
+import matplotlib as mpl
+
+
+def letterbox_image(image, size):
+    '''
+    Resize image with unchanged aspect ratio using padding
+    '''
+
+    # original image size
+    ih, iw, ic = image.shape
+
+    # given size
+    h, w = size
+
+    # scale and new size of the image
+    scale = min(w/iw, h/ih)
+    nw = int(iw*scale)
+    nh = int(ih*scale)
+    
+    # placeholder letter box
+    new_image = np.zeros((h, w, ic), dtype='uint8') + 128
+
+    # top-left corner
+    top, left = (h - nh)//2, (w - nw)//2
+
+    # paste the scaled image in the placeholder anchoring at the top-left corner
+    new_image[top:top+nh, left:left+nw, :] = cv2.resize(image, (nw, nh))
+    
+    return new_image
+
+
+def draw_detection(
+    img,
+    boxes,
+    class_names,
+    # drawing configs
+    font=cv2.FONT_HERSHEY_DUPLEX,
+    font_scale=0.5,
+    box_thickness=2,
+    border=5,
+    text_color=(255, 255, 255),
+    text_weight=1
+):
+    '''
+    Draw the bounding boxes on the image
+    '''
+    # generate some colors for different classes
+    num_classes = len(class_names) # number of classes
+    colors = [mpl.colors.hsv_to_rgb((i/num_classes, 1, 1)) * 255 for i in range(num_classes)]
+    
+    # draw the detections
+    for box in boxes:
+        x1, y1, x2, y2 = box[:4].astype(int)
+        score = box[-2]
+        label = int(box[-1])
+
+        clr = colors[label]
+
+        # draw the bounding box
+        img = cv2.rectangle(img, (x1, y1), (x2, y2), clr, box_thickness)
+
+        # text: <object class> (<confidence score in percent>%)
+        text = f'{class_names[label]} ({score*100:.0f}%)'
+
+        # get width (tw) and height (th) of the text
+        (tw, th), _ = cv2.getTextSize(text, font, font_scale, 1)
+
+        # background rectangle for the text
+        tb_x1 = x1 - box_thickness//2
+        tb_y1 = y1 - box_thickness//2 - th - 2*border
+        tb_x2 = x1 + tw + 2*border
+        tb_y2 = y1
+
+        # draw the background rectangle
+        img = cv2.rectangle(img, (tb_x1, tb_y1), (tb_x2, tb_y2), clr, -1)
+
+        # put the text
+        img = cv2.putText(img, text, (x1 + border, y1 - border), font, font_scale, text_color, text_weight, cv2.LINE_AA)
+
+    return img

src/yolo3/detect.py

@@ -0,0 +1,170 @@
+import numpy as np
+import tensorflow as tf
+
+
+def detection(
+    prediction,
+    anchor_boxes,
+    num_classes,
+    image_shape,
+    input_shape,
+    max_boxes = 20,
+    score_threshold=0.3,
+    iou_threshold=0.45,
+    classes_can_overlap=True,
+):
+    '''
+    INPUT:
+    OUTPUT:
+    '''
+    
+    all_boxes  = []
+
+    '''@ Each output layer'''
+    for output, anchors in zip( prediction, anchor_boxes ):
+
+        '''Preprocessing'''
+        '''-------------'''
+        # shapes
+        batch_size     = output.shape[0]
+        grid_h, grid_w = output.shape[1:3]
+
+        # reshape to [batch_size, grid_height, grid_width, num_anchors, box_params]
+        output = tf.reshape( output, [ -1, grid_h, grid_w, len(anchors), num_classes+5 ] )
+
+        # create a tensor for the anchor boxes
+        anchors_tensor = tf.constant(anchors, dtype=output.dtype)
+
+        '''Scaling factors'''
+        '''---------------'''
+        image_shape_tensor = tf.cast( image_shape,       output.dtype ) # actual image's shape
+        grids_shape_tensor = tf.cast( output.shape[1:3], output.dtype ) # grid_height, grid_width @ output layer
+        input_shape_tensor = tf.cast( input_shape,       output.dtype )  # yolo input image's shape
+
+        # reshape
+        image_shape_tensor = tf.reshape( image_shape_tensor, [-1, 1, 1, 1, 2] )
+        grids_shape_tensor = tf.reshape( grids_shape_tensor, [-1, 1, 1, 1, 2] )
+        input_shape_tensor = tf.reshape( input_shape_tensor, [-1, 1, 1, 1, 2] )
+
+        ### Scaling factors
+        sized_shape_tensor = tf.round( image_shape_tensor * tf.reshape( tf.reduce_min( input_shape_tensor / image_shape_tensor, axis=-1 ), [-1,1,1,1,1] ) )
+        # to scale the boxes from grid's unit to actual image's pixel unit
+        box_scaling = input_shape_tensor * image_shape_tensor / sized_shape_tensor / grids_shape_tensor
+        # to offset the boxes
+        box_offsets = (tf.expand_dims(tf.reduce_max(image_shape_tensor, axis=-1), axis=-1) - image_shape_tensor) / 2.
+
+        '''Box geometric properties'''
+        '''------------------------'''
+        grid_h, grid_w = output.shape[1:3] # grid_height, grid_width @ output layer
+
+        grid_i = tf.reshape( np.arange(grid_h), [-1, 1, 1, 1] )
+        grid_i = tf.tile( grid_i, [1, grid_w, 1, 1] )
+
+        grid_j = tf.reshape( np.arange(grid_w), [1, -1, 1, 1] )
+        grid_j = tf.tile( grid_j, [grid_h, 1, 1, 1] )
+
+        grid_ji = tf.concat( [grid_j, grid_i], axis=-1 )
+        grid_ji = tf.cast( grid_ji, output.dtype )
+
+        # Box centers
+        box_xy  = output[..., 0:2]
+        box_xy  = tf.sigmoid( box_xy ) + grid_ji
+
+        # Box sizes
+        box_wh  = output[..., 2:4]
+        box_wh  = tf.exp( box_wh ) * anchors_tensor
+
+        # scale to actual pixel unit
+        box_xy  = box_xy * box_scaling - box_offsets[...,::-1]
+        box_wh  = box_wh * box_scaling
+
+        # calculate top-left corner (x1, y1) and bottom-right corner (x2, y2) of the boxex
+        box_x1_y1 = box_xy - box_wh / 2
+        box_x2_y2 = box_xy + box_wh / 2
+
+        # top-left corner cannot be negative
+        box_x1_y1 = tf.maximum(0, box_x1_y1)
+        # bottom-right corner cannot be more than actual image size
+        box_x2_y2 = tf.minimum(box_x2_y2, image_shape_tensor[..., ::-1])
+
+        '''Box labels and confidences'''
+        '''--------------------------'''
+        # class probabilities = objectness score * conditional class probabilities
+        if classes_can_overlap:
+            # use sigmoid for the conditional class probabilities
+            classs_probs = tf.sigmoid( output[..., 4:5] ) * tf.sigmoid( output[..., 5:] )
+        else:
+            # use softmax for the conditional class probabilities
+            classs_probs = tf.sigmoid( output[..., 4:5] ) * tf.nn.softmax( output[..., 5:] )
+
+        box_cl = tf.argmax( classs_probs, axis=-1 )     # final classes
+        box_sc = tf.reduce_max( classs_probs, axis=-1 ) # confidence scores
+
+        '''Organize'''
+        '''--------'''
+        # take care of dtype and dimensions
+        box_cl = tf.cast( box_cl, output.dtype )
+        box_cl = tf.expand_dims(box_cl, axis=-1)
+        box_sc = tf.expand_dims(box_sc, axis=-1)
+
+        # store all information as: [ left(x1), top(y1), right(x2), bottom(y2),  confidence, label ]
+        boxes  = tf.reshape( tf.concat( [ box_x1_y1, box_x2_y2, box_sc, box_cl ], axis=-1 ), 
+                              [batch_size, -1, 6] )
+
+        all_boxes. append( boxes  )
+
+    # Merge across all output layers
+    all_boxes  = tf.concat( all_boxes,  axis=1 )
+
+    # To store all the final results of all images in the batch
+    all_final_boxes = []
+
+    '''For each image in the batch'''
+    for _boxes_ in all_boxes:
+
+        if classes_can_overlap:
+            '''Perform NMS for each class individually'''
+
+            # to stote the final results of this image
+            final_boxes = []
+
+            for class_id in range(num_classes):
+
+                # Get the boxes and scores for this class
+                class_boxes  = _boxes_[ _boxes_[...,-1] == class_id ]
+
+                '''Non-max-suppression'''
+                selected_idc = tf.image.non_max_suppression(
+                    class_boxes[...,:4], # boxes' (y1,x1,y2,x2)
+                    class_boxes[...,-2], # boxes' scores
+                    max_output_size = max_boxes,
+                    iou_threshold = iou_threshold,
+                    score_threshold = score_threshold
+                )
+
+                # boxes selected by nms
+                class_boxes = tf.gather( class_boxes,  selected_idc )
+                final_boxes.append( class_boxes )
+
+            # concatenate boxes for each class in the image
+            final_boxes  = tf.concat( final_boxes,  axis=0 )
+
+        else:
+            '''Perform NMS for all classes'''
+
+            # nms indices
+            selected_idc = tf.image.non_max_suppression(
+                _boxes_[...,:4], # boxes' (y1,x1,y2,x2)
+                _boxes_[...,-2], # boxes' scores
+                max_output_size = max_boxes,
+                iou_threshold = iou_threshold,
+                score_threshold = score_threshold
+            )
+            
+            # boxes selected by nms
+            final_boxes = tf.gather( _boxes_,  selected_idc )
+
+        # append final boxes for each image in the batch
+        all_final_boxes.append( final_boxes )
+        
+    return all_final_boxes

src/yolo3/model.py

@@ -0,0 +1,95 @@
+import tensorflow as tf
+import tensorflow.keras.backend as K
+
+from tensorflow.keras.layers import Input, Conv2D, Add, ZeroPadding2D, UpSampling2D, Concatenate, MaxPooling2D
+from tensorflow.keras.layers import LeakyReLU, BatchNormalization
+from tensorflow.keras.models import Sequential, Model
+from tensorflow.keras.regularizers import l2
+
+'''===================================================================================================='''
+'''BLOCKS'''
+
+'''Convolutional Block'''
+def yolo_ConvBlock (input_tensor, num_filters, filter_size, strides = (1,1) ):
+    padding = 'valid' if strides == (2,2) else 'same'
+
+    ### Layers
+    x = Conv2D( num_filters, filter_size, strides, padding, use_bias=False, kernel_regularizer=l2(5e-4) ) (input_tensor)
+    x = BatchNormalization() (x)
+    x = LeakyReLU(alpha=0.1) (x)
+    
+    return x
+
+'''Residual Block'''
+def yolo_ResidualBlocks (input_tensor, num_filters, num_blocks ):
+    
+    ### Layers
+    x = ZeroPadding2D( ((1,0),(1,0)) ) (input_tensor) # left & top padding
+    x = yolo_ConvBlock ( x, num_filters, filter_size=(3,3), strides = (2,2) )
+    
+    for _ in range( num_blocks ):
+        y = yolo_ConvBlock ( x, num_filters//2, filter_size=(1,1), strides = (1,1) )
+        y = yolo_ConvBlock ( y, num_filters   , filter_size=(3,3), strides = (1,1) )
+        x = Add() ([x, y])
+    
+    return x
+
+'''Output Block'''
+def yolo_OutputBlock (x, num_filters, out_filters ):
+    
+    ### Layers
+    x = yolo_ConvBlock ( x, 1*num_filters, filter_size=(1,1), strides = (1,1) )
+    x = yolo_ConvBlock ( x, 2*num_filters, filter_size=(3,3), strides = (1,1) )
+    x = yolo_ConvBlock ( x, 1*num_filters, filter_size=(1,1), strides = (1,1) )
+    x = yolo_ConvBlock ( x, 2*num_filters, filter_size=(3,3), strides = (1,1) )
+    x = yolo_ConvBlock ( x, 1*num_filters, filter_size=(1,1), strides = (1,1) )
+    
+    y = yolo_ConvBlock ( x, 2*num_filters, filter_size=(3,3), strides = (1,1) )
+    y = Conv2D ( filters=out_filters, kernel_size=(1,1), strides=(1,1), 
+                padding='same', use_bias=True, kernel_regularizer=l2(5e-4) ) (y)
+    
+    return x, y
+
+'''===================================================================================================='''
+'''COMPLETE MODEL'''
+
+def yolo_body (input_tensor, num_out_filters):
+    '''
+    Input: 
+        input_tensor   = Input( shape=( *input_shape, 3 ) )
+        num_out_filter = ( num_anchors // 3 ) * ( 5 + num_classes )
+    Output:
+        complete YOLO-v3 model
+    '''
+
+    # 1st Conv block
+    x = yolo_ConvBlock( input_tensor, num_filters=32, filter_size=(3,3), strides=(1,1) )
+
+    # 5 Resblocks
+    x = yolo_ResidualBlocks ( x, num_filters=  64, num_blocks=1 )
+    x = yolo_ResidualBlocks ( x, num_filters= 128, num_blocks=2 )
+    x = yolo_ResidualBlocks ( x, num_filters= 256, num_blocks=8 )
+    x = yolo_ResidualBlocks ( x, num_filters= 512, num_blocks=8 )
+    x = yolo_ResidualBlocks ( x, num_filters=1024, num_blocks=4 )
+
+    darknet = Model( input_tensor, x ) # will use it just in a moment
+
+    # 1st output block
+    x, y1 = yolo_OutputBlock( x, num_filters= 512, out_filters=num_out_filters )
+
+    # 2nd output block
+    x = yolo_ConvBlock( x, num_filters=256, filter_size=(1,1), strides=(1,1) )
+    x = UpSampling2D(2) (x)
+    x = Concatenate() ( [x, darknet.layers[152].output] )
+    x, y2 = yolo_OutputBlock( x, num_filters= 256, out_filters=num_out_filters )
+
+    # 3rd output block
+    x = yolo_ConvBlock( x, num_filters=128, filter_size=(1,1), strides=(1,1) )
+    x = UpSampling2D(2) (x)
+    x = Concatenate() ( [x, darknet.layers[92].output] )
+    x, y3 = yolo_OutputBlock( x, num_filters= 128, out_filters=num_out_filters )
+
+    # Final model
+    model = Model( input_tensor, [y1, y2, y3] )
+    
+    return model