Spaces:
Sleeping
Sleeping
import gradio as gr | |
import numpy as np | |
import os | |
os.environ['CUDA_VISIBLE_DEVICES'] = '-1' | |
from glob import glob | |
import tensorflow as tf | |
import matplotlib.pyplot as plt | |
from skimage.transform import resize | |
from skimage.io import imsave | |
from skimage.filters import threshold_otsu | |
from doodleverse_utils.prediction_imports import * | |
from doodleverse_utils.imports import * | |
#load model | |
filepath = './saved_model' | |
model = tf.keras.models.load_model(filepath, compile = True) | |
model.compile | |
#segmentation | |
def segment(input_img, use_tta, use_otsu, dims=(512, 512)): | |
N = 2 | |
if use_otsu: | |
print("Use Otsu threshold") | |
else: | |
print("No Otsu threshold") | |
if use_tta: | |
print("Use TTA") | |
else: | |
print("Do not use TTA") | |
worig, horig, channels = input_img.shape | |
w, h = dims[0], dims[1] | |
print("Original dimensions {}x{}".format(worig,horig)) | |
print("New dimensions {}x{}".format(w,h)) | |
img = standardize(input_img) | |
img = resize(img, dims, preserve_range=True, clip=True) | |
img = np.expand_dims(img,axis=0) | |
est_label = model.predict(img) | |
if use_tta: | |
#Test Time Augmentation | |
est_label2 = np.flipud(model.predict((np.flipud(img)), batch_size=1)) | |
est_label3 = np.fliplr(model.predict((np.fliplr(img)), batch_size=1)) | |
est_label4 = np.flipud(np.fliplr(model.predict((np.flipud(np.fliplr(img)))))) | |
#soft voting - sum the softmax scores to return the new TTA estimated softmax scores | |
est_label = est_label + est_label2 + est_label3 + est_label4 | |
est_label /= 4 | |
pred = np.squeeze(est_label, axis=0) | |
pred = resize(pred, (worig, horig), preserve_range=True, clip=True) | |
mask = np.argmax(pred,-1) | |
imsave("greyscale_download_me.png", mask.astype('uint8')) | |
class_label_colormap = [ | |
"#3366CC", | |
"#DC3912", | |
"#FF9900", | |
"#109618", | |
"#990099", | |
"#0099C6", | |
"#DD4477", | |
"#66AA00", | |
"#B82E2E", | |
"#316395", | |
] | |
# add classes | |
class_label_colormap = class_label_colormap[:N] | |
color_label = label_to_colors( | |
mask, | |
input_img[:, :, 0] == 0, | |
alpha=128, | |
colormap=class_label_colormap, | |
color_class_offset=0, | |
do_alpha=False, | |
) | |
imsave("color_download_me.png", color_label) | |
if use_otsu: | |
thres = threshold_otsu(mask) | |
print("Otsu threshold is {}".format(thres)) | |
water_nowater = (mask>thres).astype('uint8') | |
else: | |
water_nowater = (mask>=1).astype('uint8') | |
#overlay plot | |
plt.clf() | |
plt.subplot(121) | |
plt.imshow(input_img[:,:,-1],cmap='gray') | |
plt.imshow(color_label, alpha=0.4) | |
plt.axis("off") | |
plt.margins(x=0, y=0) | |
plt.subplot(122) | |
plt.imshow(input_img[:,:,-1],cmap='gray') | |
plt.contour(water_nowater, levels=[0], colors='r') | |
plt.axis("off") | |
plt.margins(x=0, y=0) | |
plt.savefig("overlay_download_me.png", dpi=300, bbox_inches="tight") | |
return color_label, plt , "greyscale_download_me.png", "color_download_me.png", "overlay_download_me.png" | |
with open("article.html", "r", encoding='utf-8') as f: | |
article= f.read() | |
title = "Segment Satellite imagery" | |
description = "This simple model demonstration segments 15-m Landsat-7/8 or 10-m Sentinel-2 RGB (visible spectrum) imagery into the following classes: 1. water and 2. other. Please note that, ordinarily, ensemble models are used in predictive mode. Here, we are using just one model, i.e. without ensembling. Allows upload of 3-band imagery in jpg format and download of label imagery only one at a time. " | |
examples= [[l] for l in glob('examples/*.jpg')] | |
inp = gr.Image() | |
out1 = gr.Image(type='numpy') | |
out2 = gr.Plot(type='matplotlib') | |
out3 = gr.File() | |
out4 = gr.File() | |
out5 = gr.File() | |
inp2 = gr.inputs.Checkbox(default=False, label="Use TTA") | |
inp3 = gr.inputs.Checkbox(default=False, label="Use Otsu") | |
Segapp = gr.Interface(segment, [inp, inp2, inp3], | |
[out1, out2, out3, out4, out5], | |
title = title, description = description, examples=examples, article=article, | |
theme="grass") | |
Segapp.launch(enable_queue=True) | |