<jupyter_start><jupyter_text>DEMO ALTERED. THESE CELLS GENERATE SOME DATA<jupyter_code># import the helper functions
from test_unet_helpers import *

import rasterio
import glob
import os,sys
#from skimage.transform import rescale, resize, downscale_local_mean
from matplotlib import pyplot as plt
import numpy as np
#from scipy import misc
import fiona
# import geopandas as gpd
from shapely.geometry import shape
import shapely
from rasterio.mask import mask
from pyproj import Proj, transform
# specify scene id for DG data
dg_scene_id = '1030010057062200'

# specify image for 2.0 meter analysis
gt_image_2 = '../../rasters/union_impervious_raster_2_0_0.tif' # desktop
gt_image_2 = '../../rasters/union_impervious_raster_2_0_0_wgs84.tif' # desktop
# gt_image_2 = '../rasters/union_impervious_raster_2_0_0.tif' # laptop
#gt_image_2 = 'C:/projects/smart_cities/Denver_Sample_Sites/rasters/union_impervious_raster_2_0.tif'

# specify images for 0.5 meter analysis
gt_image_05 = '../../rasters/union_impervious_raster_0_5.tif' # desktop
# gt_image_05 = '../rasters/union_impervious_raster_0_5.tif' # laptop
#gt_image_05 = 'C:/projects/smart_cities/Denver_Sample_Sites/rasters/union_impervious_raster_0_5.tif'

# specify the shapefile
# shpfile = '../Denver_Sample_Sites/union/union_grid_index256.shp'
shpfile = '../../union/union.shp' # desktop
# shpfile = '../union/union.shp' # desktop
rows = []
with fiona.open(shpfile) as shp:
    crs = shp.crs
#     for s in shp:
#         # rows.append(shape(s['geometry']))
#         pass
        
        
    shp_bounds = shp.bounds
from gbdxtools import Interface, CatalogImage

# connect to gbdx
gbdx = Interface()

# get the dask array for the 8 band MS image
img_2m = CatalogImage(dg_scene_id, band_type='MS', bbox=shp_bounds, acomp=True)
rows, cols = img_2m.shape[1:]

# get the dask array for the Pansharpened, 4 band image
image_05m = CatalogImage(dg_scene_id, bbox=shp_bounds, acomp=True, pansharpen=True)
rowsP, colsP = image_05m.shape[1:]
pan_factorR = np.floor(rowsP/rows)
pan_factorC = np.floor(colsP/cols)
os.path.exists(gt_image_2)
# generate the sample points
# we have some funny stuff from ArcGIS converting the polygons to a raster... so... 
# generate the points within the bounds of the shapefile / DG image array

import random

random.seed(10)

num = 10000
side = 64
coords = []
coords_pan = []
im_patches = []
im_patches_pan = []

with rasterio.open(gt_image_2, 'r') as src:
    
    for i in range(num):
        ran_row, ran_col = random.randrange(0,rows), random.randrange(0,cols)
        print('generating sample {} of {}'.format(i+1, num))

        try:
            r_start = ran_row - side/2
            r_end = ran_row + side/2
            c_start = ran_col - side/2
            c_end = ran_col + side/2
            #im_data = img_2m[:, r_start:r_end, c_start:c_end]

            pr_start = ran_row*pan_factorR - side/2
            pr_end = ran_row*pan_factorR + side/2
            pc_start = ran_col*pan_factorC - side/2
            pc_end = ran_col*pan_factorC + side/2
            #im_data_pan = image_05m[:, pr_start:pr_end, pc_start:pc_end]

            # append regardless and check later.

            #############################################################################################################
            ### SHOULD CHECK THE INFORMATION THRESHOLDS BELOW, BUT HERE, TO RESELECT ROWS AND COLUMNS UNTIL IT IS MET ###
            #############################################################################################################

            s_pt = calcXYfromRC(img_2m.affine, (ran_row, ran_col))
            pt_flag = checkWindow(src, s_pt)
            ct = 0
            while pt_flag:
                ran_row, ran_col = random.randrange(0,rows), random.randrange(0,cols)
                s_pt = calcXYfromRC(img_2m.affine, (ran_row, ran_col))
                pt_flag = checkWindow(src, s_pt)
                
                ct+=1
                # exit loop after 1000 tries
                if ct > 1000:
                    break
            if ct > 1000:
                print('1000 tries, moving on')
                continue
                
                
            coords.append((ran_row, ran_col))
            coords_pan.append(( int(ran_row*pan_factorR), int(ran_col*pan_factorC)))

        except Exception as e:
            print(e)

    coords_arr = np.array(coords)  
