From 2567b288c949ff31ed4ec2ba5fed5eb67e0284d9 Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Sun, 27 Aug 2023 19:19:13 +0200
Subject: [PATCH] Cleanup redundant SAM `forward()` methods (#4591)
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
---
tests/test_cuda.py | 8 +-
ultralytics/data/loaders.py | 2 +-
ultralytics/models/sam/modules/sam.py | 112 +-------------------------
ultralytics/utils/checks.py | 10 ++-
4 files changed, 14 insertions(+), 118 deletions(-)
diff --git a/tests/test_cuda.py b/tests/test_cuda.py
index da9cc084..ed645567 100644
--- a/tests/test_cuda.py
+++ b/tests/test_cuda.py
@@ -27,12 +27,8 @@ def test_checks():
@pytest.mark.skipif(not CUDA_IS_AVAILABLE, reason='CUDA is not available')
def test_train():
- YOLO(MODEL).train(data=DATA, imgsz=64, epochs=1, batch=-1, device=0) # also test AutoBatch, requires imgsz>=64
-
-
-@pytest.mark.skipif(CUDA_DEVICE_COUNT < 2, reason=f'DDP is not available, {CUDA_DEVICE_COUNT} device(s) found')
-def test_train_ddp():
- YOLO(MODEL).train(data=DATA, imgsz=64, epochs=1, device=[0, 1]) # requires imgsz>=64
+ device = 0 if CUDA_DEVICE_COUNT < 2 else [0, 1]
+ YOLO(MODEL).train(data=DATA, imgsz=64, epochs=1, batch=-1, device=device) # also test AutoBatch, requires imgsz>=64
@pytest.mark.skipif(not CUDA_IS_AVAILABLE, reason='CUDA is not available')
diff --git a/ultralytics/data/loaders.py b/ultralytics/data/loaders.py
index 8d447a65..bed7a19c 100644
--- a/ultralytics/data/loaders.py
+++ b/ultralytics/data/loaders.py
@@ -119,7 +119,7 @@ class LoadStreams:
# Wait until a frame is available in each buffer
while not all(self.imgs):
if not all(x.is_alive() for x in self.threads) or cv2.waitKey(1) == ord('q'): # q to quit
- cv2.destroyAllWindows()
+ self.close()
raise StopIteration
time.sleep(1 / min(self.fps))
diff --git a/ultralytics/models/sam/modules/sam.py b/ultralytics/models/sam/modules/sam.py
index d20b5a93..5649920c 100644
--- a/ultralytics/models/sam/modules/sam.py
+++ b/ultralytics/models/sam/modules/sam.py
@@ -6,11 +6,10 @@
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
-from typing import Any, Dict, List, Tuple
+from typing import List
import torch
from torch import nn
-from torch.nn import functional as F
from .decoders import MaskDecoder
from .encoders import ImageEncoderViT, PromptEncoder
@@ -31,6 +30,9 @@ class Sam(nn.Module):
"""
SAM predicts object masks from an image and input prompts.
+ Note:
+ All forward() operations moved to SAMPredictor.
+
Args:
image_encoder (ImageEncoderViT): The backbone used to encode the image into image embeddings that allow for
efficient mask prediction.
@@ -45,109 +47,3 @@ class Sam(nn.Module):
self.mask_decoder = mask_decoder
self.register_buffer('pixel_mean', torch.Tensor(pixel_mean).view(-1, 1, 1), False)
self.register_buffer('pixel_std', torch.Tensor(pixel_std).view(-1, 1, 1), False)
-
- @property
- def device(self) -> Any:
- return self.pixel_mean.device
-
- @torch.no_grad()
- def forward(
- self,
- batched_input: List[Dict[str, Any]],
- multimask_output: bool,
- ) -> List[Dict[str, torch.Tensor]]:
- """
- Predicts masks end-to-end from provided images and prompts. If prompts are not known in advance, using
- SamPredictor is recommended over calling the model directly.
-
- Args:
- batched_input (list(dict)): A list over input images, each a dictionary with the following keys. A prompt
- key can be excluded if it is not present.
- 'image': The image as a torch tensor in 3xHxW format, already transformed for input to the model.
- 'original_size': (tuple(int, int)) The original size of the image before transformation, as (H, W).
- 'point_coords': (torch.Tensor) Batched point prompts for this image, with shape BxNx2. Already
- transformed to the input frame of the model.
- 'point_labels': (torch.Tensor) Batched labels for point prompts, with shape BxN.
- 'boxes': (torch.Tensor) Batched box inputs, with shape Bx4. Already transformed to the input frame of
- the model.
- 'mask_inputs': (torch.Tensor) Batched mask inputs to the model, in the form Bx1xHxW.
- multimask_output (bool): Whether the model should predict multiple disambiguating masks, or return a single
- mask.
-
- Returns:
- (list(dict)): A list over input images, where each element is as dictionary with the following keys.
- 'masks': (torch.Tensor) Batched binary mask predictions, with shape BxCxHxW, where B is the number of
- input prompts, C is determined by multimask_output, and (H, W) is the original size of the image.
- 'iou_predictions': (torch.Tensor) The model's predictions of mask quality, in shape BxC.
- 'low_res_logits': (torch.Tensor) Low resolution logits with shape BxCxHxW, where H=W=256. Can be passed
- as mask input to subsequent iterations of prediction.
- """
- input_images = torch.stack([self.preprocess(x['image']) for x in batched_input], dim=0)
- image_embeddings = self.image_encoder(input_images)
-
- outputs = []
- for image_record, curr_embedding in zip(batched_input, image_embeddings):
- if 'point_coords' in image_record:
- points = (image_record['point_coords'], image_record['point_labels'])
- else:
- points = None
- sparse_embeddings, dense_embeddings = self.prompt_encoder(
- points=points,
- boxes=image_record.get('boxes', None),
- masks=image_record.get('mask_inputs', None),
- )
- low_res_masks, iou_predictions = self.mask_decoder(
- image_embeddings=curr_embedding.unsqueeze(0),
- image_pe=self.prompt_encoder.get_dense_pe(),
- sparse_prompt_embeddings=sparse_embeddings,
- dense_prompt_embeddings=dense_embeddings,
- multimask_output=multimask_output,
- )
- masks = self.postprocess_masks(
- low_res_masks,
- input_size=image_record['image'].shape[-2:],
- original_size=image_record['original_size'],
- )
- masks = masks > self.mask_threshold
- outputs.append({
- 'masks': masks,
- 'iou_predictions': iou_predictions,
- 'low_res_logits': low_res_masks, })
- return outputs
-
- def postprocess_masks(
- self,
- masks: torch.Tensor,
- input_size: Tuple[int, ...],
- original_size: Tuple[int, ...],
- ) -> torch.Tensor:
- """
- Remove padding and upscale masks to the original image size.
-
- Args:
- masks (torch.Tensor): Batched masks from the mask_decoder, in BxCxHxW format.
- input_size (tuple(int, int)): The size of the model input image, in (H, W) format. Used to remove padding.
- original_size (tuple(int, int)): The original image size before resizing for input to the model, in (H, W).
-
- Returns:
- (torch.Tensor): Batched masks in BxCxHxW format, where (H, W) is given by original_size.
- """
- masks = F.interpolate(
- masks,
- (self.image_encoder.img_size, self.image_encoder.img_size),
- mode='bilinear',
- align_corners=False,
- )
- masks = masks[..., :input_size[0], :input_size[1]]
- return F.interpolate(masks, original_size, mode='bilinear', align_corners=False)
-
- def preprocess(self, x: torch.Tensor) -> torch.Tensor:
- """Normalize pixel values and pad to a square input."""
- # Normalize colors
- x = (x - self.pixel_mean) / self.pixel_std
-
- # Pad
- h, w = x.shape[-2:]
- padh = self.image_encoder.img_size - h
- padw = self.image_encoder.img_size - w
- return F.pad(x, (0, padw, 0, padh))
diff --git a/ultralytics/utils/checks.py b/ultralytics/utils/checks.py
index d1de180b..28cad0dc 100644
--- a/ultralytics/utils/checks.py
+++ b/ultralytics/utils/checks.py
@@ -519,9 +519,13 @@ def cuda_device_count() -> int:
# Run the nvidia-smi command and capture its output
output = subprocess.check_output(['nvidia-smi', '--query-gpu=count', '--format=csv,noheader,nounits'],
encoding='utf-8')
- return int(output.strip())
- except (subprocess.CalledProcessError, FileNotFoundError):
- # If the command fails or nvidia-smi is not found, assume no GPUs are available
+
+ # Take the first line and strip any leading/trailing white space
+ first_line = output.strip().split('\n')[0]
+
+ return int(first_line)
+ except (subprocess.CalledProcessError, FileNotFoundError, ValueError):
+ # If the command fails, nvidia-smi is not found, or output is not an integer, assume no GPUs are available
return 0