You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
52 lines
2.1 KiB
52 lines
2.1 KiB
import numpy as np |
|
import torch |
|
|
|
from ultralytics.yolo.engine.predictor import BasePredictor |
|
from ultralytics.yolo.engine.results import Results |
|
from ultralytics.yolo.utils.torch_utils import select_device |
|
|
|
from .modules.mask_generator import SamAutomaticMaskGenerator |
|
|
|
|
|
class Predictor(BasePredictor): |
|
|
|
def preprocess(self, im): |
|
"""Prepares input image for inference.""" |
|
# TODO: Only support bs=1 for now |
|
# im = ResizeLongestSide(1024).apply_image(im[0]) |
|
# im = torch.as_tensor(im, device=self.device) |
|
# im = im.permute(2, 0, 1).contiguous()[None, :, :, :] |
|
return im[0] |
|
|
|
def setup_model(self, model): |
|
"""Set up YOLO model with specified thresholds and device.""" |
|
device = select_device(self.args.device) |
|
model.eval() |
|
self.model = SamAutomaticMaskGenerator(model.to(device), |
|
pred_iou_thresh=self.args.conf, |
|
box_nms_thresh=self.args.iou) |
|
self.device = device |
|
# TODO: Temporary settings for compatibility |
|
self.model.pt = False |
|
self.model.triton = False |
|
self.model.stride = 32 |
|
self.model.fp16 = False |
|
self.done_warmup = True |
|
|
|
def postprocess(self, preds, path, orig_imgs): |
|
"""Postprocesses inference output predictions to create detection masks for objects.""" |
|
names = dict(enumerate(list(range(len(preds))))) |
|
results = [] |
|
# TODO |
|
for i, pred in enumerate([preds]): |
|
masks = torch.from_numpy(np.stack([p['segmentation'] for p in pred], axis=0)) |
|
orig_img = orig_imgs[i] if isinstance(orig_imgs, list) else orig_imgs |
|
path = self.batch[0] |
|
img_path = path[i] if isinstance(path, list) else path |
|
results.append(Results(orig_img=orig_img, path=img_path, names=names, masks=masks)) |
|
return results |
|
|
|
# def __call__(self, source=None, model=None, stream=False): |
|
# frame = cv2.imread(source) |
|
# preds = self.model.generate(frame) |
|
# return self.postprocess(preds, source, frame)
|
|
|