lonlat_MS = [calcXYfromRC(img_2m.affine, pair) for pair in coords]
lonlat_PAN = [calcXYfromRC(image_05m.affine, pair) for pair in coords_pan]
import pyproj
       
def assignRC2(rio_obj, samp_pt, window_size=64, inproj='epsg:4326', outproj='epsg:32613'):
    # project the point to source crs
    outProj = Proj(init=outproj)
    inProj = Proj(init=inproj)
    x1,y1 = samp_pt        
    x2,y2 = pyproj.transform(inProj, outProj ,x1,y1)
    x3,y3 = pyproj.transform(outProj, inProj, x2,y2)
    
    # get the row column
    temp = rio_obj.index(x2,y2)
    r,c = [int(c) for c in temp]
    
    return ((x3,y3), samp_pt, r,c)

def assignRC3(rio_obj, samp_pt, window_size=64, inproj='epsg:4326', outproj='epsg:32613'):
    # project the point to source crs
    #outProj = Proj(init=outproj)
    #inProj = Proj(init=inproj)
    x1,y1 = samp_pt        
    #x2,y2 = pyproj.transform(inProj, outProj ,x1,y1)
    #x3,y3 = pyproj.transform(outProj, inProj, x2,y2)
    
    # get the row column
    temp = rio_obj.index(x1,y1)
    r,c = [int(c) for c in temp]
    
    return ((x1,y1), samp_pt, r,c)

# use the lonlats_* lists to sample the ground truth image. return the histogram if the window is on the border,
# which should be evident if the window extends across
with rasterio.open(gt_image_2, 'r') as src:
    #print(src.crs)

    # cull the lon/lat points by checking if they include nodata pixels
    gt_pts_MS = [assignRC3(src, samp_pt) for samp_pt in lonlat_MS]

# do the same for the 0.5 meter ground truth raster
with rasterio.open(gt_image_05, 'r') as src:

    # cull the lon/lat points by checking if they include nodata pixels
    gt_pts_PAN = [assignRC(src, samp_pt) for samp_pt in lonlat_MS]



# rand_inds = random.sample(range(num),10)

# for sample_ind in rand_inds:
#     print('\n\nSample ID: {}'.format(sample_ind))
#     # check the images
#     gt_2m_rowcol = [x[2:] for x in gt_pts_MS]
#     gt_dataset = gtDatasetSampler(gt_image_2, gt_2m_rowcol) # 64 pixel dimension windows
#     # gt_dataset = gtDatasetSampler(gt_image_2, gt_pts_PAN) # 64 pixel dimension windows
#     sample_gt = gt_dataset[sample_ind]
#     print(sample_gt.shape)
#     #plt.imshow(sample_gt)
#     #plt.colorbar()
#     #plt.show()

#     dg_dataset = DigitalGlobeSampler(img_2m, coords)
#     # dg_dataset = DigitalGlobeSampler(img_pan)
#     sample = np.rollaxis(dg_dataset[sample_ind],0,3)
#     print(sample.shape)
#     #plt.imshow(misc.bytescale(sample[:,:,[4,3,1]]))
#     #plt.show()
    
#     plt.imshow(misc.bytescale(sample[:,:,[4,3,1]]))
#     plt.imshow(sample_gt, alpha = 0.15)
#     plt.show()
    

