%tensorflow_version 1.x
import tensorflow as tf
print(tf.__version__)

# Silence deprecation warnings for now.
tf.logging.set_verbosity(tf.logging.ERROR)

device_name = tf.test.gpu_device_name()
if device_name != '/device:GPU:0':
  print('GPU device not found')
  gpu = False
else:
  print('Found GPU at: {}'.format(device_name))
  gpu = True

"""
Explanation: Load TensorFlow
Go to Edit->Notebook settings to confirm you have a GPU accelerated kernel.
End of explanation
"""


!git clone https://github.com/google/ffn.git
%cd ffn

from google.protobuf import text_format
from ffn.inference import inference
from ffn.inference import inference_pb2

# Download the example datasets.
!mkdir -p third_party/neuroproof_examples
!gsutil rsync -r -x ".*.gz" gs://ffn-flyem-fib25/ third_party/neuroproof_examples/

"""
Explanation: Set up FFN code and sample data
Colab already provides most of the dependencies.
End of explanation
"""


config = '''image {
  hdf5: "third_party/neuroproof_examples/training_sample2/grayscale_maps.h5:raw"
}
image_mean: 128
image_stddev: 33
checkpoint_interval: 1800
seed_policy: "PolicyPeaks"
model_checkpoint_path: "models/fib25/model.ckpt-27465036"
model_name: "convstack_3d.ConvStack3DFFNModel"
model_args: "{\\"depth\\": 12, \\"fov_size\\": [33, 33, 33], \\"deltas\\": [8, 8, 8]}"
segmentation_output_dir: "results/fib25/training2"
inference_options {
  init_activation: 0.95
  pad_value: 0.05
  move_threshold: 0.9
  min_boundary_dist { x: 1 y: 1 z: 1}
  segment_threshold: 0.6
  min_segment_size: 1000
}'''
req = inference_pb2.InferenceRequest()
_ = text_format.Parse(config, req)

runner = inference.Runner()
runner.start(req)
canvas, alignment = runner.make_canvas((0, 0, 0), (250, 250, 250))

# Create a single segment, starting from the specified origin point.
if gpu:
  vis_update = 20
else:
  vis_update = 1
canvas.segment_at((125, 125, 125),  # zyx
                  dynamic_image=inference.DynamicImage(),
                  vis_update_every=vis_update)

"""
Explanation: Run inference
End of explanation
"""


