์ฝ˜ํ…์ธ ๋กœ ๊ฑด๋„ˆ๋›ฐ๊ธฐ

์ฐธ์กฐ ultralytics/models/fastsam/predict.py

์ฐธ๊ณ 

์ด ํŒŒ์ผ์€ https://github.com/ultralytics/ ultralytics/blob/main/ ultralytics/models/ fastsam/predict .py์—์„œ ํ™•์ธํ•  ์ˆ˜ ์žˆ์Šต๋‹ˆ๋‹ค. ๋ฌธ์ œ๋ฅผ ๋ฐœ๊ฒฌํ•˜๋ฉด ํ’€ ๋ฆฌํ€˜์ŠคํŠธ ๐Ÿ› ๏ธ ์— ๊ธฐ์—ฌํ•˜์—ฌ ๋ฌธ์ œ๋ฅผ ํ•ด๊ฒฐํ•˜๋„๋ก ๋„์™€์ฃผ์„ธ์š”. ๊ฐ์‚ฌํ•ฉ๋‹ˆ๋‹ค ๐Ÿ™!



ultralytics.models.fastsam.predict.FastSAMPredictor

๋ฒ ์ด์Šค: DetectionPredictor

FastSAMPredictor๋Š” Ultralytics YOLO ํ”„๋ ˆ์ž„์›Œํฌ์—์„œ ๋น ๋ฅธ SAM (Segment Anything Model) ์„ธ๊ทธ๋จผํŠธ ์˜ˆ์ธก ์ž‘์—…์— ํŠนํ™”๋˜์–ด ์žˆ์Šต๋‹ˆ๋‹ค.

์ด ํด๋ž˜์Šค๋Š” DetectionPredictor๋ฅผ ํ™•์žฅํ•˜์—ฌ ์˜ˆ์ธก ํŒŒ์ดํ”„๋ผ์ธ์„ ํŠน๋ณ„ํžˆ ๋น ๋ฅธ SAM ์— ๋งž๊ฒŒ ์‚ฌ์šฉ์ž ์ •์˜ํ•ฉ๋‹ˆ๋‹ค. ๋งˆ์Šคํฌ ์˜ˆ์ธก๊ณผ ๋น„์ตœ๋Œ€ ์–ต์ œ๋ฅผ ํ†ตํ•ฉํ•˜๋„๋ก ํ›„์ฒ˜๋ฆฌ ๋‹จ๊ณ„๋ฅผ ์กฐ์ •ํ•˜๋Š” ๋™์‹œ์— ๋‹จ์ผ ํด๋ž˜์Šค ์„ธ๋ถ„ํ™”๋ฅผ ์œ„ํ•ด ์ตœ์ ํ™”ํ•ฉ๋‹ˆ๋‹ค. ๋ฅผ ์ตœ์ ํ™”ํ•˜๋ฉด์„œ ํ›„์ฒ˜๋ฆฌ ๋‹จ๊ณ„๋ฅผ ์กฐ์ •ํ•ฉ๋‹ˆ๋‹ค.

์†์„ฑ:

์ด๋ฆ„ ์œ ํ˜• ์„ค๋ช…
cfg dict

์˜ˆ์ธก์„ ์œ„ํ•œ ๊ตฌ์„ฑ ๋งค๊ฐœ๋ณ€์ˆ˜.

overrides dict

์‚ฌ์šฉ์ž ์ง€์ • ๋™์ž‘์— ๋Œ€ํ•œ ์„ ํƒ์  ๋งค๊ฐœ๋ณ€์ˆ˜ ์˜ค๋ฒ„๋ผ์ด๋“œ.

_callbacks dict

์˜ˆ์ธก ์ค‘์— ํ˜ธ์ถœํ•  ์ฝœ๋ฐฑ ํ•จ์ˆ˜์˜ ์„ ํƒ์  ๋ชฉ๋ก์ž…๋‹ˆ๋‹ค.

