基于FastSAM模型的realsense实时图像分割

mfrobot

1202人浏览 · 2024-12-10 19:57:49

mfrobot · 2024-12-10 19:57:49 发布

前言

在上一小节介绍了海康MVS相机如何与SAM图像分割模型打通接口之后，本小节将记录realsense相机与FastSAM的结合，在保证图像分割精度的基础上提升分割的实时性。

一、SAM与FastSAM

Segment Anything Model （SAM）被提出。它被认为是一个里程碑式的愿景基础模型。它可以在各种可能的用户交互提示的引导下分割图像中的任何对象。SAM 利用在广泛的 SA-1B 数据集上训练的 Transformer 模型，使其能够巧妙地处理各种场景和对象。SAM 为一项激动人心的新任务打开了大门，称为 Segment Anything。由于其可推广性和潜力，这项任务具有成为未来广泛视觉任务的基石的所有条件。然而，尽管取得了这些进步，并且SAM和后续模型在处理细分市场任何任务方面都显示出可喜的结果，但其实际应用仍然具有挑战性。最突出的问题是与 Transformer （ViT）模型相关的大量计算资源需求，而 Transformer 模型是 SAM 架构的主要部分。与卷积对应物相比，ViT因其繁重的计算资源需求而脱颖而出，这给它们的实际部署带来了障碍，尤其是在实时应用中。因此，这种限制阻碍了细分任何任务的进展和潜力。
FastSAM为分割一切设计了一个实时解决方案。我们将任务细分为两个顺序阶段，即全实例分割和提示引导选择。第一阶段取决于基于卷积神经网络（CNN）的检测器的实现，它生成图像中所有实例的分割掩码。然后在第二阶段，它输出与提示相对应的感兴趣区域。通过利用CNN的计算效率，该方法证明了任何模型的实时分割都是可以实现的，而不会对性能质量产生太大影响。希望所提出的方法能够促进分割任何东西的基础任务的工业应用。

二、环境配置

1.FastSAM环境配置

FastSAM的环境配置可以参考其github开源地址：
https://github.com/CASIA-IVA-Lab/FastSAM
基本上，只要完整配置过一遍SAM环境，FastSAM的环境配置流程应该就非常熟悉了，这里不再赘述。

2.realsense环境配置

开发环境在pycharm PYQT5下进行，因此在pycharm终端输入命令：

pip install pyrealsense2

即可成功配置

三.代码移植

realsense采集图像的部分，参考如下博客：
《英特尔Realsense学习笔记二: pyqt5 实时显示 Realsense D415 深度图像和彩色图像并支持按钮保存图像》
https://blog.csdn.net/qq_24815615/article/details/112986401
这里我们只保留彩色图像color_image。
代码的核心部分----图像保存，因为当realsense相机采集到图像之后需要先保存到指定路径，然后FastSAM端再调取分割，并最终显示在UI上，代码如下：

now_date = dt.datetime.now().strftime('%F')
                    path_ok = os.path.exists(now_date)
                    if (path_ok == False):
                        os.mkdir(now_date)

                    if (os.path.isdir(now_date)):
                        id = self.lineEdit_id.text()

                        if (re.match('^[\u4E00-\u9FA5a-zA-Z0-9_]*$', id) and (id != '')):

                            color_full_path = os.path.join('./', now_date, id + '_color.png')
                            print(color_full_path)
                        else:
                            color_full_path = os.path.join('./', now_date + '_color.png')
                            print(color_full_path)

                        cv2.imencode('.png', color_image)[1].tofile(color_full_path)

然后，FastSAM端通过获取该图像路径，进行分割。分割的种类包括全景分割，点分割，文本分割等，可以根据需求进行选择

#SAM分割部分
        IMAGE_PATH = color_full_path
        everything_results = model(IMAGE_PATH, device=DEVICE, retina_masks=True, imgsz=1024, conf=0.4,
                                   iou=0.9, )
        prompt_process = FastSAMPrompt(IMAGE_PATH, everything_results, device=DEVICE)

        # point propmt
        # ann = prompt_process.point_prompt(points=[[420, 350]], pointlabel=[1])
        # everything prompt
        ann = prompt_process.everything_prompt()
        prompt_process.plot(annotations=ann,
                            output_path='/home/yyc/SAM_series/FastSAM-main/output/seg.jpg')
        image = cv2.imread('/home/yyc/SAM_series/FastSAM-main/output/seg.jpg')
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        qimage2 = QImage(image, 1280, 720, QImage.Format_BGR888)
        pixmap2 = QPixmap.fromImage(qimage2)

最终，点击按钮采集并保存图像，再点击分割按钮，即可在UI上显示分割的结果。根据终端显示的信息以及直观感受，从采集图像到显示分割结果整个过程只用了3-4秒，和原始SAM模型相比实时性有了非常大的提升。

在这里插入图片描述
完整代码如下：

