fix bug

.idea/ISAT_with_segment_anything.iml

@@ -2,7 +2,7 @@
 <module type="PYTHON_MODULE" version="4">
   <component name="NewModuleRootManager">
     <content url="file://$MODULE_DIR$" />
-    <orderEntry type="jdk" jdkName="Python 3.8 (ISAT_with_segment_anything)" jdkType="Python SDK" />
+    <orderEntry type="jdk" jdkName="Python 3.6 (segment_anything_env)" jdkType="Python SDK" />
     <orderEntry type="sourceFolder" forTests="false" />
   </component>
 </module>

.idea/misc.xml

@@ -1,4 +1,4 @@
 <?xml version="1.0" encoding="UTF-8"?>
 <project version="4">
-  <component name="ProjectRootManager" version="2" project-jdk-name="Python 3.8 (ISAT_with_segment_anything)" project-jdk-type="Python SDK" />
+  <component name="ProjectRootManager" version="2" project-jdk-name="Python 3.6 (segment_anything_env)" project-jdk-type="Python SDK" />
 </project>

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；

UpdateLog.md

@@ -7,4 +7,4 @@
     但最终保存还会以ISAT格式的json保存。
 
 - 添加了显示/隐藏按钮（快捷键V），用于显示或隐藏所有多边形
-- 添加了
+- 添加了GPU显存占用

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))
 

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

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()