์˜ ์†Œ์Šค ์ฝ”๋“œ ultralytics/models/fastsam/predict.py
class FastSAMPredictor(DetectionPredictor):
    """
    FastSAMPredictor is specialized for fast SAM (Segment Anything Model) segmentation prediction tasks in Ultralytics
    YOLO framework.

    This class extends the DetectionPredictor, customizing the prediction pipeline specifically for fast SAM.
    It adjusts post-processing steps to incorporate mask prediction and non-max suppression while optimizing
    for single-class segmentation.

    Attributes:
        cfg (dict): Configuration parameters for prediction.
        overrides (dict, optional): Optional parameter overrides for custom behavior.
        _callbacks (dict, optional): Optional list of callback functions to be invoked during prediction.
    """

    def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
        """
        Initializes the FastSAMPredictor class, inheriting from DetectionPredictor and setting the task to 'segment'.

        Args:
            cfg (dict): Configuration parameters for prediction.
            overrides (dict, optional): Optional parameter overrides for custom behavior.
            _callbacks (dict, optional): Optional list of callback functions to be invoked during prediction.
        """
        super().__init__(cfg, overrides, _callbacks)
        self.args.task = "segment"

    def postprocess(self, preds, img, orig_imgs):
        """
        Perform post-processing steps on predictions, including non-max suppression and scaling boxes to original image
        size, and returns the final results.

        Args:
            preds (list): The raw output predictions from the model.
            img (torch.Tensor): The processed image tensor.
            orig_imgs (list | torch.Tensor): The original image or list of images.

        Returns:
            (list): A list of Results objects, each containing processed boxes, masks, and other metadata.
        """
        p = ops.non_max_suppression(
            preds[0],
            self.args.conf,
            self.args.iou,
            agnostic=self.args.agnostic_nms,
            max_det=self.args.max_det,
            nc=1,  # set to 1 class since SAM has no class predictions
            classes=self.args.classes,
        )
        full_box = torch.zeros(p[0].shape[1], device=p[0].device)
        full_box[2], full_box[3], full_box[4], full_box[6:] = img.shape[3], img.shape[2], 1.0, 1.0
        full_box = full_box.view(1, -1)
        critical_iou_index = bbox_iou(full_box[0][:4], p[0][:, :4], iou_thres=0.9, image_shape=img.shape[2:])
        if critical_iou_index.numel() != 0:
            full_box[0][4] = p[0][critical_iou_index][:, 4]
            full_box[0][6:] = p[0][critical_iou_index][:, 6:]
            p[0][critical_iou_index] = full_box

        if not isinstance(orig_imgs, list):  # input images are a torch.Tensor, not a list
            orig_imgs = ops.convert_torch2numpy_batch(orig_imgs)

        results = []
        proto = preds[1][-1] if len(preds[1]) == 3 else preds[1]  # second output is len 3 if pt, but only 1 if exported
        for i, pred in enumerate(p):
            orig_img = orig_imgs[i]
            img_path = self.batch[0][i]
            if not len(pred):  # save empty boxes
                masks = None
            elif self.args.retina_masks:
                pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape)
                masks = ops.process_mask_native(proto[i], pred[:, 6:], pred[:, :4], orig_img.shape[:2])  # HWC
            else:
                masks = ops.process_mask(proto[i], pred[:, 6:], pred[:, :4], img.shape[2:], upsample=True)  # HWC
                pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape)
            results.append(Results(orig_img, path=img_path, names=self.model.names, boxes=pred[:, :6], masks=masks))
        return results

__init__(cfg=DEFAULT_CFG, overrides=None, _callbacks=None)

๊ฐ์ง€ ์˜ˆ์ธก์ž์—์„œ ์ƒ์†ํ•˜๊ณ  ์ž‘์—…์„ '์„ธ๊ทธ๋จผํŠธ'๋กœ ์„ค์ •ํ•˜์—ฌ FastSAMPredictor ํด๋ž˜์Šค๋ฅผ ์ดˆ๊ธฐํ™”ํ•ฉ๋‹ˆ๋‹ค.

๋งค๊ฐœ๋ณ€์ˆ˜:

์ด๋ฆ„ ์œ ํ˜• ์„ค๋ช… ๊ธฐ๋ณธ๊ฐ’
cfg dict

์˜ˆ์ธก์„ ์œ„ํ•œ ๊ตฌ์„ฑ ๋งค๊ฐœ๋ณ€์ˆ˜.

DEFAULT_CFG
overrides dict

์‚ฌ์šฉ์ž ์ง€์ • ๋™์ž‘์— ๋Œ€ํ•œ ์„ ํƒ์  ๋งค๊ฐœ๋ณ€์ˆ˜ ์˜ค๋ฒ„๋ผ์ด๋“œ.

None
_callbacks dict

์˜ˆ์ธก ์ค‘์— ํ˜ธ์ถœํ•  ์ฝœ๋ฐฑ ํ•จ์ˆ˜์˜ ์„ ํƒ์  ๋ชฉ๋ก์ž…๋‹ˆ๋‹ค.

None
์˜ ์†Œ์Šค ์ฝ”๋“œ ultralytics/models/fastsam/predict.py
def __init__(self, cfg=DEFAULT_CFG, overrides=None, _callbacks=None):
    """
    Initializes the FastSAMPredictor class, inheriting from DetectionPredictor and setting the task to 'segment'.

    Args:
        cfg (dict): Configuration parameters for prediction.
        overrides (dict, optional): Optional parameter overrides for custom behavior.
        _callbacks (dict, optional): Optional list of callback functions to be invoked during prediction.
    """
    super().__init__(cfg, overrides, _callbacks)
    self.args.task = "segment"

postprocess(preds, img, orig_imgs)

