deploy: 59dde48

janelia-cellmap · Nov 4, 2024 · 26cf718 · 26cf718
1 parent 5e45ba6
commit 26cf718
Show file tree

Hide file tree

Showing 25 changed files with 3,250 additions and 62 deletions.
diff --git a/_modules/cellmap_data/dataloader.html b/_modules/cellmap_data/dataloader.html
@@ -246,10 +246,10 @@ <h1></h1>
 
   <h1>Source code for cellmap_data.dataloader</h1><div class="highlight"><pre>
 <span></span><span class="kn">import</span> <span class="nn">torch</span>
-<span class="kn">from</span> <span class="nn">torch.utils.data</span> <span class="kn">import</span> <span class="n">DataLoader</span><span class="p">,</span> <span class="n">Sampler</span>
+<span class="kn">from</span> <span class="nn">torch.utils.data</span> <span class="kn">import</span> <span class="n">DataLoader</span><span class="p">,</span> <span class="n">Sampler</span><span class="p">,</span> <span class="n">Subset</span>
 <span class="kn">from</span> <span class="nn">.dataset</span> <span class="kn">import</span> <span class="n">CellMapDataset</span>
 <span class="kn">from</span> <span class="nn">.multidataset</span> <span class="kn">import</span> <span class="n">CellMapMultiDataset</span>
-<span class="kn">from</span> <span class="nn">.subdataset</span> <span class="kn">import</span> <span class="n">CellMapSubset</span>
+<span class="kn">from</span> <span class="nn">.dataset_writer</span> <span class="kn">import</span> <span class="n">CellMapDatasetWriter</span>
 <span class="kn">from</span> <span class="nn">typing</span> <span class="kn">import</span> <span class="n">Callable</span><span class="p">,</span> <span class="n">Iterable</span><span class="p">,</span> <span class="n">Optional</span><span class="p">,</span> <span class="n">Sequence</span>
 
 
@@ -281,7 +281,7 @@ <h1>Source code for cellmap_data.dataloader</h1><div class="highlight"><pre>
 <a class="viewcode-back" href="../../cellmap_data.CellMapDataLoader.html#cellmap_data.CellMapDataLoader.__init__">[docs]</a>
     <span class="k">def</span> <span class="fm">__init__</span><span class="p">(</span>
         <span class="bp">self</span><span class="p">,</span>
-        <span class="n">dataset</span><span class="p">:</span> <span class="n">CellMapMultiDataset</span> <span class="o">|</span> <span class="n">CellMapDataset</span> <span class="o">|</span> <span class="n">CellMapSubset</span><span class="p">,</span>
+        <span class="n">dataset</span><span class="p">:</span> <span class="n">CellMapMultiDataset</span> <span class="o">|</span> <span class="n">CellMapDataset</span> <span class="o">|</span> <span class="n">Subset</span> <span class="o">|</span> <span class="n">CellMapDatasetWriter</span><span class="p">,</span>
         <span class="n">classes</span><span class="p">:</span> <span class="n">Iterable</span><span class="p">[</span><span class="nb">str</span><span class="p">],</span>
         <span class="n">batch_size</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">1</span><span class="p">,</span>
         <span class="n">num_workers</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">0</span><span class="p">,</span>
@@ -373,7 +373,7 @@ <h1>Source code for cellmap_data.dataloader</h1><div class="highlight"><pre>
 <a class="viewcode-back" href="../../cellmap_data.dataloader.html#cellmap_data.CellMapDataLoader.collate_fn">[docs]</a>
     <span class="k">def</span> <span class="nf">collate_fn</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">batch</span><span class="p">:</span> <span class="nb">list</span><span class="p">[</span><span class="nb">dict</span><span class="p">])</span> <span class="o">-&gt;</span> <span class="nb">dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">]:</span>
 <span class="w">        </span><span class="sd">&quot;&quot;&quot;Combine a list of dictionaries from different sources into a single dictionary for output.&quot;&quot;&quot;</span>
-        <span class="n">outputs</span><span class="p">:</span> <span class="nb">dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">]</span> <span class="o">=</span> <span class="p">{}</span>
+        <span class="n">outputs</span> <span class="o">=</span> <span class="p">{}</span>
         <span class="k">for</span> <span class="n">b</span> <span class="ow">in</span> <span class="n">batch</span><span class="p">:</span>
             <span class="k">for</span> <span class="n">key</span><span class="p">,</span> <span class="n">value</span> <span class="ow">in</span> <span class="n">b</span><span class="o">.</span><span class="n">items</span><span class="p">():</span>
                 <span class="k">if</span> <span class="n">key</span> <span class="ow">not</span> <span class="ow">in</span> <span class="n">outputs</span><span class="p">:</span>

