Spaces:
Sleeping
Sleeping
| # Ultralytics YOLO 🚀, AGPL-3.0 license | |
| 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) | |