退购1.1定位算法

ultralytics/vit/sam/predict.py

@@ -0,0 +1,52 @@
+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)