<jupyter_start><jupyter_text>Load the model and the model state from the .pt file. It may need to be downloaded from AWS S3<jupyter_code>model_path = 'runs/debug/model_e100_b32_no_aug_GPUPAR.pt' # RGB batch size 32
model_path = 'runs/debug/model_e100_b8_no_aug_GPUPAR.pt' # RGB batch size 8
model_path = 'runs/debug/MS_model_e100_b8_no_aug_GPUPAR.pt' # MS model
model_path = 'runs/debug/model_e44_s34938.pt' # RGB early exit (large batch size, i think)

model = unet11(pretrained=False)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # get device for gpu or cpu
if torch.cuda.device_count() > 1:
    print("Let's use", torch.cuda.device_count(), "GPUs!")
    model = nn.DataParallel(model)
    model.to(device)
    
elif torch.cuda.is_available():
    model.to(device)

model.eval()


#load model
if os.path.exists(model_path):
    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))
img_2m.shape, torch.cuda.is_available()
## try to the full dataset
# input_img = torch.unsqueeze(variable(dg_dataset[test_ind], volatile=True), dim=0)



img_arr = img_2m[[1,2,4], :2048,:2048].compute() #BGR
# img_arr = img_2m[:, :2048,:2048].compute() # all 8
img_arr = img_transform(torch.from_numpy(img_arr))

input_img = torch.unsqueeze(variable(img_arr, volatile=True), dim=0)
del img_arr
# model.cpu()
# big_mask = model(input_img.cpu())

big_mask = model(input_img)
big_mask.cpu().detach().numpy().squeeze().shape
## write out with gdal
import gdal, osr

def array2raster(newRasterfn,rasterOrigin,pixelWidth,pixelHeight,array):

    cols = array.shape[1]
    rows = array.shape[0]
    originX = rasterOrigin[0]
    originY = rasterOrigin[1]

    driver = gdal.GetDriverByName('GTiff')
    outRaster = driver.Create(newRasterfn, cols, rows, 1, gdal.GDT_Float32)
    outRaster.SetGeoTransform((originX, pixelWidth, 0, originY, 0, pixelHeight))
    outband = outRaster.GetRasterBand(1)
    outband.WriteArray(array)
    outRasterSRS = osr.SpatialReference()
    outRasterSRS.ImportFromEPSG(4326)
    outRaster.SetProjection(outRasterSRS.ExportToWkt())
    outband.FlushCache()
    
rname = 'runs/debug/napoli_test_eval_BGR_2011.tif'
aff = img_2m.affine
raster_origin = (aff.c, aff.f)
pixel_height = aff.e
pixel_width = aff.a
array2raster(rname, raster_origin, pixel_width, pixel_height, big_mask.cpu().detach().numpy().squeeze())<jupyter_output><empty_output><jupyter_text>plt.figure(figsize=(20,20))big_mask_arr = big_mask.detach().numpy().squeeze()big_mask_arr_bin = big_mask_arr.copy()big_mask_arr_bin[big_mask_arr > 0] = 1big_mask_arr_bin[big_mask_arr <= 0] = 0plt.imshow(big_mask_arr_bin)<jupyter_code>dg_ids = [u'1030010056565000',
  u'103005004AF92E00',
  u'1030050034F99000',
  u'1030010003390200',
  u'1030010007914D00',
  u'10300100048FDE00',
  u'103001002CB22C00',
  u'103001002480AC00',
  u'1030010024043000',
  u'10300100254DF300',
  u'10300100266A6200',
  u'10300100267ACD00',
  u'10300100252EBE00',
  u'103001001E10B700',
  u'1030010010AF3000',
  u'1030010010429400',
  u'103001000ED36500',
  u'103001000F3D1D00',
  u'103001000D952D00',
  u'103001000FD5E300',
  u'103001000FD41000',
  u'10300100088CF100',
  u'1030010010205A00',
  u'1030010010B5F700',
  u'1030010012D61500',
  u'1030010012B84D00',
  u'10300100104F0B00',
  u'1030010011A83400',
  u'1030010010965400',
  u'1030010019087B00',
  u'103001001D438400',
  u'1030010010BCF200',
  u'103001001DB85200',
  u'103001001E11FB00',
  u'10300100187B1F00',
  u'103001001A425F00',
  u'103001001ABE2000',
  u'103001001ACB8500']

