├── requirements.txt ├── images ├── scene_1.png ├── scene_2.png ├── scene_3.png └── scene_4.png ├── train.py ├── model.py ├── detector.py └── README.md /requirements.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/souravp-d/ShipDetection/HEAD/requirements.txt -------------------------------------------------------------------------------- /images/scene_1.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/souravp-d/ShipDetection/HEAD/images/scene_1.png -------------------------------------------------------------------------------- /images/scene_2.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/souravp-d/ShipDetection/HEAD/images/scene_2.png -------------------------------------------------------------------------------- /images/scene_3.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/souravp-d/ShipDetection/HEAD/images/scene_3.png -------------------------------------------------------------------------------- /images/scene_4.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/souravp-d/ShipDetection/HEAD/images/scene_4.png -------------------------------------------------------------------------------- /train.py: -------------------------------------------------------------------------------- 1 | """ 2 | Train and export machine learning model using ShipsNet dataset 3 | """ 4 | 5 | import sys 6 | import json 7 | import numpy as np 8 | from tflearn.data_utils import to_categorical 9 | from model import model 10 | 11 | def train(fname, out_fname): 12 | """ Train and save CNN model on ShipsNet dataset 13 | 14 | Args: 15 | fname (str): Path to ShipsNet JSON dataset 16 | out_fname (str): Path to output Tensorflow model file (.tfl) 17 | """ 18 | 19 | # Load shipsnet data 20 | f = open(fname) 21 | shipsnet = json.load(f) 22 | f.close() 23 | 24 | # Preprocess image data and labels for input 25 | X = np.array(shipsnet['data']) / 255. 26 | X = X.reshape([-1,3,80,80]).transpose([0,2,3,1]) 27 | Y = np.array(shipsnet['labels']) 28 | Y = to_categorical(Y, 2) 29 | 30 | # Train the model 31 | model.fit(X, Y, n_epoch=50, shuffle=True, validation_set=.2, 32 | show_metric=True, batch_size=128, run_id='shipsnet') 33 | 34 | # Save trained model 35 | model.save(out_fname) 36 | 37 | 38 | # Main function 39 | if __name__ == "__main__": 40 | 41 | # Train using input file 42 | train(sys.argv[1], sys.argv[2]) 43 | -------------------------------------------------------------------------------- /model.py: -------------------------------------------------------------------------------- 1 | """ 2 | Define Convolutional Nerual Network model for ShipsNet input 3 | """ 4 | 5 | import tflearn 6 | from tflearn.data_utils import shuffle, to_categorical 7 | from tflearn.layers.core import input_data, dropout, fully_connected 8 | from tflearn.layers.conv import conv_2d, max_pool_2d 9 | from tflearn.layers.estimator import regression 10 | from tflearn.data_preprocessing import ImagePreprocessing 11 | from tflearn.data_augmentation import ImageAugmentation 12 | 13 | # Real-time data preprocessing 14 | img_prep = ImagePreprocessing() 15 | img_prep.add_featurewise_zero_center() 16 | img_prep.add_featurewise_stdnorm() 17 | 18 | # Real-time data augmentation 19 | img_aug = ImageAugmentation() 20 | img_aug.add_random_flip_leftright() 21 | img_aug.add_random_flip_updown() 22 | img_aug.add_random_rotation(max_angle=25.) 23 | 24 | # Convolutional network building 25 | network = input_data(shape=[None, 80, 80, 3], 26 | data_preprocessing=img_prep, 27 | data_augmentation=img_aug) 28 | network = conv_2d(network, 32, 3, activation='relu') 29 | network = max_pool_2d(network, 2) 30 | network = conv_2d(network, 64, 3, activation='relu') 31 | network = conv_2d(network, 64, 3, activation='relu') 32 | network = max_pool_2d(network, 2) 33 | network = fully_connected(network, 512, activation='relu') 34 | network = dropout(network, 0.5) 35 | network = fully_connected(network, 2, activation='softmax') 36 | network = regression(network, optimizer='adam', 37 | loss='categorical_crossentropy', 38 | learning_rate=0.001) 39 | 40 | # Define model 41 | model = tflearn.DNN(network, tensorboard_verbose=0) 42 | -------------------------------------------------------------------------------- /detector.py: -------------------------------------------------------------------------------- 1 | ''' 2 | Apply trained machine learning model to an entire image scene using 3 | a sliding window. 4 | ''' 5 | 6 | import sys 7 | import os 8 | import numpy as np 9 | from PIL import Image 10 | from scipy import ndimage 11 | from model import model 12 | 13 | def detector(model_fname, in_fname, out_fname=None): 14 | """ Perform a sliding window detector on an image. 15 | 16 | Args: 17 | model_fname (str): Path to Tensorflow model file (.tfl) 18 | in_fname (str): Path to input image file 19 | out_fname (str): Path to output image file. Default of None. 20 | 21 | """ 22 | 23 | # Load trained model 24 | model.load(model_fname) 25 | 26 | # Read input image data 27 | im = Image.open(in_fname) 28 | arr = np.array(im)[:,:,0:3] 29 | shape = arr.shape 30 | 31 | # Set output fname 32 | if not out_fname: 33 | out_fname = os.path.splitext(in_fname)[0] + '_detection.png' 34 | 35 | # Create detection variables 36 | detections = np.zeros((shape[0], shape[1]), dtype='uint8') 37 | output = np.copy(arr) 38 | 39 | # Sliding window parameters 40 | step = 2 41 | win = 80 42 | 43 | # Loop through pixel positions 44 | print('Processing...') 45 | for i in range(0, shape[0]-win, step): 46 | print('row %1.0f of %1.0f' % (i, (shape[0]-win-1))) 47 | 48 | for j in range(0, shape[1]-win, step): 49 | 50 | # Extract sub chip 51 | chip = arr[i:i+win,j:j+win,:] 52 | 53 | # Predict chip label 54 | prediction = model.predict_label([chip / 255.])[0][0] 55 | 56 | # Record positive detections 57 | if prediction == 1: 58 | detections[i+int(win/2), j+int(win/2)] = 1 59 | 60 | # Process detection locations 61 | dilation = ndimage.binary_dilation(detections, structure=np.ones((3,3))) 62 | labels, n_labels = ndimage.label(dilation) 63 | center_mass = ndimage.center_of_mass(dilation, labels, np.arange(n_labels)+1) 64 | 65 | # Loop through detection locations 66 | if type(center_mass) == tuple: center_mass = [center_mass] 67 | for i, j in center_mass: 68 | i = int(i - win/2) 69 | j = int(j - win/2) 70 | 71 | # Draw bouding boxes in output array 72 | output[i:i+win, j:j+2, 0:3] = [255,0,0] 73 | output[i:i+win, j+win-2:j+win, 0:3] = [255,0,0] 74 | output[i:i+2, j:j+win, 0:3] = [255,0,0] 75 | output[i+win-2:i+win, j:j+win, 0:3] = [255,0,0] 76 | 77 | # Save output image 78 | outIm = Image.fromarray(output) 79 | outIm.save(out_fname) 80 | 81 | 82 | # Main function 83 | if __name__ == "__main__": 84 | 85 | # Run detection function with command line inputs 86 | if len(sys.argv) == 3: 87 | detector(sys.argv[1], sys.argv[2]) 88 | else: 89 | detector(sys.argv[1], sys.argv[2], sys.argv[3]) -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # shipsnet-detector 2 | This repository contains scripts that enable the automatic detection of container ships in [Planet](https://www.planet.com/) imagery using machine learning techniques. Included are files which define a machine learning model, train it using the ShipsNet dataset, and apply it across an entire image scene to highlight ship detections. 3 | 4 | ## Methodology 5 | [ShipsNet](https://www.kaggle.com/rhammell/ships-in-satellite-imagery) is a labeled training dataset consiting of image chips extracted from Planet satellite imagery. It contains hundreds of 80x80 pixel RGB image chips labeled with either a "ship" or "no-ship" classification. Machine learning models can be trained against this data to classify any given input chip into either one of these classes. 6 | 7 | With an accurately trained model, this classification process can be extended to a full Planet image scene by using a sliding window technique. A 80x80 pixel window is moved across each pixel position in the image, extracted, and classified by the model. Neighboring window poistions that are classified as "ship" are then clustered into a single detection. These detections are highlighted with a bounding box in a copy of the original Planet scene. 8 | 9 | See an example of the results below. 10 |