#     # see that they are the same size

# img_file_name = 'data/cropped_train/instrument_dataset_3/images/frame004.jpg' # use GBDX
# gt_file_name = 'data/cropped_train/instrument_dataset_3/binary_masks/frame004.png' # use gt sampler
# # img = load_image(img_file_name)
# img = sample[:,:,[4,3,1]] # DG image

# # gt = cv2.imread(gt_file_name, 0) > 0
# gt = sample_gt
# get the average mean and standard deviation across the bands to apply to the transform
# dg_dataset = DigitalGlobeSampler(img_2m, coords)
# means, stds = [], []
# for i in range(num):
#     temp = np.rollaxis(dg_dataset[i],0,3)
#     temp2 = temp[:,:,[4,3,1]] #rgb
    
#     mean_a = list(temp2.mean(axis=(0,1)))
#     std_a = list(temp2.std(axis=(0,1)))
    
#     means.append(mean_a)
#     stds.append(std_a)
    

# np.array(means).shape, np.array(stds).shape
# np.array(means).mean(axis=0), np.array(stds).mean(axis=0)
    
img_transform = Compose([
    Normalize(mean=[1630.7322, 1574.9552, 1549.1031], std=[837.0847 , 766.4114 , 604.51605])
])
# input_img = torch.unsqueeze(variable(img_transform(a), volatile=True), dim=0)
# mask = model(input_img.cpu())
# mask.shape
# mask_array = mask.data[0].cpu().numpy()[0]
# imshow(mask_array)
# plt.colorbar()
# plt.show()

# imshow(misc.bytescale(img.compute()))
# plt.show()
# dg_dataset2 = DigitalGlobeSamplerTensor(img_2m, coords, transform=img_transform) # will return RGB by default
# dg_dataset2[0].shape, dg_dataset[1].shape<jupyter_output><empty_output><jupyter_text>Train the initial model with a subset of the sample points.. need to edit the source code to use the input data classes instead of data loaders. to use the dataloader, may need to create some attributes on the class such as self.__len__<jupyter_code>val_test_percentage = 0.2
train_ind = int(num - num * val_test_percentage)
val_ind = int((num - train_ind)*0.5)
val_ind = train_ind + val_ind
train_ind, val_ind

## establish the data loaders
from torch.utils.data import DataLoader

# ground truth
gt_transform = Compose([
    ToTensor()
    ])
gt_dataset_train = gtDatasetSampler2(gt_image_2, coords[:train_ind], transform=gt_transform) 
gt_dataset_train = gtDatasetSampler2(gt_image_2, coords[:100], transform=gt_transform) # debug 
gt_dataset_val = gtDatasetSampler2(gt_image_2, coords[train_ind:val_ind], transform=gt_transform)
gt_dataset_test = gtDatasetSampler2(gt_image_2, coords[val_ind:], transform=gt_transform)

gt_dl_train = DataLoader(gt_dataset_train, batch_size=10, shuffle=False)
gt_dl_val = DataLoader(gt_dataset_val, batch_size=10, shuffle=False)


#DG
dg_dataset_train = DigitalGlobeSamplerTensor(img_2m, coords[:train_ind], transform=img_transform) # will return RGB by default
dg_dataset_train = DigitalGlobeSamplerTensor(img_2m, coords[:100], transform=img_transform) # debug
dg_dataset_val = DigitalGlobeSamplerTensor(img_2m, coords[train_ind:val_ind], transform=img_transform)
dg_dataset_test = DigitalGlobeSamplerTensor(img_2m, coords[val_ind:], transform=img_transform)

dg_dl_train = DataLoader(dg_dataset_train, batch_size=10, shuffle=False)
dg_dl_val = DataLoader(dg_dataset_val, batch_size=10, shuffle=False)
gt_dataset_train[0], dg_dataset_train[0].shape
class objectview(object):
    def __init__(self, d):
        self.__dict__ = d
        
