file dict ssaved and loaded

This commit is contained in:
Till Raab 2023-11-15 09:22:41 +01:00
parent 33f6bb8683
commit 3f235f02f3

View File

@ -13,6 +13,8 @@ import torchvision.transforms.functional as F
from IPython import embed
import pathlib
NONE_CHECKED_CHECKER_COLORS = ['red', 'black']
class ImageWithBbox(pg.GraphicsLayoutWidget):
# class ImageWithBbox(pg.ImageView):
def __init__(self, img_path, parent=None):
@ -83,10 +85,7 @@ class Bbox_correct_UI(QMainWindow):
super(Bbox_correct_UI, self).__init__(parent)
self.data_path = data_path
files = sorted(list(pathlib.Path(self.data_path).absolute().rglob('*images/*.png')))
# embed()
# quit()
self.files = sorted(list(pathlib.Path(self.data_path).absolute().rglob('*images/*.png')))
rec = QApplication.desktop().screenGeometry()
height = rec.height()
@ -102,7 +101,7 @@ class Bbox_correct_UI(QMainWindow):
self.setCentralWidget(self.central_widget)
# image widget
self.current_img = ImageWithBbox(files[0], parent=self)
self.current_img = ImageWithBbox(self.files[0], parent=self)
self.central_Layout.addWidget(self.current_img, 4)
# image select widget
self.scroll = QScrollArea() # Scroll Area which contains the widgets, set as the centralWidget
@ -111,13 +110,8 @@ class Bbox_correct_UI(QMainWindow):
self.highlighted_label = None
self.all_labels = []
# embed()
# quit()
self.file_dict = pd.DataFrame(
{'name': [f.name for f in files],
'files': files,
'text_color': ['red' for i in range(len(files))]} # change this to locked
)
self.load_or_create_file_dict()
for i in range(len(self.file_dict)):
label = QLabel(f'{self.file_dict["name"][i]}')
@ -129,11 +123,11 @@ class Bbox_correct_UI(QMainWindow):
if i == 0:
label.setStyleSheet("border: 2px solid black; "
"color : %s;" % (self.file_dict['text_color'][i]))
"color : %s;" % (NONE_CHECKED_CHECKER_COLORS[self.file_dict['checked'][i]]))
self.highlighted_label = label
else:
label.setStyleSheet("border: 1px solid gray; "
"color : %s;" % (self.file_dict['text_color'][i]))
"color : %s;" % (NONE_CHECKED_CHECKER_COLORS[self.file_dict['checked'][i]]))
self.vbox.addWidget(label)
self.all_labels.append(label)
@ -148,21 +142,57 @@ class Bbox_correct_UI(QMainWindow):
self.add_actions()
def load_or_create_file_dict(self):
csvs_im_data_path = list(pathlib.Path(self.data_path).absolute().rglob('*file_dict.csv'))
if len(csvs_im_data_path) == 0:
self.file_dict = pd.DataFrame(
{'name': [f.name for f in self.files],
'files': self.files,
'checked': np.zeros(len(self.files), dtype=int)} # change this to locked
)
else:
self.file_dict = pd.read_csv(csvs_im_data_path[0], sep=',')
def save_file_dict(self):
self.file_dict.to_csv(pathlib.Path(self.data_path)/'file_dict.csv', sep=',', index=False)
def label_clicked(self, clicked_label):
if self.highlighted_label:
hl_mask = self.file_dict['name'] == self.highlighted_label.text()
hl_label_name, _, hl_label_text_color = self.file_dict[hl_mask].values[0]
hl_label_name, _, hl_checked = self.file_dict[hl_mask].values[0]
self.highlighted_label.setStyleSheet("border: 1px solid gray; "
"color: %s;" % (hl_label_text_color))
"color: %s;" % (NONE_CHECKED_CHECKER_COLORS[hl_checked]))
mask = self.file_dict['name'] == clicked_label.text()
new_name, new_file, new_text_color = self.file_dict[mask].values[0]
new_name, new_file, new_checked = self.file_dict[mask].values[0]
clicked_label.setStyleSheet("border: 2px solid black;"
"color: %s;" % (new_text_color))
"color: %s;" % (NONE_CHECKED_CHECKER_COLORS[new_checked]))
self.highlighted_label = clicked_label
self.switch_to_new_file(new_file, new_name)
def lock_file(self):
hl_mask = self.file_dict['name'] == self.highlighted_label.text()
hl_label_name, _, hl_checked = self.file_dict[hl_mask].values[0]
# ToDo: do everything with the index instead of mask
df_idx = self.file_dict.loc[self.file_dict['name'] == self.highlighted_label.text()].index.values[0]
self.file_dict.at[df_idx, 'checked'] = 1
self.highlighted_label.setStyleSheet("border: 1px solid gray; "
"color: 'black';")
new_idx = df_idx + 1 if df_idx < len(self.file_dict)-1 else 0
new_name, new_file, new_checked = self.file_dict.iloc[new_idx].values
self.all_labels[new_idx].setStyleSheet("border: 2px solid black;"
"color: %s;" % (NONE_CHECKED_CHECKER_COLORS[new_checked]))
self.highlighted_label = self.all_labels[new_idx]
self.switch_to_new_file(new_file, new_name)
def switch_to_new_file(self, new_file, new_name):
self.readout_rois()
@ -189,38 +219,18 @@ class Bbox_correct_UI(QMainWindow):
new_labels = np.array(new_labels)
np.savetxt(self.current_img.label_path, new_labels)
def lock_file(self):
hl_mask = self.file_dict['name'] == self.highlighted_label.text()
hl_label_name, _, hl_label_text_color = self.file_dict[hl_mask].values[0]
# ToDo: do everything with the index instead of mask
df_idx = self.file_dict.loc[self.file_dict['name'] == self.highlighted_label.text()].index.values[0]
self.file_dict.iloc[df_idx]['text_color'] = 'black'
self.highlighted_label.setStyleSheet("border: 1px solid gray; "
"color: 'black';")
new_idx = df_idx + 1 if df_idx < len(self.file_dict)-1 else 0
new_name, new_file, new_color = self.file_dict.iloc[new_idx].values
self.all_labels[new_idx].setStyleSheet("border: 2px solid black;"
"color: %s;" % (new_color))
self.highlighted_label = self.all_labels[new_idx]
self.switch_to_new_file(new_file, new_name)
# embed()
# quit()
# go to next
pass
def add_actions(self):
self.lock_file_act = QAction('loc', self)
self.lock_file_act.triggered.connect(self.lock_file)
self.lock_file_act.setShortcut(Qt.Key_Space)
self.addAction(self.lock_file_act)
def closeEvent(self, *args, **kwargs):
super(QMainWindow, self).closeEvent(*args, **kwargs)
self.readout_rois()
self.save_file_dict()
print ("you just closed the pyqt window!!! you are awesome!!!")
def main_UI():
app = QApplication(sys.argv) # create application
data_path = sys.argv[1]