switched to n5 for better mem mang.

mr_robot/mr_robot.py

@@ -4,8 +4,10 @@
 import torch.nn as nn
 import torch.nn.functional as f
 from sklearn.metrics import mean_squared_error
-from model import DUNet2D
+#from model import DUNet2D
 import h5py
+import z5py
+from z5py.converter import convert_from_h5 
 from scipy.ndimage import convolve
 from torch.autograd import Variable
 from collections import OrderedDict
@@ -17,17 +19,21 @@
 from utils import *
 
 patch_size = 16
-
+img_dim = 32
 
 class MrRobot:
     def __init__(self):
         # start the server
         self.new_server = TikTorchServer()
 
+
     def load_data(self):
-        with h5py.File("train.hdf", "r") as f:
-            x = np.array(f.get("volumes/labels/neuron_ids"))
-            y = np.array(f.get("volumes/raw"))
+        self.f = z5py.File('train.n5')
+        return self.f
+        """
+        #with h5py.File("train.hdf", "r") as f:
+        #    x = np.array(f.get("volumes/labels/neuron_ids"))
+        #    y = np.array(f.get("volumes/raw"))
 
         self.labels = []
         self.ip = []
@@ -40,6 +46,7 @@ def load_data(self):
         self.ip = NDArray(np.asarray(self.ip)[:, :, 0:patch_size, 0:patch_size])
         print("data loaded")
         return (ip, labels)
+        """
 
     def load_model(self):
         # load the model
@@ -60,10 +67,12 @@ def resume(self):
         print("training resumed")
 
     def predict(self):
-        self.op = new_server.forward(self.ip)
+        self.ip = np.expand_dims(self.f['volume'][0,0:img_dim, 0:img_dim], axis = 0)
+        #self.label = np.expand_dims(self.f['volumes/labels/neuron_ids'][0,0:img_dim,0:img_dim], axis=0)
+        self.op = self.new_server.forward(self.ip)
         self.op = op.result().as_numpy()
         print("prediction run")
-        return (self.op, self.labels)
+        return self.op
 
     def add(self, row, column):
         self.ip = self.ip.as_numpy()[
@@ -79,32 +88,38 @@ def add(self, row, column):
     def dense_annotate(self, x, y, label, image):
         raise NotImplementedError
 
-    def terminate():
-        new_server.shutdown()
+    def terminate(self):
+        self.new_server.shutdown()
 
 
 class BaseStrategy:
-    def __init__():
-        raise NotImplementedError
+    def __init__(self, file, op):
+        self.f = file
+        self.op = op
 
     # compute loss for a given patch
     def base_loss(self, patch, label):
+        label = label[0][0]
+        patch = patch[0][0]
         result = mean_squared_error(label, patch)  # CHECK THIS
         return result
 
 
 class Strategy1(BaseStrategy):
-    def __init__(self, op, labels):
-        super().__init__()
-        pred_idx = tile_image2D(op[0, 0].shape, 16)
-        actual_idx = tile_image2D(labels[0, 0].shape, 16)
-        w, h, self.row, self.column = 32, 32, -1, -1
+    def __init__(self, file, op):
+        super().__init__(file,op)
+
+    def run(self):
+        idx = tile_image(self.op.shape, patch_size)
+        label = np.expand_dims(self.f['volumes/labels/neuron_ids'][0,0:img_dim,0:img_dim], axis=0)
+        #idx = tile_image(label.shape, patch_size)
+        w, h, self.row, self.column = img_dim, img_dim, -1, -1
         error = 1e7
-        for i in range(len(pred_patches)):
+        for i in range(len(idx)):
             # print(pred_patches[i].shape, actual_patches[i].shape)
-            curr_loss = self.loss(
-                op[0, 0, pred_idx[i][0] : pred_idx[i][1], pred_idx[i][2] : pred_idx[i][3]],
-                labels[0, 0, actual_idx[i][0] : actual_idx[i][1], actual_idx[i][2] : actual_idx[i][3]],
+            curr_loss = super().base_loss(
+                self.op[idx[i][0]: idx[i][1], idx[i][2]:idx[i][3], idx[i][4] : idx[i][5], idx[i][6] : idx[i][7]],
+                labels[idx[i][0]: idx[i][1], idx[i][2]:idx[i][3], idx[i][4] : idx[i][5], idx[i][6] : idx[i][7]],
             )
             print(curr_loss)
             if error > curr_loss:
@@ -128,18 +143,20 @@ def __init__():
 
 
 if __name__ == "__main__":
-
+    
     robo = MrRobot()
-    robo.load_data()
+    file = robo.load_data()
     robo.load_model()
     robo.resume()  # resume training
 
     # run prediction
-    op, label = robo.predict()
+    op = robo.predict()
 
-    metric = Strategy1(op, label)
+    metric = Strategy1(file, op)
+    metric.run()
     row, column = metric.get_patch()
     robo.add(row, column)
 
     # shut down server
     robo.terminate()
+