✅ Address review comments

Author: shaneahmed

tests/engines/test_nucleus_detection_engine.py

@@ -14,7 +14,6 @@
 from tiatoolbox.annotation.storage import SQLiteStore
 from tiatoolbox.models.engine.nucleus_detector import (
     NucleusDetector,
-    # _flatten_predictions_to_dask,
 )
 from tiatoolbox.utils import env_detection as toolbox_env
 from tiatoolbox.utils.misc import imwrite

tiatoolbox/models/engine/nucleus_detector.py

@@ -465,8 +465,8 @@ def post_process_wsi(
         )
 
         centroid_maps = da.map_overlap(
-            rechunked_prediction_map,
             self.model.postproc,
+            rechunked_prediction_map,
             min_distance=min_distance,
             threshold_abs=threshold_abs,
             threshold_rel=threshold_rel,

tiatoolbox/models/engine/semantic_segmentor.py

@@ -1369,7 +1369,7 @@ def prepare_full_batch(
     full_output_dict = {tuple(row): i for i, row in enumerate(full_output_locs)}
     matches = [full_output_dict[tuple(row)] for row in batch_locs]
 
-    total_size = np.max(matches).astype(np.uint16) + 1
+    total_size = np.max(matches).astype(np.uint32) + 1
 
     # Initialize full output array
     full_batch_output = np.zeros(