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
import os
import numpy as np
import yaml
class TOCOCO(QThread):
message = pyqtSignal(int, int, str)
@ -40,6 +40,14 @@ class TOCOCO(QThread):
coco_anno['categories'] = []
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')]
num_jsons = len(jsons)
@ -157,4 +165,4 @@ class TOCOCO(QThread):
self.message.emit(None, None, '*** Finished! ***')
def __del__(self):
self.wait()
self.wait()