import sys
from PyQt5.QtWidgets import QApplication, QMainWindow, QMessageBox
from PyQt5.QtGui import QImage, QPixmap
from PyQt5.QtCore import QObject, pyqtSignal
from untitled import Ui_MainWindow

from fastsam import FastSAM, FastSAMPrompt
import os

import time
import matplotlib.pyplot as plt
import datetime as dt
import numpy as np
import threading as th
import ctypes
import inspect

import re

# First import the library
import pyrealsense2 as rs
import cv2
from skimage import io

envpath = '换成自己的路径'
os.environ['QT_QPA_PLATFORM_PLUGIN_PATH'] = envpath


DELAY = 0

model = FastSAM('换成自己的路径')
DEVICE = 'cpu'


color_full_path = '换成自己的路径'

class MainWindow(QMainWindow, Ui_MainWindow):

    def __init__(self):
        super(MainWindow, self).__init__()

        # Set up the user interface from Designer.
        self.setupUi(self)

        self.dis_update.connect(self.camera_view)
        self.pushButton_takephotos.clicked.connect(self.pushButton_takephotos_clicked)
        self.pushButton_segment.clicked.connect(self.seg)


        self.thread_camera = None
        self.takePhotos = False

    # 在对应的页面类的内部，与def定义的函数同级
    dis_update = pyqtSignal(QPixmap)

    def pushButton_takephotos_clicked(self):
        self.takePhotos = True

    # 添加一个退出的提示事件
    def closeEvent(self, event):
        """我们创建了一个消息框，上面有俩按钮：Yes和No.第一个字符串显示在消息框的标题栏，第二个字符串显示在对话框，
              第三个参数是消息框的俩按钮，最后一个参数是默认按钮，这个按钮是默认选中的。返回值在变量reply里。"""

        reply = QMessageBox.question(self, 'Message', "Are you sure to quit?",
                                     QMessageBox.Yes | QMessageBox.No, QMessageBox.No)
        # 判断返回值，如果点击的是Yes按钮，我们就关闭组件和应用，否则就忽略关闭事件
        if reply == QMessageBox.Yes:
            self.stop_thread(self.thread_camera)
            event.accept()
        else:
            event.ignore()

    def open_camera(self):
        # target选择开启摄像头的函数
        self.thread_camera = th.Thread(target=self.open_realsense)
        self.thread_camera.start()
        print('Open Camera')

    def camera_view(self, c):
        # 调用setPixmap函数设置显示Pixmap
        self.label_show.setPixmap(c)
        # 调用setScaledContents将图像比例化显示在QLabel上
        self.label_show.setScaledContents(True)
        # 调用setPixmap函数设置显示Pixmap
        # self.label_seg.setPixmap(c)
        # # 调用setScaledContents将图像比例化显示在QLabel上
        # self.label_seg.setScaledContents(True)

    def show_anns(anns):
        if len(anns) == 0:
            return
        sorted_anns = sorted(anns, key=(lambda x: x['area']), reverse=True)
        ax = plt.gca()
        ax.set_autoscale_on(False)

        img = np.ones((sorted_anns[0]['segmentation'].shape[0], sorted_anns[0]['segmentation'].shape[1], 4))
        img[:, :, 3] = 0
        for ann in sorted_anns:
            m = ann['segmentation']
            color_mask = np.concatenate([np.random.random(3), [0.35]])
            img[m] = color_mask
        ax.imshow(img)

    # def camera_seg(self, seg):


    def _async_raise(self, tid, exctype):
        """raises the exception, performs cleanup if needed"""
        tid = ctypes.c_long(tid)
        if not inspect.isclass(exctype):
            exctype = type(exctype)
        res = ctypes.pythonapi.PyThreadState_SetAsyncExc(tid, ctypes.py_object(exctype))
        if res == 0:
            raise ValueError("invalid thread id")
        elif res != 1:
            # """if it returns a number greater than one, you're in trouble,
            # and you should call it again with exc=NULL to revert the effect"""
            ctypes.pythonapi.PyThreadState_SetAsyncExc(tid, None)
            raise SystemError("PyThreadState_SetAsyncExc failed")

    def stop_thread(self, thread):
        self._async_raise(thread.ident, SystemExit)

    def open_realsense(self):
        print('open_realsense')

        # Create a pipeline
        pipeline = rs.pipeline()

        # Create a config and configure the pipeline to stream
        #  different resolutions of color and depth streams
        config = rs.config()
        config.enable_stream(rs.stream.depth, 1280, 720, rs.format.z16, 30)
        config.enable_stream(rs.stream.color, 1280, 720, rs.format.bgr8, 30)

        # Start streaming
        profile = pipeline.start(config)

        # Getting the depth sensor's depth scale (see rs-align example for explanation)
        depth_sensor = profile.get_device().first_depth_sensor()
        depth_scale = depth_sensor.get_depth_scale()
        print("Depth Scale is: ", depth_scale)

        # We will be removing the background of objects more than
        #  clipping_distance_in_meters meters away
        clipping_distance_in_meters = 100000  # 1 meter
        clipping_distance = clipping_distance_in_meters / depth_scale

        # Color Intrinsics
        # intr = color_frame.profile.as_video_stream_profile().intrinsics

        # Create an align object
        # rs.align allows us to perform alignment of depth frames to others frames
        # The "align_to" is the stream type to which we plan to align depth frames.
        align_to = rs.stream.color
        align = rs.align(align_to)

        # Streaming loop
        try:
            while True:
                # Get frameset of color and depth
                frames = pipeline.wait_for_frames()
                # frames.get_depth_frame() is a 640x360 depth image

                # Align the depth frame to color frame
                aligned_frames = align.process(frames)

                # Get aligned frames
                aligned_depth_frame = aligned_frames.get_depth_frame()  # aligned_depth_frame is a 640x480 depth image
                color_frame = aligned_frames.get_color_frame()

                # Validate that both frames are valid
                if not aligned_depth_frame or not color_frame:
                    continue

                depth_image = np.asanyarray(aligned_depth_frame.get_data())
                color_image = np.asanyarray(color_frame.get_data())

                if (self.takePhotos == True):

                    now_date = dt.datetime.now().strftime('%F')
                    path_ok = os.path.exists(now_date)
                    if (path_ok == False):
                        os.mkdir(now_date)

                    if (os.path.isdir(now_date)):
                        id = self.lineEdit_id.text()

                        if (re.match('^[\u4E00-\u9FA5a-zA-Z0-9_]*$', id) and (id != '')):

                            color_full_path = os.path.join('./', now_date, id + '_color.png')
                            print(color_full_path)
                        else:
                            color_full_path = os.path.join('./', now_date + '_color.png')
                            print(color_full_path)

                        cv2.imencode('.png', color_image)[1].tofile(color_full_path)
                    self.takePhotos = False

                # 显示图像部分
                qimage = QImage(color_image, 1280, 720, QImage.Format_BGR888)
                pixmap = QPixmap.fromImage(qimage)
                self.dis_update.emit(pixmap)

                time.sleep(DELAY)
        finally:
            pipeline.stop()

    def seg(self):
        now_date = dt.datetime.now().strftime('%F')

        path_ok = os.path.exists(now_date)
        if (path_ok == False):
            os.mkdir(now_date)

        if (os.path.isdir(now_date)):
            id = self.lineEdit_id.text()

            if (re.match('^[\u4E00-\u9FA5a-zA-Z0-9_]*$', id) and (id != '')):

                color_full_path = os.path.join('./', now_date, id + '_color.png')
                # print(color_full_path)
            else:
                color_full_path = os.path.join('./', now_date + '_color.png')
                # print(color_full_path)
        #SAM分割部分
        IMAGE_PATH = color_full_path
        everything_results = model(IMAGE_PATH, device=DEVICE, retina_masks=True, imgsz=1024, conf=0.4,
                                   iou=0.9, )
        prompt_process = FastSAMPrompt(IMAGE_PATH, everything_results, device=DEVICE)

        # point propmt
        # ann = prompt_process.point_prompt(points=[[420, 350]], pointlabel=[1])
        # everything prompt
        ann = prompt_process.everything_prompt()
        prompt_process.plot(annotations=ann,
                            output_path='换成自己的路径')
        image = cv2.imread('换成自己的路径')
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        qimage2 = QImage(image, 1280, 720, QImage.Format_BGR888)
        pixmap2 = QPixmap.fromImage(qimage2)
        self.dis_update.emit(pixmap2)
        self.label_seg.setPixmap(pixmap2)
        # 调用setScaledContents将图像比例化显示在QLabel上
        self.label_seg.setScaledContents(True)


if __name__ == "__main__":
    app = QApplication(sys.argv)
    w = MainWindow()
    w.show()

    w.open_camera()

    # thread_camera = th.Thread(target=w.open_realsense)
    # thread_camera.start()

    print('Hello World!')

    sys.exit(app.exec_())

四.总结

以上就是本小节的全部实验。工程难度的确不算大，主要是一个优化分割实时性的过程，对于后期应用到解决一些具体项目或者实际问题上肯定也会有很大的帮助，继续加油干！

参考链接：
1.英特尔Realsense学习笔记二: pyqt5 实时显示 Realsense D415 深度图像和彩色图像并支持按钮保存图像
https://blog.csdn.net/qq_24815615/article/details/112986401
2.中科院版「分割一切」Fast SAM模型 | 精度相当，速度提升50倍！目前GitHub6.9K+星
https://blog.csdn.net/Java_college/article/details/138350227
3.FastSAM分割一切
https://blog.csdn.net/qq_30061245/article/details/140147454