11 | 12 | 13 |

14 | 15 | [Additional Results](https://imgur.com/a/fvo7g) 16 | 17 | ## Setup 18 | Python 3.5+ is required for compatability with all required modules 19 | 20 | ```bash 21 | # Clone this repository 22 | git clone https://github.com/rhammell/shipsnet-detector.git 23 | 24 | # Go into the repository 25 | cd shipsnet-detector 26 | 27 | # Install required modules 28 | pip install -r requirements.txt 29 | ``` 30 | 31 | ## Model 32 | A convolutional neural network (CNN) is defined within the `model.py` module using the [TFLearn](http://tflearn.org/) library. This model supports the 80x80x3 input dimensions of the ShipsNet image data. 33 | 34 | ## Training 35 | The defined CNN can be trained with the JSON version of the ShipsNet dataset and saved to a Tensorflow .tfl file for later use. Train the model by running `train.py` and passing the path to `shipsnet.json` and the path to the output .tfl file as arguments. 36 | 37 | ```bash 38 | # Train the model 39 | mkdir models 40 | python train.py "shipsnet.json" "models/model.tfl" 41 | ``` 42 | 43 | The latest version of `shipsnet.json` is available through the [ShipsNet](https://www.kaggle.com/rhammell/ships-in-satellite-imagery) Kaggle page, which has further information describing the dataset layout. 44 | 45 | ## Detector 46 | A trained model can be applied across entire images using the sliding window detector function `detector.py`, which takes the model file path, input image path, and optional output image path as arguments. The output image will cluster positive detections and draw a bounding box around their center point. 47 | 48 | Example images are contained in the `images` directory. 49 | ```bash 50 | # Run on demo image with default output path 51 | python detector.py "models/model.tfl" "images/scene_1.png" 52 | 53 | # Run on demo image with defined output path 54 | python detector.py "models/model.tfl" "image/scene_1.png" "image/scene_1_detections.png" 55 | ``` 56 | --------------------------------------------------------------------------------