From 02bdfe663c6b7d254332d24a6bac80b5b367d7a5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Cl=C3=A9ment=20H=2E=20Benedetti?= Date: Mon, 30 Sep 2024 18:28:58 +0200 Subject: [PATCH] Added boxes counter --- .../_widget_yolo_annotations.py | 114 ++++++++++++++---- 1 file changed, 90 insertions(+), 24 deletions(-) diff --git a/src/microglia_analyzer/_widget_yolo_annotations.py b/src/microglia_analyzer/_widget_yolo_annotations.py index aa8dc35..75e35cf 100644 --- a/src/microglia_analyzer/_widget_yolo_annotations.py +++ b/src/microglia_analyzer/_widget_yolo_annotations.py @@ -49,7 +49,11 @@ def __init__(self, napari_viewer): # Active Napari viewer. self.viewer = napari_viewer # Folder in which are located the 'images' and 'annotations' folders. + self.root_directory = None + # Name of the folder in which images are located self.sources_directory = None + # Name of the folder in which annotations are located + self.annotations_directory = None # List of images ('.tif') in the 'images' folder. self.images_list = [] @@ -64,10 +68,16 @@ def add_media_management_group_ui(self): layout = QVBoxLayout() box.setLayout(layout) + # Label + button to select the source directory: + self.select_sources_directory_button = QPushButton("📂 Sources directory") + self.select_sources_directory_button.clicked.connect(self.select_sources_directory) + layout.addWidget(self.select_sources_directory_button) + # Label + text box for the inputs sub-folder's name: inputs_name_label = QLabel("Inputs sub-folder:") - self.inputs_name = QLineEdit() - self.inputs_name.setText("inputs") + self.inputs_name = QComboBox() + self.inputs_name.currentIndexChanged.connect(self.set_sources_directory) + self.inputs_name.addItem("---") h_layout = QHBoxLayout() h_layout.addWidget(inputs_name_label) h_layout.addWidget(self.inputs_name) @@ -75,18 +85,17 @@ def add_media_management_group_ui(self): # Label + text box for the annotations sub-folder's name: annotations_name_label = QLabel("Annotations sub-folder:") - self.annotations_name = QLineEdit() - self.annotations_name.setText("labels") + self.annotations_name = QLabel() + self.annotations_name.setText("---") + self.annotations_name.setMaximumHeight(20) + font = self.annotations_name.font() + font.setBold(True) + self.annotations_name.setFont(font) h_layout = QHBoxLayout() h_layout.addWidget(annotations_name_label) h_layout.addWidget(self.annotations_name) layout.addLayout(h_layout) - # Label + button to select the source directory: - self.select_sources_directory_button = QPushButton("📂 Sources directory") - self.select_sources_directory_button.clicked.connect(self.select_sources_directory) - layout.addWidget(self.select_sources_directory_button) - self.layout.addWidget(box) def add_classes_management_group_ui(self): @@ -103,6 +112,17 @@ def add_classes_management_group_ui(self): h_laytout.addWidget(self.add_yolo_class_button) layout.addLayout(h_laytout) + # Label showing the number of boxes in each class + self.counts_label = QLabel("Counts:") + font = self.counts_label.font() + font.setBold(True) + self.counts_label.setFont(font) + layout.addWidget(self.counts_label) + self.count_display_label = QLabel("") + self.counts_label.setMaximumHeight(20) + self.count_display_label.setMaximumHeight(10) + layout.addWidget(self.count_display_label) + self.layout.addWidget(box) def add_annotations_management_group_ui(self): @@ -134,13 +154,11 @@ def select_sources_directory(self): """ Select the folder containing the "images" and "labels" sub-folders. """ - directory = QFileDialog.getExistingDirectory(self, "Select sources directory") + directory = QFileDialog.getExistingDirectory(self, "Select root directory") if not os.path.isdir(directory): show_info("Invalid directory.") return - if not self.set_sources_directory(directory): - show_info("No input directory found.") - return + self.set_root_directory(directory) def upper_corner(self, box): """ @@ -203,13 +221,13 @@ def layer2yolo(self, layer_name, index): return tuples def write_annotations(self, tuples): - labels_folder = os.path.join(self.sources_directory, self.annotations_name.text()) + labels_folder = os.path.join(self.root_directory, self.annotations_directory) current_as_txt = TIFF_REGEX.match(self.image_selector.currentText()).group(1) + ".txt" labels_path = os.path.join(labels_folder, current_as_txt) with open(labels_path, "w") as f: for row in tuples: f.write(" ".join(map(str, row)) + "\n") - with open(os.path.join(self.sources_directory, "classes.txt"), "w") as f: + with open(os.path.join(self.root_directory, "classes.txt"), "w") as f: for c in self.get_classes(): f.write(c + "\n") show_info("Annotations saved.") @@ -223,6 +241,7 @@ def save_state(self): lines += self.layer2yolo(l.name, count) count += 1 self.write_annotations(lines) + self.count_boxes() def get_new_class_name(self): """ @@ -259,15 +278,32 @@ def add_yolo_class(self): opacity=0.8, edge_width=3 ) - - def set_sources_directory(self, directory): - inputs_path = os.path.join(directory, self.inputs_name.text()) - annotations_path = os.path.join(directory, self.annotations_name.text()) + + def set_root_directory(self, directory): + folders = sorted([f for f in os.listdir(directory) if os.path.isdir(os.path.join(directory, f))]) + folders = ["---"] + folders + self.inputs_name.clear() + self.inputs_name.addItems(folders) + self.root_directory = directory + + def set_sources_directory(self): + source_folder = self.inputs_name.currentText() + annotations_folder = source_folder + "-labels" + if (source_folder is None) or (source_folder == "---") or (source_folder == ""): + return + inputs_path = os.path.join(self.root_directory, source_folder) + annotations_path = os.path.join(self.root_directory, annotations_folder) if not os.path.isdir(inputs_path): return False if not os.path.isdir(annotations_path): os.makedirs(annotations_path) - self.sources_directory = directory + self.sources_directory = source_folder + self.annotations_directory = annotations_folder + self.annotations_name.setText(annotations_folder) + self.open_sources_directory() + + def open_sources_directory(self): + inputs_path = os.path.join(self.root_directory, self.sources_directory) self.images_list = sorted([f for f in os.listdir(inputs_path) if TIFF_REGEX.match(f) is not None]) if len(self.images_list) == 0: # Didn't find any file in the folder. show_info("Didn't find any TIFF file in the provided folder.") @@ -290,7 +326,7 @@ def clear_classes_layers(self): self.viewer.layers[n].data = [] def restore_classes_layers(self): - classes_path = os.path.join(self.sources_directory, "classes.txt") + classes_path = os.path.join(self.root_directory, "classes.txt") if not os.path.isfile(classes_path): show_info("No classes file found.") return @@ -388,11 +424,11 @@ def open_image(self): """ current_image = self.image_selector.currentText() # Check that the name is valid. - if (self.sources_directory is None) or (current_image is None) or (current_image == "---") or (current_image == ""): + if (self.root_directory is None) or (current_image is None) or (current_image == "---") or (current_image == ""): return - image_path = os.path.join(self.sources_directory, self.inputs_name.text(), current_image) + image_path = os.path.join(self.root_directory, self.sources_directory, current_image) current_as_txt = TIFF_REGEX.match(current_image).group(1) + ".txt" # Remove extension + adds ".txt" extension. - labels_path = os.path.join(self.sources_directory, self.annotations_name.text(), current_as_txt) + labels_path = os.path.join(self.root_directory, self.annotations_directory, current_as_txt) if not os.path.isfile(image_path): print(f"The image: '{current_image}' doesn't exist.") return @@ -406,3 +442,33 @@ def open_image(self): self.clear_classes_layers() if os.path.isfile(labels_path): # If some annotations already exist for this image. self.load_annotations(labels_path) + self.count_boxes() + + def count_boxes(self): + """ + Counts the number of boxes in each class to make sure annotations are balanced. + """ + classes = self.get_classes() + annotations_path = os.path.join(self.root_directory, self.annotations_directory) + counts = dict() + for f in os.listdir(annotations_path): + if not f.endswith(".txt"): + continue + with open(os.path.join(annotations_path, f), "r") as file: + lines = file.read().split('\n') + for line in lines: + if line == "": + continue + c, _, _, _, _ = line.split(" ") + c = int(c) + counts[c] = counts.get(c, 0) + 1 + + text = "" + for class_idx, class_name in enumerate(classes): + if class_idx not in counts: + text += f'{class_name}: 0
' + else: + text += f'{classes[class_idx]}: {counts[class_idx]}
' + + self.count_display_label.setText(text) + self.count_display_label.setMaximumHeight(50 * len(counts)) \ No newline at end of file