args = {}
args['lr'] = 0.0001
args['n_epochs'] = 4
args['batch_size'] = 10
args['root'] = 'runs/debug'
args['fold']=0
args['jaccard_weight'] = 1

# convert dict keys to objects and specify lr
args = objectview(args)
lr = args.lr




import tqdm
from datetime import datetime

def write_event(log, step: int, **data):
    data['step'] = step
    data['dt'] = datetime.now().isoformat()
    log.write(json.dumps(data, sort_keys=True))
    log.write('\n')
    log.flush()
    
    
def train(args, model, criterion, train_loader, valid_loader, validation, init_optimizer, n_epochs=None, fold=None, num_classes=None):
    lr = args.lr
    #n_epochs = n_epochs or args.n_epochs
    n_epochs = n_epochs
    optimizer = init_optimizer(lr)

    root = Path(args.root)
    model_path = root / 'model_{fold}.pt'.format(fold=fold)
    if model_path.exists():
        state = torch.load(str(model_path))
        epoch = state['epoch']
        step = state['step']
        model.load_state_dict(state['model'])
        print('Restored model, epoch {}, step {:,}'.format(epoch, step))
    else:
        epoch = 1
        step = 0

    save = lambda ep: torch.save({
        'model': model.state_dict(),
        'epoch': ep,
        'step': step,
    }, str(model_path))

    report_each = 10
    log = root.joinpath('train_{fold}.log'.format(fold=fold)).open('at', encoding='utf8')
    valid_losses = []
    
    model=model.cuda()    
    
    # record some data
    mean_loss_list = []
    mean_losses = []
    val_loss_list = []
    #for epoch in range(epoch, n_epochs + 1):
    for epoch in range(1,5):
        print('on epoch {}'.format(epoch))
        
        model.train() #try moving above loop
        random.seed()
        #tq = tqdm.tqdm(total=(len(train_loader) * args.batch_size))
        tq = tqdm.tqdm(total=(8000))
        tq.set_description('Epoch {}, lr {}'.format(epoch, lr))
        losses = []
        tl = train_loader
        
        try:
            mean_loss = 0
            for i, (inputs, targets) in enumerate(tl):
                #print('inputs shape: {}'.format(inputs.shape))
                #print('targets shape: {}'.format(targets.shape))
                inputs, targets = variable(inputs).cuda(), variable(targets).cuda()
                outputs = model(inputs)
                loss = criterion(outputs, targets)
                optimizer.zero_grad()
                batch_size = inputs.size(0)
                loss.backward()
                optimizer.step()
                step += 1
                tq.update(batch_size)
                losses.append(loss.data[0])
                mean_loss = np.mean(losses[-report_each:])
                tq.set_postfix(loss='{:.5f}'.format(mean_loss))
                
                #record losses
                mean_losses.append(mean_loss)
                
                if i and i % report_each == 0:
                    write_event(log, step, loss=mean_loss)
                    #pass
                
            write_event(log, step, loss=mean_loss)
            tq.close()
            save(epoch + 1)
            valid_metrics = validation(model, criterion, valid_loader, num_classes)
            write_event(log, step, **valid_metrics)
            valid_loss = valid_metrics['valid_loss']
            valid_losses.append(valid_loss)
            
            
            # record mean losses
            mean_loss_list.append(mean_losses)
            val_loss_list.append(valid_losses)
            
        except KeyboardInterrupt:
            tq.close()
            print('Ctrl+C, saving snapshot')
            save(epoch)
            print('done.')
            return
        
        
        except Exception as e:
            print(e)
        
    return mean_loss_list, val_loss_list
    
