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 = 4 | |
| 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: | |
| c1 = pred[:,:,0] | |
| c2 = pred[:,:,1] | |
| water = c1+c2 | |
| water /= water.max() | |
| thres = threshold_otsu(water) | |
| print("Otsu threshold is {}".format(thres)) | |
| water_nowater = (water>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 (unbroken water); 2. whitewater (surf, active wave breaking); 3. sediment (natural deposits of sand. gravel, mud, etc), and 4. other (development, bare terrain, vegetated terrain, etc). 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) | |