【FastSAM】超簡単、超高速何でもセグメンテーション FastSAMを使ってみる

高速セグメンテーションモデルFastSAM

画像上に写っている様々なものを高速にセグメンテーションできるFastSAMの使い方紹介です。
モデルのほうで何でも認識してセグメンテーションしてくれるので、ユーザーは事前学習等を考慮する必要がありません。そのため通常SAMは処理時間がネックなのですが、FastSAMはその名の通りとても早くセグメンテーションできます。

早速使ってみる

大体1フレーム当たり100ms～200msくらいでセグメンテーションできていました。
環境は以下です。
Intel(R) Core(TM) i9-14900HX　24コア　2.20 GHz
NVIDIA GeForce RTX 4080 Laptop GPU

　　
きれいにセグメンテーションできてますね。

　　
動画はこちらを使わせていただきました。

インストール

ここは公式サイトそのままに

git clone https://github.com/CASIA-IVA-Lab/FastSAM.git
cd FastSAM
pip install -r requirements.txt
pip install git+https://github.com/ultralytics/CLIP.git

コード

from FastSAM.fastsam import FastSAM, FastSAMPrompt 
import torch
import cv2

class FASTSAMmodl:
    
    def __init__(self, SEGMENTATION):
        self.model = FastSAM(SEGMENTATION["model"])
        self.retina_masks = SEGMENTATION["param"]["retina_masks"]
        self.imgsz = SEGMENTATION["param"]["imgsz"]
        self.conf = SEGMENTATION["param"]["conf"]
        self.iou = SEGMENTATION["param"]["iou"]

    def Segment(self, frame):

        if torch.cuda.is_available():
            device = 'cuda'
        else:
            device = 'cpu'

        everything_results = self.model(frame, device = device, retina_masks = self.retina_masks, imgsz = self.imgsz, conf = self.conf, iou = self.iou)

        prompt_process = FastSAMPrompt(frame, everything_results)
        segment_filter = prompt_process.everything_prompt()

        if(len(segment_filter) > 0):
            result = prompt_process.plot_to_result(annotations = segment_filter)
        else:
            result = frame

        return segment_filter, result
    

par = {
        "model" : "FastSAM/weights/FastSAM-x.pt",
        "param" : {
            "retina_masks" : True,
            "imgsz" : 1000,
            "conf" : 0.7,
            "iou" : 0.9
        }
    }
FS = FASTSAMmodl(par)

filepath = "input.mp4"
video = cv2.VideoCapture(filepath)
framerate = video.get(cv2.CAP_PROP_FPS)

while True:
    ret, image = video.read()
    if not ret:
        break

    segment_filter, segmented_image = FS.Segment(image)

    cv2.imshow("image", segmented_image)
    cv2.waitKey(int(1000.0/framerate))

video.release()
cv2.destroyAllWindows()