def train2(args, model, criterion, train_loader, valid_loader, validation, init_optimizer, n_epochs=None, fold=None, num_classes=None):
    
    
    lr = args.lr
    #n_epochs = n_epochs or args.n_epochs
    n_epochs = n_epochs
    #optimizer = init_optimizer(lr)
    optimizer = optim.Adam(model.parameters(), lr=lr)

    # model on GPU
    model=model.cuda()    
    
    # record some data
    mean_loss_list = []
    mean_losses = []
    val_loss_list = []
    for epoch in [1,2]:
        print('on epoch {}'.format(epoch))
        
        #model.train() #try moving above loop
        #random.seed()
                
        mean_loss = 0.0
        running_loss = 0.0
        print('training...')
        
        tl = train_loader
        for i, (inputs, targets) in enumerate(tl):
            
            # get the inputs
            inputs, targets = variable(inputs).cuda(), variable(targets).cuda()

            # zero parameter gradients
            optimizer.zero_grad()

            # forward + backward + optimize
            outputs = model(inputs)
            loss = criterion(outputs, targets)                
            loss.backward()
            optimizer.step()

            # print statistics
            running_loss += loss
            print('[%d, %5d] loss: %.3f' %(epoch, i + 1, loss))
                    
        
        
    return inputs, targets

def train3(args, model, criterion, train_loader, valid_loader, validation, init_optimizer, n_epochs=None, fold=None, num_classes=None):
    
    
    lr = args.lr
    #n_epochs = n_epochs or args.n_epochs
    n_epochs = n_epochs
    #optimizer = init_optimizer(lr)
    optimizer = optim.Adam(model.parameters(), lr=lr)

    # model on GPU
    model=model.cuda()    
    
    # record some data
    mean_loss_list = []
    mean_losses = []
    val_loss_list = []
    
    
    # try manually going through 2 epochs

    #model.train() #try moving above loop
    #random.seed()

    mean_loss = 0.0
    running_loss = 0.0
    print('training...')

    epoch = 1
    tl = train_loader
    for i, (inputs, targets) in enumerate(tl):

        # get the inputs
        inputs, targets = variable(inputs).cuda(), variable(targets).cuda()

        # zero parameter gradients
        optimizer.zero_grad()

        # forward + backward + optimize
        outputs = model(inputs)
        loss = criterion(outputs, targets)                
        loss.backward()
        optimizer.step()

        # print statistics
        running_loss += loss
        print('[%d, %5d] loss: %.3f' %(epoch, i + 1, loss))
        
    print('manually on epoch 2')
    epoch=2
    tl = train_loader
    for i, (inputs, targets) in enumerate(tl):

        # get the inputs
        inputs, targets = variable(inputs).cuda(), variable(targets).cuda()

        # zero parameter gradients
        optimizer.zero_grad()

        # forward + backward + optimize
        outputs = model(inputs)
        loss = criterion(outputs, targets)                
        loss.backward()
        optimizer.step()

        # print statistics
        running_loss += loss
        print('[%d, %5d] loss: %.3f' %(epoch, i + 1, loss))

        
        
    return inputs, targets

print('gt shape: {}'.format(gt_dataset_train[0].shape))
print(dg_dataset_train[0].shape)
import utils as pyt_utils
from torch.optim import Adam
import torch.backends.cudnn as cudnn
from pathlib import Path
from validation import validation_binary
from loss import LossBinary
import json

# need to change this to DICE loss!
loss = LossBinary(jaccard_weight=args.jaccard_weight)
#criterion = nn.CrossEntropyLoss()


cudnn.benchmark = True

root = Path(args.root)
root.mkdir(exist_ok=True, parents=True)
root.joinpath('params.json').write_text(
    json.dumps(vars(args), indent=True, sort_keys=True))

valid = validation_binary

num_classes=1
# set a model from scratch
model = unet11(pretrained=False)
in_,out_ = train2(
        init_optimizer=lambda lr: Adam(model.parameters(), lr=lr),
        args=args,
        model=model,
        criterion=loss,
        train_loader=zip(dg_dl_train, gt_dl_train),
        valid_loader=zip(dg_dl_val, gt_dl_val),
        validation=valid,
        fold=args.fold,
        num_classes=num_classes,
        n_epochs=6
    )

