import os
import numpy as np
from osgeo import gdal
import cv2
import tensorflow as tf
import time
def normalize(x):
    x = x/float(np.max(x))
    offset = np.mean(x, 0)
    scale = np.std(x, 0).clip(min=1)
    return (x - offset) / scale

def run(img_path):
    bands_data = []
    raster_dataset = gdal.Open(img_path, gdal.GA_ReadOnly)
    geo_transform = raster_dataset.GetGeoTransform()
    proj = raster_dataset.GetProjectionRef()
    for b in range(1, raster_dataset.RasterCount+1):
        band = raster_dataset.GetRasterBand(b)
        bands_data.append(band.ReadAsArray())
        data = np.dstack(bands_data)
    data = np.expand_dims(normalize(cv2.resize(data,(80,80))),axis = 0)

    with tf.Session() as sess:
       new_saver = tf.train.import_meta_graph("weights/model.ckpt.meta")
       new_saver.restore(sess, tf.train.latest_checkpoint('weights/'))
       graph = tf.get_default_graph()
       inputs = graph.get_tensor_by_name("inputs:0")
       prediction = graph.get_tensor_by_name("out/BiasAdd:0")
       pred = tf.nn.sigmoid(prediction)
       out_img = sess.run([pred],feed_dict={inputs: data})

    ''' can get node of the load graph'''
    # for op in tf.get_default_graph().get_operations():
    # 	print str(op.name)
    cv2.imwrite("out_img_15.jpg",np.array(out_img)[0,0,:,:,0]*255)

run("/home/ispladmin/GIS_workspace/classification/FCN_on_new_data/final_waterbody_extraction/Sentinel_Rawdata/TrainingSamples- 15.tif")