def save_data(data):
try:
(img, time, data) = data
file_name = time.strftime("%H_%M_%S")
write_img(img, DAY_DIR / (file_name + ".jpg"))
(processed_img, robot) = data
processed_img: robot_finder.ProcessedImg = processed_img
write_img(processed_img.img, DAY_DIR / (file_name + "_proc.jpg"))
with (DAY_DIR / (file_name + ".json")).open('w') as fp:
json.dump(
{
"blue_pts": processed_img.blue_pts,
"red_pts": processed_img.red_pts,
"robot": robot,
},
fp,
indent=4,
cls=JSONEncoder,
)
except Exception as e:
logging.fatal("Error while saving data...", exc_info=e)
def predict(data_dir, output_dir, X, mode, pred):
if not os.path.exists(output_dir):
os.mkdir(output_dir)
for i in range(0, 15000, 5000):
for j in range(0, 15000, 5000):
one_data = '%s/tile_%d_%d.tif' % (data_dir, i, j)
one_data_output = '%s/pred_tile_%d_%d.tif' % (output_dir, i, j)
print ("predict tile_%d_%d.tif" % (i, j))
proj, geotrans, input_img = read_img(one_data)
label_pred_tiny = predict_img_with_smooth_windowing(
input_img,
window_size=768,
subdivisions=2,
batch_size=8,
pred_func=(
lambda img: sess.run(pred, feed_dict={X: img, mode: False})
)
)
label_pred_tiny = label_pred_tiny[:, :, 0]
# label_pred_tiny[np.where(label_pred_tiny >= 0.3)] = 255
label_pred_tiny[label_pred_tiny<0.1] = 0
label_pred_tiny = label_pred_tiny.astype(np.float32)
write_img(one_data_output, proj, geotrans, label_pred_tiny)
def post_process(prediction, class_idx, z_mask, out_path):
# print('post processing ...')
# 4. make predicted map and label map have the same index
# print('prediction', np.unique(np.ravel(prediction)))
# ind_mask = []
class_idx.reverse()
for i, num in enumerate(class_idx):
ind_mask = (prediction == len(class_idx)-i-1)
# print(i, np.sum(ind_mask[-1]))
prediction[ind_mask] = num
# print('prediction', np.unique(np.ravel(prediction)))
prediction[z_mask.astype(np.bool)] = 0
# print('prediction', np.unique(np.ravel(prediction)))
# 5. store prediction result
im_projection, im_geo_transformation = read_img(flags.image, only_coordinate=True)
write_img(filename=out_path, im_proj=im_projection, im_geotrans=im_geo_transformation, im_data=prediction)
# print('post process done ! now store the image')
# io.imsave(out_path, prediction)
class_idx.reverse()
return prediction
# make predicted map and label map have the same index
ind_mask = []
for i, num in enumerate(class_idx):
ind_mask.append(prediction == i)
# print(i, np.sum(ind_mask[-1]))
for i, num in enumerate(class_idx):
prediction[ind_mask[i]] = num
prediction[zero_mask] = 0
print(np.unique(np.ravel(prediction)))
# generate pseudo color
colors = np.random.randint(0,
255,
size=(len(class_idx) * 3, ),
dtype=np.uint8).reshape(-1, 3)
pseudo = np.zeros((prediction.shape[0], prediction.shape[1], 3),
dtype=np.uint8)
print(colors)
for i, idx in enumerate(class_idx):
pseudo[prediction == idx, :] = colors[i, :]
im_proj, im_geotrans = read_img(flags.image, only_coordinate=True)
write_img(filename=flags.output,
im_proj=im_proj,
im_geotrans=im_geotrans,
im_data=prediction)
# io.imsave(flags.output, prediction)
io.imsave('./pseudo.png', pseudo)
def predict(data_dir, output_dir, X, mode, pred):
print("reading image...")
proj, geotrans, input_img = read_img(data_dir)
# input_img[:, :, [2, 0]] = input_img[:, :, [0, 2]]
input_img = input_img[:, :, :3]
print(input_img.shape)
input_img = np.asarray(input_img, dtype='uint8')
step = 32
height, width, _ = input_img.shape
print(height // step)
print(width // step)
for i in range(height // step):
if i == height // step and height // step != 0:
continue
for j in range(width // step + 1):
if j == width // step and width // step != 0:
continue
last_height_value = (
i + 1) * step if i + 1 != height // step else height
last_width_value = (j +
1) * step if j + 1 != width // step else width
cv2.imwrite(
'%s/tmp_%d_%d.jpg' % (tmp_dir, i, j),
input_img[i * step:last_height_value,
j * step:last_width_value, :])
print('%s/tmp_%d_%d.jpg write done' % (tmp_dir, i, j))
del input_img
gc.collect()
print("done1")
label_pred = np.zeros((height, width), dtype=np.uint8)
for i in range(height // step):
if i == height // step and height // step != 0:
continue
for j in range(width // step + 1):
if j == width // step and width // step != 0:
continue
last_height_value = (
i + 1) * step if i + 1 != height // step else height
last_width_value = (j +
1) * step if j + 1 != width // step else width
print(last_height_value, last_width_value)
print('%s/tmp_%d_%d.jpg' % (tmp_dir, i, j))
img_tiny = cv2.imread('%s/tmp_%d_%d.jpg' % (tmp_dir, i, j))
print(img_tiny)
label_pred_tiny = predict_img_with_smooth_windowing(
img_tiny,
window_size=32,
subdivisions=2,
batch_size=64,
pred_func=(
lambda img: sess.run(pred, feed_dict={
X: img,
mode: False
})))
label_pred_tiny = label_pred_tiny[:, :, 0]
label_pred_tiny[np.where(label_pred_tiny >= 0.5)] = 1
label_pred_tiny[np.where(label_pred_tiny < 0.5)] = 0
label_pred_tiny = label_pred_tiny.astype(np.uint8)
label_pred[i * step:last_height_value,
j * step:last_width_value] = label_pred_tiny
gc.collect()
print("done2")
# cv2.imwrite(output_dir, label_pred)
write_img(output_dir, proj, geotrans, label_pred)
print("done3")
import tifffile as tif
from sklearn.neighbors import KNeighborsClassifier
import numpy as np
import cv2
from tools import read_img, write_img
if __name__ == "__main__":
pred_label_file = "../data/pine/pred_classes.tif"
test_label_file = "../data/pine/test_classes.tif"
pred_label_file_png = "../data/pine/pred_classes.png"
proj, geotrans, input_img = read_img(test_label_file)
pred = cv2.imread(pred_label_file)
write_img(pred_label_file, proj, geotrans, pred)