Spaces:
Build error
Build error
| import torch | |
| import scipy.io as sio | |
| from argparse import Namespace | |
| import numpy as np | |
| import streamlit as st | |
| import matplotlib.pyplot as plt | |
| from PIL import Image, ImageEnhance | |
| import multispectral as mspec | |
| def GetRGB( snp, scale=2 ): | |
| try: | |
| snp = np.transpose( snp, [ 1, 2, 0 ] ).astype( float ) | |
| except: | |
| pass # not a multi-channel image | |
| snp = np.round( 255*( snp - snp.min() ) / ( snp.max() - snp.min() ) ).astype( np.uint8 ) | |
| im = Image.fromarray( snp ) | |
| im = im.resize( ( im.size[0]//scale, im.size[1]//scale ), resample=Image.Resampling.NEAREST ) | |
| if np.ndim( snp ) == 3: | |
| enhancer = ImageEnhance.Brightness( im ) | |
| im = enhancer.enhance( 5 ) | |
| return im | |
| def plot_zoomed_region( image, x1, y1, x2, y2, title='Zoomed in region' ): | |
| # Crop the image to the zoom box coordinates | |
| cropped_image = image.crop( ( x1, y1, x2, y2 ) ) | |
| # Convert cropped image to numpy array for plotting | |
| img_array = np.array( cropped_image ) | |
| # Plot the cropped image | |
| fig, ax = plt.subplots() | |
| ax.imshow( img_array ) | |
| ax.set_title( title ) | |
| ax.axis( 'off' ) | |
| return fig | |
| def dashboard_header(): | |
| cultural_reference = 'https://www.youtube.com/watch?v=ACmydtFDTGs' | |
| st.markdown( '# `hotdog-nothotdog v2.0`' ) | |
| st.markdown( '## - Jian Yang (possibly)' ) | |
| st.markdown( 'Cultural reference: [this clip](%s) from _Silicon Valley_. '%cultural_reference ) | |
| return | |
| def dashboard_intro(): | |
| st.markdown( | |
| ''' | |
| # TL;DR | |
| An interactive, trained-from-scratch semantic segmenter for 6-channel multi-spectral imagery data, trained solely on the RIT-18 dataset. | |
| This is a work in progress. | |
| <a href="#app">Go to the app. </a> | |
| # Dataset description | |
| RIT-18 consists of 19 classes within the scan range of a multi-spectral drone camera over a New York state park. It is highly unbalanced. | |
| # Training methodology | |
| This model was trained repeatedly on randomized $256 \\times 256$ pixel slices of the same swathe region of the drone, but with the pixel labels at each iteration increasingly diverse. | |
| For example, for the very first iteration, the dataset was relabeled from 19 classes to only "Vegetation" and "Other" (_i.e._, a binary segmenter). | |
| The resulting model became the starting point for the training of the next iteration, in which the pixels with the "Other" label were further re-labeled "Water" and "Other", for a total of 3 classes. | |
| Similarly, in later iterations, "Vegetation" will be broken down to "Trees" and "Grass" and so on. | |
| In this manner, the goal is to obtain a final model trained on all 19 classes of the RIT-18 dataset. | |
| # Goal | |
| To see if it's possible to incrementally train the model in this manner. | |
| If so, how influential are the near infrared (NIR) bands in the segmentation outcome? | |
| - **Dataset**: [RIT-18](https://arxiv.org/abs/1703.06452) | |
| - **Model architecture**: U-net | |
| # The app<a name="app"> | |
| ''', | |
| unsafe_allow_html=True | |
| ) | |
| return | |
| def get_model( num_classes, model_state ): | |
| mynn = mspec.MultiSpectNet( | |
| img_size=[ 256, 256 ], | |
| num_channels=7, | |
| num_classes=num_classes, | |
| learning_rate=1.e-3 | |
| ) | |
| mynn.model.load_state_dict( torch.load( model_state ) ) | |
| return mynn | |