dg_dates = [u'2016-05-19T10:02:40.151Z',
  u'2016-04-06T09:49:07.389Z',
  u'2014-10-06T10:08:32.167Z',
  u'2010-02-01T09:58:57.200Z',
  u'2010-09-19T10:23:29.763Z',
  u'2010-03-14T10:04:17.721Z',
  u'2014-03-03T10:11:36.261Z',
  u'2013-06-26T10:18:35.209Z',
  u'2013-06-26T10:18:23.514Z',
  u'2013-07-23T10:22:50.466Z',
  u'2013-08-11T10:21:48.153Z',
  u'2013-08-22T10:15:37.287Z',
  u'2013-08-03T10:16:47.951Z',
  u'2013-01-06T10:18:53.821Z',
  u'2011-12-20T10:29:03.008Z',
  u'2011-12-31T10:24:33.624Z',
  u'2011-11-12T10:25:46.383Z',
  u'2011-11-20T10:31:50.994Z',
  u'2011-08-17T10:18:39.418Z',
  u'2011-11-15T10:15:40.186Z',
  u'2011-12-01T10:27:39.012Z',
  u'2011-01-13T10:10:29.597Z',
  u'2011-12-23T10:18:49.813Z',
  u'2012-01-19T10:25:39.551Z',
  u'2012-04-04T10:28:13.432Z',
  u'2012-03-27T10:22:50.220Z',
  u'2012-01-30T10:21:00.566Z',
  u'2012-01-27T10:31:17.962Z',
  u'2012-01-08T10:30:27.435Z',
  u'2012-04-26T10:18:05.764Z',
  u'2012-12-26T10:24:28.605Z',
  u'2012-01-11T10:20:00.939Z',
  u'2012-11-26T10:29:57.162Z',
  u'2012-12-26T10:24:18.205Z',
  u'2012-05-12T10:28:26.586Z',
  u'2012-07-27T10:28:45.976Z',
  u'2012-07-19T10:23:21.164Z',
  u'2012-07-11T10:18:30.952Z']
import gdal, osr

def array2raster(newRasterfn,rasterOrigin,pixelWidth,pixelHeight,array):

    cols = array.shape[1]
    rows = array.shape[0]
    originX = rasterOrigin[0]
    originY = rasterOrigin[1]

    driver = gdal.GetDriverByName('GTiff')
    outRaster = driver.Create(newRasterfn, cols, rows, 1, gdal.GDT_Float32)
    outRaster.SetGeoTransform((originX, pixelWidth, 0, originY, 0, pixelHeight))
    outband = outRaster.GetRasterBand(1)
    outband.WriteArray(array)
    outRasterSRS = osr.SpatialReference()
    outRasterSRS.ImportFromEPSG(4326)
    outRaster.SetProjection(outRasterSRS.ExportToWkt())
    outband.FlushCache()
model_path = 'runs/debug/model_e44_s34938.pt' # RGB early exit (large batch size, i think)

model = unet11(pretrained=False)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") # get device for gpu or cpu
if torch.cuda.device_count() > 1:
    print("Let's use", torch.cuda.device_count(), "GPUs!")
    model = nn.DataParallel(model)
    model.to(device)
    
elif torch.cuda.is_available():
    model.to(device)

model.eval()


#load model
if os.path.exists(model_path):
    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))
# connect to gbdx
gbdx = Interface()
shp_bounds = (14.159864978661743, 40.81140987011493, 14.302707008009895, 40.95396890196416) # napoli

for d_id, d_dt in zip(dg_ids, dg_dates):
    
    try:
        # get the dask array for the 8 band MS image
        img_2m = CatalogImage(d_id, band_type='MS', bbox=shp_bounds, acomp=True)
        rows, cols = img_2m.shape[1:]
        print(d_id, d_dt, img_2m.shape)
        print('calculating stats\n')
        
        stats = img_2m.display_stats
        means = stats['mean']
        stds = stats['sigma']

        # construct the normalization transform
        img_transform = Compose([
            Normalize(mean=torch.Tensor([means[1], means[2], means[4]]), std=torch.Tensor([stds[1], stds[2], stds[4]]))
        ])

        # subset the image array
        img_arr = img_2m[[1,2,4], :2048,:2048].compute() #BGR
        img_arr = img_transform(torch.from_numpy(img_arr))

        # convert to torch
        input_img = torch.unsqueeze(variable(img_arr, volatile=True), dim=0)
        del img_arr

        # inference
        big_mask = model(input_img)


        ## write out with gdal
        rname = 'runs/debug/napoli_{}_{}_test_eval_BGR.tif'.format(d_id, d_dt[0:10])
        aff = img_2m.affine
        raster_origin = (aff.c, aff.f)
        pixel_height = aff.e
        pixel_width = aff.a
        array2raster(rname, raster_origin, pixel_width, pixel_height, big_mask.cpu().detach().numpy().squeeze())
        
        del input_img
        del img_2m
        del big_mask
        
        torch.cuda.empty_cache()
        
    except Exception as e:
        print(e)
        print('something happened with image {}, passing...'.format(d_id))
        torch.cuda.empty_cache()
        continue
    
print('finished')<jupyter_output>finished
{'fraction_non_alphanumeric': 0.1503079049893176, 'fraction_numerical': 0.1744376021113485, 'mean_word_length': 4.180989583333333, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '421377', 'license': 'MIT', 'n_tokens_mistral': 4196, 'n_tokens_neox': 3364, 'n_words': 474, 'path': 'napoli_test_model_full_subset.ipynb', 'repo_name': 'joemcglinchy/IS_segmentation', 'size': '7957'}