- 为视频分割功能添加了进度条

Author: yatengLG

ISAT/segment_any/model_zoo.py

@@ -119,7 +119,7 @@
                 'https://www.modelscope.cn/api/v1/models/yatengLG/ISAT_with_segment_anything_checkpoints/repo?Revision=master&FilePath=checkpoints/sam2_hiera_large.pt'
             ],
             'memory': '4000M',
-            'bf16_memory': '2500M',
+            'bf16_memory': '2800M',
             'params': '900M'
         },
     'sam2_hiera_base_plus.pt':
@@ -128,8 +128,8 @@
                 'https://huggingface.co/yatengLG/ISAT_with_segment_anything_checkpoints/resolve/main/sam2_hiera_base_plus.pt',
                 'https://www.modelscope.cn/api/v1/models/yatengLG/ISAT_with_segment_anything_checkpoints/repo?Revision=master&FilePath=checkpoints/sam2_hiera_base_plus.pt'
             ],
-            'memory': '4000M',
-            'bf16_memory': '2000M',
+            'memory': '2800M',
+            'bf16_memory': '2200M',
             'params': '324M'
         },
     'sam2_hiera_small.pt':
@@ -138,8 +138,8 @@
                 'https://huggingface.co/yatengLG/ISAT_with_segment_anything_checkpoints/resolve/main/sam2_hiera_small.pt',
                 'https://www.modelscope.cn/api/v1/models/yatengLG/ISAT_with_segment_anything_checkpoints/repo?Revision=master&FilePath=checkpoints/sam2_hiera_small.pt'
             ],
-            'memory': '3000M',
-            'bf16_memory': '1500M',
+            'memory': '2500M',
+            'bf16_memory': '1800M',
             'params': '185M'
         },
     'sam2_hiera_tiny.pt':
@@ -148,8 +148,8 @@
                 'https://huggingface.co/yatengLG/ISAT_with_segment_anything_checkpoints/resolve/main/sam2_hiera_tiny.pt',
                 'https://www.modelscope.cn/api/v1/models/yatengLG/ISAT_with_segment_anything_checkpoints/repo?Revision=master&FilePath=checkpoints/sam2_hiera_tiny.pt'
             ],
-            'memory': '2400M',
-            'bf16_memory': '1200M',
+            'memory': '2200M',
+            'bf16_memory': '1500M',
             'params': '156M'
         },
 }

ISAT/widgets/mainwindow.py

@@ -37,6 +37,26 @@
 from skimage.draw.draw import polygon
 
 
+class QtBoxStyleProgressBar(QtWidgets.QProgressBar):
+    # copy from qtbox
+    def __init__(self):
+        super(QtBoxStyleProgressBar, self).__init__()
+        self.setTextVisible(False)
+        self.setStyleSheet("""
+        QProgressBar {
+            border: 2px solid #888783;
+            border-radius: 5px;
+        }
+
+        QProgressBar::chunk {
+            background-color: #74d65f;
+            border-radius: 2px;
+            width: 9px;
+            margin: 0.5px;
+        }
+        """)
+
+
 def calculate_area(points):
     area = 0
     num_points = len(points)
@@ -146,6 +166,11 @@ def run(self):
         print('self.start_frame_idx: ', self.start_frame_idx)
         print('self.max_frame_num_to_track: ', self.max_frame_num_to_track)
 
+        if self.max_frame_num_to_track is not None:
+            total = self.max_frame_num_to_track
+        else:
+            total = len(self.mainwindow.files_list) - self.start_frame_idx + 1
+
         with torch.inference_mode(), torch.autocast(self.mainwindow.segany_video.device,
                                                     dtype=self.mainwindow.segany_video.model_dtype):
 
@@ -194,7 +219,7 @@ def run(self):
                     group_object_dict[group]['mask'] = group_object_dict[group]['mask'] + mask
 
             if len(group_object_dict) < 1:
-                self.tag.emit(0, self.max_frame_num_to_track, True, True, 'Please label objects before video segment.')
+                self.tag.emit(0, total, True, True, 'Please label objects before video segment.')
                 return
             try:
                 for group, object_dict in group_object_dict.items():
@@ -274,12 +299,12 @@ def run(self):
 
                     annotation.objects = objects
                     annotation.save_annotation()
-                    self.tag.emit(index, self.max_frame_num_to_track, False, False, '')
+                    self.tag.emit(index, total, False, False, '')
 
-                self.tag.emit(index, self.max_frame_num_to_track, True, False, '')
+                self.tag.emit(index, total, True, False, '')
 
             except Exception as e:
-                self.tag.emit(index, self.max_frame_num_to_track, True, True, '{}'.format(e))
+                self.tag.emit(index, total, True, True, '{}'.format(e))
 
 
 class InitSegAnyThread(QThread):
@@ -531,6 +556,7 @@ def seg_video_start(self, max_frame_num_to_track=None):
             return
 
         self.setEnabled(False)
+        self.statusbar_change_status(is_message=False)
         self.segany_video_thread.start_frame_idx = self.current_index
         self.segany_video_thread.max_frame_num_to_track=max_frame_num_to_track
         self.segany_video_thread.start()
@@ -540,7 +566,11 @@ def seg_video_finish(self, current, total, finished, is_error, message):
             QtWidgets.QMessageBox.warning(self, 'warning', message)
 
         print('Segment video: {}/{}'.format(current, total))
+        self.progressbar.setMaximum(total)
+        self.progressbar.setValue(current)
         if finished:
+            self.statusbar_change_status(is_message=True)
+            self.progressbar.setValue(0)
             self.setEnabled(True)
 
 
@@ -629,6 +659,12 @@ def init_ui(self):
         """)
         self.statusbar.addPermanentWidget(self.modeState)
 
+        self.progressbar = QtBoxStyleProgressBar()
+        self.progressbar.setTextVisible(False)
+        self.progressbar.setFixedWidth(500)
+        self.progressbar.setVisible(False)
+        self.statusbar.addPermanentWidget(self.progressbar)
+
         self.update_menuSAM()
 
         # mask alpha
@@ -697,6 +733,13 @@ def init_ui(self):
 
         self.trans = QtCore.QTranslator()
 
+    def statusbar_change_status(self, is_message:bool=True):
+        self.labelGPUResource.setVisible(is_message)
+        self.labelData.setVisible(is_message)
+        self.labelCoord.setVisible(is_message)
+        self.modeState.setVisible(is_message)
+        self.progressbar.setVisible(not is_message)
+
     def update_menuSAM(self):
         #
         self.menuSAM_model.clear()

ISAT/widgets/model_manager_dialog.py

@@ -113,7 +113,7 @@ def __init__(self, parent, mainwindow):
         self.download_thread_dict = {}
         self.init_gui()
 
-        self.tableWidget.setColumnWidth(0, 200)
+        self.tableWidget.setColumnWidth(0, 300)
         self.pushButton_clear_tmp.clicked.connect(self.clear_tmp)
         self.checkBox_use_bfloat16.stateChanged.connect(self.use_bfloat16)