๋น„์ตœ๋Œ€ ์–ต์ œ ๋ฐ ์›๋ณธ ์ด๋ฏธ์ง€ ํฌ๊ธฐ๋กœ ์ƒ์ž ํฌ๊ธฐ ์กฐ์ • ๋“ฑ ์˜ˆ์ธก์— ๋Œ€ํ•œ ํ›„์ฒ˜๋ฆฌ ๋‹จ๊ณ„๋ฅผ ์ˆ˜ํ–‰ํ•ฉ๋‹ˆ๋‹ค. ํฌ๊ธฐ๋กœ ์กฐ์ •ํ•˜๋Š” ๋“ฑ์˜ ํ›„์ฒ˜๋ฆฌ ๋‹จ๊ณ„๋ฅผ ์ˆ˜ํ–‰ํ•˜๊ณ  ์ตœ์ข… ๊ฒฐ๊ณผ๋ฅผ ๋ฐ˜ํ™˜ํ•ฉ๋‹ˆ๋‹ค.

๋งค๊ฐœ๋ณ€์ˆ˜:

์ด๋ฆ„ ์œ ํ˜• ์„ค๋ช… ๊ธฐ๋ณธ๊ฐ’
preds list

๋ชจ๋ธ์˜ ์›์‹œ ์ถœ๋ ฅ ์˜ˆ์ธก์ž…๋‹ˆ๋‹ค.

ํ•„์ˆ˜
img Tensor

์ฒ˜๋ฆฌ๋œ ์ด๋ฏธ์ง€ tensor.

ํ•„์ˆ˜
orig_imgs list | Tensor

์›๋ณธ ์ด๋ฏธ์ง€ ๋˜๋Š” ์ด๋ฏธ์ง€ ๋ชฉ๋ก์ž…๋‹ˆ๋‹ค.

ํ•„์ˆ˜

๋ฐ˜ํ™˜ํ•ฉ๋‹ˆ๋‹ค:

์œ ํ˜• ์„ค๋ช…
list

์ฒ˜๋ฆฌ๋œ ์ƒ์ž, ๋งˆ์Šคํฌ ๋ฐ ๊ธฐํƒ€ ๋ฉ”ํƒ€๋ฐ์ดํ„ฐ๊ฐ€ ๊ฐ๊ฐ ํฌํ•จ๋œ ๊ฒฐ๊ณผ ๊ฐœ์ฒด ๋ชฉ๋ก์ž…๋‹ˆ๋‹ค.

์˜ ์†Œ์Šค ์ฝ”๋“œ ultralytics/models/fastsam/predict.py
def postprocess(self, preds, img, orig_imgs):
    """
    Perform post-processing steps on predictions, including non-max suppression and scaling boxes to original image
    size, and returns the final results.

    Args:
        preds (list): The raw output predictions from the model.
        img (torch.Tensor): The processed image tensor.
        orig_imgs (list | torch.Tensor): The original image or list of images.

    Returns:
        (list): A list of Results objects, each containing processed boxes, masks, and other metadata.
    """
    p = ops.non_max_suppression(
        preds[0],
        self.args.conf,
        self.args.iou,
        agnostic=self.args.agnostic_nms,
        max_det=self.args.max_det,
        nc=1,  # set to 1 class since SAM has no class predictions
        classes=self.args.classes,
    )
    full_box = torch.zeros(p[0].shape[1], device=p[0].device)
    full_box[2], full_box[3], full_box[4], full_box[6:] = img.shape[3], img.shape[2], 1.0, 1.0
    full_box = full_box.view(1, -1)
    critical_iou_index = bbox_iou(full_box[0][:4], p[0][:, :4], iou_thres=0.9, image_shape=img.shape[2:])
    if critical_iou_index.numel() != 0:
        full_box[0][4] = p[0][critical_iou_index][:, 4]
        full_box[0][6:] = p[0][critical_iou_index][:, 6:]
        p[0][critical_iou_index] = full_box

    if not isinstance(orig_imgs, list):  # input images are a torch.Tensor, not a list
        orig_imgs = ops.convert_torch2numpy_batch(orig_imgs)

    results = []
    proto = preds[1][-1] if len(preds[1]) == 3 else preds[1]  # second output is len 3 if pt, but only 1 if exported
    for i, pred in enumerate(p):
        orig_img = orig_imgs[i]
        img_path = self.batch[0][i]
        if not len(pred):  # save empty boxes
            masks = None
        elif self.args.retina_masks:
            pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape)
            masks = ops.process_mask_native(proto[i], pred[:, 6:], pred[:, :4], orig_img.shape[:2])  # HWC
        else:
            masks = ops.process_mask(proto[i], pred[:, 6:], pred[:, :4], img.shape[2:], upsample=True)  # HWC
            pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape)
        results.append(Results(orig_img, path=img_path, names=self.model.names, boxes=pred[:, :6], masks=masks))
    return results





์ƒ์„ฑ 2023-11-12, ์—…๋ฐ์ดํŠธ 2023-11-25
์ž‘์„ฑ์ž: glenn-jocher (3)