-
Notifications
You must be signed in to change notification settings - Fork 0
/
predict.py
88 lines (69 loc) · 2.8 KB
/
predict.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
import numpy as np
import pandas as pd
import tensorflow as tf
from tensorflow.keras import layers
from PIL import Image
from tqdm import tqdm
import matplotlib.pyplot as plt
import os, glob
target_size = [256, 256] # size of image for model
shuffle_buffer = 32
batch_size = 4 # batch size
MODEL_DIR = './' #dir where saving models or weights
train_dir = 'train_v2' # folder with train image
train_mask_dir = 'train_masks' # folder with masks
test_dir = 'test_v2' # folder with test image
model = tf.keras.models.load_model(os.path.join(MODEL_DIR, 'airbus_model.h5'))
def get_predict_mask(img, threshold = 0.5, pred_size = 256):
"""
Predict mask
img - input image (array)
threshold - threshold for classes
pred_size - size of window for cut
Return mask (array)
"""
# how many parts cut the picture, it will be necessary to create a mask
mask_parts = [np.ceil(img.shape[0]/pred_size).astype(int), np.ceil(img.shape[1]/pred_size).astype(int)]
#mask with a multiple of pred_size sides
full_mask = np.zeros((mask_parts[0]*pred_size, mask_parts[1]*pred_size))
patch = np.zeros((pred_size, pred_size, 3))
stack_parts = []
# Pass through the analyzed image, pulling out sections of pred_size in size
for i in range(0, img.shape[0], pred_size):
for j in range(0, img.shape[1], pred_size):
patch_img = img[i:i+pred_size, j:j+pred_size]
patch = np.zeros((pred_size, pred_size, 3))
patch[:patch_img.shape[0], :patch_img.shape[1]] = patch_img
stack_parts.append(patch / 127.5 - 1)
# predict mask
pred = model.predict(np.array(stack_parts))
#Use threshold
pred[pred>=threshold] = 1
pred[pred<threshold] = 0
for k in range(pred.shape[0]):
i = k//mask_parts[0]*pred_size
j = k%mask_parts[1]*pred_size
full_mask[i:i+pred_size, j:j+pred_size] = np.squeeze(pred[k])
# Crop the mask to fit the image
full_mask = full_mask[:img.shape[0], :img.shape[1]]
return full_mask
def decode_mask(mask, shape=(768, 768)):
#from image(mask) to run-length encoding
pixels = mask.T.flatten()
pixels = np.concatenate([[0], pixels, [0]])
runs = np.where(pixels[1:] != pixels[:-1])[0] + 1
runs[1::2] -= runs[::2]
return ' '.join(str(x) for x in runs)
#predict all test images
test_data = os.listdir(test_dir)
ship_list_dict = []
for name in tqdm(test_data):
img = plt.imread(os.path.join(test_dir, name))
predict_mask = get_predict_mask(img,0.5)
decode = decode_mask(predict_mask)
ship_list_dict.append({'ImageId':name,'EncodedPixels':decode})
pred_df = pd.DataFrame(ship_list_dict)
pred_df.to_csv('submission.csv', index=False)
pred_df
#load submission to kaggle
#!kaggle competitions submit -c airbus-ship-detection -f submission.csv -m "20 epoch, treshold=0.5"