Memory management issues

[oscarbau]

Hi all,

I am in the process of segmenting a drone orthophoto (226Mb) that covers around 5ha, reprocessed at 10cm per pixel.

I have tested the text prompt segmentation and pin prompt segmentation with poor results.

Now I am testing the general segmentation of the whole mosaic. But I have found errors related to memory use when conducting the process. I am using sam.generate(image, mask)

I have around 20Gb of ram memory available and I am not using GPUs.

I am wondering if there is a way to limit memory usage to avoid this kind of issue.

See the error below:

`

RuntimeError Traceback (most recent call last)
Cell In[9], line 3
1 #mask = 'Ut_NL_segment.tiff'
2 mask = 'C:/UAS/Delgesso5/3_dsm_ortho/2_mosaic/mosaic_segment_all.tiff'
----> 3 sam.generate(image, mask)

File ~\miniconda3\envs\geo\lib\site-packages\samgeo\samgeo.py:208, in SamGeo.generate(self, source, output, foreground, batch, erosion_kernel, mask_multiplier, unique, **kwargs)
206 self.image = image # Store the input image as a numpy array
207 mask_generator = self.mask_generator # The automatic mask generator
--> 208 masks = mask_generator.generate(image) # Segment the input image
209 self.masks = masks # Store the masks as a list of dictionaries
210 self.batch = False

File ~\miniconda3\envs\geo\lib\site-packages\torch\autograd\grad_mode.py:27, in _DecoratorContextManager.call..decorate_context(*args, **kwargs)
24 @functools.wraps(func)
25 def decorate_context(*args, **kwargs):
26 with self.clone():
---> 27 return func(*args, **kwargs)

File ~\miniconda3\envs\geo\lib\site-packages\segment_anything\automatic_mask_generator.py:163, in SamAutomaticMaskGenerator.generate(self, image)
138 """
139 Generates masks for the given image.
140
(...)
159 the mask, given in XYWH format.
160 """
162 # Generate masks
--> 163 mask_data = self._generate_masks(image)
165 # Filter small disconnected regions and holes in masks
166 if self.min_mask_region_area > 0:

File ~\miniconda3\envs\geo\lib\site-packages\segment_anything\automatic_mask_generator.py:206, in SamAutomaticMaskGenerator._generate_masks(self, image)
204 data = MaskData()
205 for crop_box, layer_idx in zip(crop_boxes, layer_idxs):
--> 206 crop_data = self._process_crop(image, crop_box, layer_idx, orig_size)
207 data.cat(crop_data)
209 # Remove duplicate masks between crops

File ~\miniconda3\envs\geo\lib\site-packages\segment_anything\automatic_mask_generator.py:245, in SamAutomaticMaskGenerator._process_crop(self, image, crop_box, crop_layer_idx, orig_size)
243 data = MaskData()
244 for (points,) in batch_iterator(self.points_per_batch, points_for_image):
--> 245 batch_data = self._process_batch(points, cropped_im_size, crop_box, orig_size)
246 data.cat(batch_data)
247 del batch_data

File ~\miniconda3\envs\geo\lib\site-packages\segment_anything\automatic_mask_generator.py:297, in SamAutomaticMaskGenerator._process_batch(self, points, im_size, crop_box, orig_size)
295 if self.pred_iou_thresh > 0.0:
296 keep_mask = data["iou_preds"] > self.pred_iou_thresh
--> 297 data.filter(keep_mask)
299 # Calculate stability score
300 data["stability_score"] = calculate_stability_score(
301 data["masks"], self.predictor.model.mask_threshold, self.stability_score_offset
302 )

File ~\miniconda3\envs\geo\lib\site-packages\segment_anything\utils\amg.py:49, in MaskData.filter(self, keep)
47 self._stats[k] = None
48 elif isinstance(v, torch.Tensor):
---> 49 self._stats[k] = v[torch.as_tensor(keep, device=v.device)]
50 elif isinstance(v, np.ndarray):
51 self._stats[k] = v[keep.detach().cpu().numpy()]

RuntimeError: [enforce fail at C:\b\abs_f0dma8qm3d\croot\pytorch_1669187301762\work\c10\core\impl\alloc_cpu.cpp:81] data. DefaultCPUAllocator: not enough memory: you tried to allocate 49592930604 bytes.
`

Any ideas or leads are appreciated.

Thank you!

[giswqs]

This seems a SAM issue rather than samgeo. Under the hood, samgeo class segment_anything\automatic_mask_generator.py to segment the image. The memory is managed by SAM. You can open an issue at https://github.com/facebookresearch/segment-anything