in_.shape, out_.shape
### call outside training method, no loop

model = unet11(pretrained=False).cuda()
lr = 0.0001
optimizer = Adam(model.parameters(), lr=lr)
criterion = LossBinary(jaccard_weight=args.jaccard_weight)

mean_loss = 0.0
running_loss = 0.0
print('training...')

epoch = 1
tl = zip(dg_dl_train, gt_dl_train)
for i, (inputs, targets) in enumerate(tl):

    # get the inputs
    inputs, targets = variable(inputs).cuda(), variable(targets).cuda()

    # zero parameter gradients
    optimizer.zero_grad()

    # forward + backward + optimize
    outputs = model(inputs)
    i_loss = criterion(outputs, targets)                
    i_loss.backward()
    optimizer.step()

    # print statistics
    running_loss += i_loss
    print('[%d, %5d] loss: %.3f' %(epoch, i + 1, i_loss))

print('manually on epoch 2')
epoch=2
tl = zip(dg_dl_train, gt_dl_train)
for i, (inputs, targets) in enumerate(tl):

    # get the inputs
    inputs, targets = variable(inputs).cuda(), variable(targets).cuda()

    # zero parameter gradients
    optimizer.zero_grad()

    # forward + backward + optimize
    outputs = model(inputs)
    i_loss = criterion(outputs, targets)                
    i_loss.backward()
    optimizer.step()

    # print statistics
    running_loss += i_loss
    print('[%d, %5d] loss: %.3f' %(epoch, i + 1, i_loss))
### call outside training method, in loop

model = unet11(pretrained=False).cuda()
lr = 0.0001
optimizer = Adam(model.parameters(), lr=lr)
criterion = LossBinary(jaccard_weight=args.jaccard_weight)

mean_loss = 0.0
running_loss = 0.0
print('training...')

#tl = zip(dg_dl_train, gt_dl_train)
for i in range(2):
    epoch = i+1
    print('on epoch {}'.format(epoch))
    tl = zip(dg_dl_train, gt_dl_train) # this can't be defined outside the loop. This suggests it is...cashing out??
    for i, (inputs, targets) in enumerate(tl):

        # get the inputs
        inputs, targets = variable(inputs).cuda(), variable(targets).cuda()

        # zero parameter gradients
        optimizer.zero_grad()

        # forward + backward + optimize
        outputs = model(inputs)
        i_loss = criterion(outputs, targets)                
        i_loss.backward()
        optimizer.step()

        # print statistics
        running_loss += i_loss
        print('[%d, %5d] loss: %.3f' %(epoch, i + 1, i_loss))


model.parameters
# check one of the test images
test_ind = 100
input_img = torch.unsqueeze(variable(dg_dataset_val[test_ind], volatile=True), dim=0) 
mask = model(input_img)


mask_im = mask.squeeze().data.cpu().numpy()
mask_binary = mask_im.copy()

thresh = 0
mask_binary[mask_im<=thresh] = 0
mask_binary[mask_im>thresh] = 1
plt.imshow(mask_binary)
plt.colorbar()
plt.show()

source_im = dg_dataset_val[test_ind].numpy()
plt.imshow(misc.bytescale(np.rollaxis(source_im, 0, 3)))
plt.show()

gt_im = gt_dataset_val[test_ind].numpy().squeeze()
plt.imshow(gt_im)
plt.show()
gt_im.shape<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1262514688601645, 'fraction_numerical': 0.01945945945945946, 'mean_word_length': 2.885317750182615, 'pattern_counts': {'":': 0, '<': 10, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '99876', 'license': 'MIT', 'n_tokens_mistral': 7570, 'n_tokens_neox': 7207, 'n_words': 1602, 'path': 'Demo_edit_to_use_DG_data.ipynb', 'repo_name': 'joemcglinchy/IS_segmentation', 'size': '21275'}