wangtong/ISAT_with_sam
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Update toCOCO.py

Browse Source 
fix bug
This commit is contained in:
yatengLG 2023-06-25 11:00:11 +08:00 committed by GitHub
parent 7e2b31bedf
commit 5678f61bc6
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 10 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
tools/toCOCO.py
View File

@ -5,7 +5,7 @@ from PyQt5.QtCore import QThread, pyqtSignal
from json import load, dump from json import load, dump
import os import os
import numpy as np import numpy as np
import yaml
class TOCOCO(QThread): class TOCOCO(QThread):
message = pyqtSignal(int, int, str) message = pyqtSignal(int, int, str)
@ -40,6 +40,14 @@ class TOCOCO(QThread):
coco_anno['categories'] = [] coco_anno['categories'] = []
categories_dict = {} categories_dict = {}
# categories_dict from isat.yaml, https://github.com/yatengLG/ISAT_with_segment_anything/issues/36
yaml_path = os.path.join(self.isat_json_root, 'isat.yaml')
if os.path.exists(yaml_path):
with open(yaml_path, 'rb')as f:
cfg = yaml.load(f.read(), Loader=yaml.FullLoader)
for index, label_dict in enumerate(cfg.get('label', [])):
label = label_dict.get('name', 'UNKNOW')
categories_dict[label] = index
jsons = [f for f in os.listdir(self.isat_json_root) if f.endswith('.json')] jsons = [f for f in os.listdir(self.isat_json_root) if f.endswith('.json')]
num_jsons = len(jsons) num_jsons = len(jsons)
@ -157,4 +165,4 @@ class TOCOCO(QThread):
self.message.emit(None, None, '*** Finished! ***') self.message.emit(None, None, '*** Finished! ***')
def __del__(self): def __del__(self):
self.wait() self.wait()