diff --git a/_modules/cellmap_data/dataset.html b/_modules/cellmap_data/dataset.html
@@ -663,16 +663,18 @@ <h1>Source code for cellmap_data.dataset</h1><div class="highlight"><pre>
 <a class="viewcode-back" href="../../cellmap_data.dataset.html#cellmap_data.CellMapDataset.__getitem__">[docs]</a>
     <span class="k">def</span> <span class="fm">__getitem__</span><span class="p">(</span><span class="bp">self</span><span class="p">,</span> <span class="n">idx</span><span class="p">:</span> <span class="nb">int</span><span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">dict</span><span class="p">[</span><span class="nb">str</span><span class="p">,</span> <span class="n">torch</span><span class="o">.</span><span class="n">Tensor</span><span class="p">]:</span>
 <span class="w">        </span><span class="sd">&quot;&quot;&quot;Returns a crop of the input and target data as PyTorch tensors, corresponding to the coordinate of the unwrapped index.&quot;&quot;&quot;</span>
+        <span class="n">idx</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">array</span><span class="p">(</span><span class="n">idx</span><span class="p">)</span>
+        <span class="n">idx</span><span class="p">[</span><span class="n">idx</span> <span class="o">&lt;</span> <span class="mi">0</span><span class="p">]</span> <span class="o">=</span> <span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span> <span class="o">+</span> <span class="n">idx</span><span class="p">[</span><span class="n">idx</span> <span class="o">&lt;</span> <span class="mi">0</span><span class="p">]</span>
         <span class="k">try</span><span class="p">:</span>
             <span class="n">center</span> <span class="o">=</span> <span class="n">np</span><span class="o">.</span><span class="n">unravel_index</span><span class="p">(</span>
                 <span class="n">idx</span><span class="p">,</span> <span class="p">[</span><span class="bp">self</span><span class="o">.</span><span class="n">sampling_box_shape</span><span class="p">[</span><span class="n">c</span><span class="p">]</span> <span class="k">for</span> <span class="n">c</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">axis_order</span><span class="p">]</span>
             <span class="p">)</span>
         <span class="k">except</span> <span class="ne">ValueError</span><span class="p">:</span>
+            <span class="c1"># TODO: This is a hacky temprorary fix. Need to figure out why this is happening</span>
             <span class="n">logger</span><span class="o">.</span><span class="n">error</span><span class="p">(</span>
                 <span class="sa">f</span><span class="s2">&quot;Index </span><span class="si">{</span><span class="n">idx</span><span class="si">}</span><span class="s2"> out of bounds for dataset </span><span class="si">{</span><span class="bp">self</span><span class="si">}</span><span class="s2"> of length </span><span class="si">{</span><span class="nb">len</span><span class="p">(</span><span class="bp">self</span><span class="p">)</span><span class="si">}</span><span class="s2">&quot;</span>
             <span class="p">)</span>
             <span class="n">logger</span><span class="o">.</span><span class="n">warning</span><span class="p">(</span><span class="sa">f</span><span class="s2">&quot;Returning closest index in bounds&quot;</span><span class="p">)</span>
-            <span class="c1"># TODO: This is a hacky temprorary fix. Need to figure out why this is happening</span>
             <span class="n">center</span> <span class="o">=</span> <span class="p">[</span><span class="bp">self</span><span class="o">.</span><span class="n">sampling_box_shape</span><span class="p">[</span><span class="n">c</span><span class="p">]</span> <span class="o">-</span> <span class="mi">1</span> <span class="k">for</span> <span class="n">c</span> <span class="ow">in</span> <span class="bp">self</span><span class="o">.</span><span class="n">axis_order</span><span class="p">]</span>
         <span class="n">center</span> <span class="o">=</span> <span class="p">{</span>
             <span class="n">c</span><span class="p">:</span> <span class="n">center</span><span class="p">[</span><span class="n">i</span><span class="p">]</span> <span class="o">*</span> <span class="bp">self</span><span class="o">.</span><span class="n">largest_voxel_sizes</span><span class="p">[</span><span class="n">c</span><span class="p">]</span> <span class="o">+</span> <span class="bp">self</span><span class="o">.</span><span class="n">sampling_box</span><span class="p">[</span><span class="n">c</span><span class="p">][</span><span class="mi">0</span><span class="p">]</span>