From c5c6c78c2c62a841ad2bbdf7474ff5afb3a04971 Mon Sep 17 00:00:00 2001
From: ryanlu41 <ryanlu41@gmail.com>
Date: Sat, 18 Apr 2020 15:31:22 -0700
Subject: [PATCH] asdf

---
 brainlit/preprocessing/features/base.py | 36 +++++++++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/brainlit/preprocessing/features/base.py b/brainlit/preprocessing/features/base.py
index 8ba79f89c..1c08f1af8 100644
--- a/brainlit/preprocessing/features/base.py
+++ b/brainlit/preprocessing/features/base.py
@@ -18,6 +18,9 @@ def __init__(self, url, size=[1, 1, 1], offset=[15, 15, 15]):
         self.url = url
         self.size = size
         self.offset = offset
+        self.download_time = 0
+        self.conversion_time = 0
+        self.write_time = 0
 
     @abstractmethod
     def _convert_to_features(self, img):
@@ -86,6 +89,8 @@ def fit(
         ngl = NeuroglancerSession(self.url)
         ngl_skel = NeuroglancerSession(self.url + "_segments")
 
+
+
         if start_seg is not None:
             seg_ids = seg_ids[seg_ids.index(start_seg) :]
 
@@ -97,12 +102,25 @@ def fit(
                 verts = segment.vertices[start_vert:]
             start_vert = 0
             for v_id, vertex in enumerate(verts):
+
+                start = time.time()
+
                 img, bounds, voxel = ngl.pull_voxel(
                     seg_id, v_id, self.size[0], self.size[1], self.size[2]
                 )
                 img_off = ngl.pull_bounds_img(bounds + self.offset)
+
+                end = time.time()
+                self.download_time += (end-start)
+
+                start = time.time()
+
                 features = self._convert_to_features(img, include_neighborhood)
                 features_off = self._convert_to_features(img_off, include_neighborhood)
+
+                end = time.time()
+                self.conversion_time += (end-start)
+
                 voxel_dict[counter] = {
                     **{"Segment": int(seg_id), "Vertex": int(v_id), "Label": 0},
                     **features,
@@ -127,10 +145,19 @@ def fit(
                             + str(v_id)
                             + ".feather"
                         )
+
+                        start = time.time()
+
                         feather.write_dataframe(df, path)
+
+                        end = time.time()
+                        self.write_time += (end-start)
+
                         voxel_dict = {}
                         batch_id += 1
 
+
+
         if file_path is None:
             df = pd.DataFrame.from_dict(voxel_dict, "index")
             return df
@@ -148,4 +175,13 @@ def fit(
                     + str(v_id)
                     + ".feather"
                 )
+
+                start = time.time()
+
                 feather.write_dataframe(df, path)
+
+                end = time.time()
+                self.write_time += (end-start)
+    
+    def time(self):
+        return self.download_time, self.conversion_time, self.write_time