From 0a307dcdc1194b520dda06ab127d3ed5be710dcb Mon Sep 17 00:00:00 2001 From: yatengLG <767624851@qq.com> Date: Tue, 25 Apr 2023 13:22:32 +0800 Subject: [PATCH] fix bug --- .idea/ISAT_with_segment_anything.iml | 2 +- .idea/misc.xml | 2 +- README.md | 41 ++++++++++++++-------------- UpdateLog.md | 2 +- segment_any/gpu_resource.py | 16 +++++++++-- segment_any/segment_any.py | 8 +++--- widgets/mainwindow.py | 2 +- 7 files changed, 42 insertions(+), 31 deletions(-) diff --git a/.idea/ISAT_with_segment_anything.iml b/.idea/ISAT_with_segment_anything.iml index e7438cb..a34782b 100644 --- a/.idea/ISAT_with_segment_anything.iml +++ b/.idea/ISAT_with_segment_anything.iml @@ -2,7 +2,7 @@ - + \ No newline at end of file diff --git a/.idea/misc.xml b/.idea/misc.xml index 3f589e5..b47cefa 100644 --- a/.idea/misc.xml +++ b/.idea/misc.xml @@ -1,4 +1,4 @@ - + \ No newline at end of file diff --git a/README.md b/README.md index e185ec0..02aaf39 100644 --- a/README.md +++ b/README.md @@ -21,41 +21,40 @@ ## 安装 ### 1. 源码运行 +#### (1) 创建虚拟环境 ```shell -# 创建虚拟环境 conda create -n ISAT_with_segment_anything python==3.8 conda activate ISAT_with_segment_anything ``` - +#### (2) 安装Segment anything ```shell -# 安装Segment anything git clone git@github.com:facebookresearch/segment-anything.git cd segment-anything pip install -e . cd .. ``` - +#### (3) 安装ISAT_with_segment_anything ```shell -# 安装ISAT_with_segment_anything git clone https://github.com/yatengLG/ISAT_with_segment_anything.git cd ISAT_with_segment_anything pip install -r requirements.txt ``` - -```text -# 下载Segment anything预训练模型 +#### (4) 下载Segment anything预训练模型 下载任一模型，并将模型存放于ISAT_with_segment_anything/segment_any目录下 -模型链接： - https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth - https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth - https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth -h模型最大，效果也最好； -b模型最小，效果也最差； -请按照硬件下载合适的模型，h模型在示例样本上的显存需求约8G. -``` +请按照硬件下载合适的模型. +- H模型:[sam_vit_h_4b8939.pth](https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth) + + 模型最大，效果也最好，显存至少需求8G，演示时软件实际占用7305M； +- L模型:[sam_vit_l_0b3195.pth](https://dl.fbaipublicfiles.com/segment_anything/sam_vit_l_0b3195.pth) + + 模型适中，效果也适中，显存至少需求8G，演示时软件实际占用5855M； +- B模型:[sam_vit_b_01ec64.pth](https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth) + + 模型最小，效果也最差，显存至少需求6G，演示时软件实际占用4149M； + +#### (5) 运行软件 ```shell -# 运行软件 python main.py ``` @@ -69,6 +68,8 @@ python main.py ## 注意事项 1. 自动分割效果受segment anything模型分割效果限制，如需更为精确的分割效果，可通过手动绘制多边形实现。 -2. 如只需要使用手动绘制多边形标注，推荐使用[ISAT](https://github.com/yatengLG/ISAT)。 -3. 如果没有GPU，不建议使用ISAT_with_segment_anything，载入图片花费时间较长。 -4. 如果GPU显存较小，建议使用sam_vit_b_01ec64模型。 +2. 如果没有GPU或只需要使用手动绘制多边形标注，推荐使用[ISAT](https://github.com/yatengLG/ISAT)。 +3. 软件对GPU显存有最低限制： + - h模型最大，效果也最好，显存至少需求8G，演示时软件实际占用7305M； + - l模型适中，效果也适中，显存至少需求8G，演示时软件实际占用5855M； + - b模型最小，效果也最差，显存至少需求6G，演示时软件实际占用4149M； \ No newline at end of file diff --git a/UpdateLog.md b/UpdateLog.md index 316ec62..5f77e30 100644 --- a/UpdateLog.md +++ b/UpdateLog.md @@ -7,4 +7,4 @@ 但最终保存还会以ISAT格式的json保存。 - 添加了显示/隐藏按钮（快捷键V），用于显示或隐藏所有多边形 -- 添加了 +- 添加了GPU显存占用 diff --git a/segment_any/gpu_resource.py b/segment_any/gpu_resource.py index 1dfa950..e6e8c5a 100644 --- a/segment_any/gpu_resource.py +++ b/segment_any/gpu_resource.py @@ -3,6 +3,9 @@ from PyQt5.QtCore import QThread, pyqtSignal import os +import platform + +osplatform = platform.system() class GPUResource_Thread(QThread): @@ -13,9 +16,16 @@ class GPUResource_Thread(QThread): self.gpu_id = None self.callback = None - self.keep = True + if osplatform == 'Windows': + self.command = 'nvidia-smi -q -d MEMORY -i 0 | findstr' + elif osplatform == 'Linux': + self.command = 'nvidia-smi -q -d MEMORY -i 0 | grep' + elif osplatform == 'Darwin': + self.command = 'nvidia-smi -q -d MEMORY -i 0 | grep' + else: + self.command = 'nvidia-smi -q -d MEMORY -i 0 | grep' try: - r = os.popen('nvidia-smi -q -d MEMORY -i 0 | grep Total').readline() + r = os.popen('{} Total'.format(self.command)).readline() self.total = r.split(':')[-1].strip().split(' ')[0] except Exception as e: print(e) @@ -23,7 +33,7 @@ class GPUResource_Thread(QThread): def run(self): while True: - r = os.popen('nvidia-smi -q -d MEMORY -i 0 | grep Used').readline() + r = os.popen('{} Used'.format(self.command)).readline() used = r.split(':')[-1].strip().split(' ')[0] self.message.emit("cuda: {}/{}MiB".format(used, self.total)) diff --git a/segment_any/segment_any.py b/segment_any/segment_any.py index adb4218..bbdffab 100644 --- a/segment_any/segment_any.py +++ b/segment_any/segment_any.py @@ -9,16 +9,16 @@ import numpy as np class SegAny: def __init__(self, checkpoint): if 'vit_b' in checkpoint: - model_type = "vit_b" + self.model_type = "vit_b" elif 'vit_l' in checkpoint: - model_type = "vit_l" + self.model_type = "vit_l" elif 'vit_h' in checkpoint: - model_type = "vit_h" + self.model_type = "vit_h" else: raise ValueError('The checkpoint named {} is not supported.'.format(checkpoint)) self.device = 'cuda' if torch.cuda.is_available() else 'cpu' - sam = sam_model_registry[model_type](checkpoint=checkpoint) + sam = sam_model_registry[self.model_type](checkpoint=checkpoint) sam.to(device=self.device) self.predictor = SamPredictor(sam) self.image = None diff --git a/widgets/mainwindow.py b/widgets/mainwindow.py index b8fc838..e7a0390 100644 --- a/widgets/mainwindow.py +++ b/widgets/mainwindow.py @@ -22,6 +22,7 @@ from PIL import Image import functools import imgviz from segment_any.segment_any import SegAny +from segment_any.gpu_resource import GPUResource_Thread, osplatform import icons_rc @@ -73,7 +74,6 @@ class MainWindow(QtWidgets.QMainWindow, Ui_MainWindow): if self.use_segment_anything: if self.segany.device != 'cpu': - from segment_any.gpu_resource import GPUResource_Thread self.gpu_resource_thread = GPUResource_Thread() self.gpu_resource_thread.message.connect(self.labelGPUResource.setText) self.gpu_resource_thread.start()