[Feature] Add support for displaying and exporting shape-level inform…

…ation

anylabeling/views/labeling/widgets/overview_dialog.py

@@ -26,6 +26,7 @@ def __init__(self, parent):
         self.image_file_list = self.get_image_file_list()
         self.start_index = 1
         self.end_index = len(self.image_file_list)
+        self.showing_label_infos = True
         if self.image_file_list:
             self.init_ui()
 
@@ -58,7 +59,7 @@ def init_ui(self):
         self.from_input.setMinimum(1)
         self.from_input.setMaximum(len(self.image_file_list))
         self.from_input.setSingleStep(1)
-        self.from_input.setValue(self.start_index) 
+        self.from_input.setValue(self.start_index)
         range_layout.addWidget(from_label)
         range_layout.addWidget(self.from_input)
 
@@ -67,7 +68,7 @@ def init_ui(self):
         self.to_input.setMinimum(1)
         self.to_input.setMaximum(len(self.image_file_list))
         self.to_input.setSingleStep(1)
-        self.to_input.setValue(len(self.image_file_list)) 
+        self.to_input.setValue(len(self.image_file_list))
         range_layout.addWidget(to_label)
         range_layout.addWidget(self.to_input)
 
@@ -81,7 +82,12 @@ def init_ui(self):
         # Add export button for exporting data
         self.export_button = QPushButton(self.tr("Export"))
 
+        # Add toggle button to switch between label_infos and shape_infos
+        self.toggle_button = QPushButton(self.tr("Show Shape Infos"))
+        self.toggle_button.clicked.connect(self.toggle_info)
+
         range_and_export_layout = QHBoxLayout()
+        range_and_export_layout.addWidget(self.toggle_button, 0, Qt.AlignLeft)
         range_and_export_layout.addStretch(1)
         range_and_export_layout.addLayout(range_layout)
         range_and_export_layout.addStretch(1)
@@ -110,6 +116,7 @@ def get_image_file_list(self):
     def get_label_infos(self, start_index: int = -1, end_index: int = -1):
         initial_nums = [0 for _ in range(len(self.supported_shape))]
         label_infos = {}
+        shape_infos = []
 
         progress_dialog = QProgressDialog(
             self.tr("Loading..."),
@@ -145,6 +152,7 @@ def get_label_infos(self, start_index: int = -1, end_index: int = -1):
                 continue
             with open(label_file, "r", encoding="utf-8") as f:
                 data = json.load(f)
+            filename = data["imagePath"]
             shapes = data.get("shapes", [])
             for shape in shapes:
                 if "label" not in shape or "shape_type" not in shape:
@@ -154,18 +162,38 @@ def get_label_infos(self, start_index: int = -1, end_index: int = -1):
                     print(f"Invalid shape_type {shape_type} of {label_file}!")
                     continue
                 label = shape["label"]
+                score = shape.get("score", 0.0)
+                flags = shape.get("flags", {})
+                points = shape.get("points", [])
+                group_id = shape.get("group_id", -1)
+                difficult = shape.get("difficult", False)
+                description = shape.get("description", "")
+                kie_linking = shape.get("kie_linking", [])
                 if label not in label_infos:
                     label_infos[label] = dict(zip(self.supported_shape, initial_nums))
                 label_infos[label][shape_type] += 1
+                current_shape = dict(
+                    filename=filename,
+                    label=label,
+                    score=score,
+                    flags=flags,
+                    points=points,
+                    group_id=group_id,
+                    difficult=difficult,
+                    shape_type=shape_type,
+                    description=description,
+                    kie_linking=kie_linking,
+                )
+                shape_infos.append(current_shape)
             progress_dialog.setValue(i)
             if progress_dialog.wasCanceled():
                 break
         progress_dialog.close()
         label_infos = {k: label_infos[k] for k in sorted(label_infos)}
-        return label_infos
+        return label_infos, shape_infos
 
     def get_total_infos(self, start_index: int = -1, end_index: int = -1):
-        label_infos = self.get_label_infos(start_index, end_index)
+        label_infos, _ = self.get_label_infos(start_index, end_index)
         total_infos = [["Label"] + self.supported_shape + ["Total"]]
         shape_counter = [0 for _ in range(len(self.supported_shape) + 1)]
 
@@ -181,20 +209,62 @@ def get_total_infos(self, start_index: int = -1, end_index: int = -1):
         total_infos.append(["Total"] + shape_counter)
         return total_infos
 
-    def populate_table(self, start_index: int = -1, end_index: int = -1):
-        total_infos = self.get_total_infos(start_index, end_index)
-        rows = len(total_infos) - 1
-        cols = len(total_infos[0])
-        self.table.setRowCount(rows)
-        self.table.setColumnCount(cols)
-        self.table.setHorizontalHeaderLabels(total_infos[0])
-
-        data = [list(map(str, info)) for info in total_infos[1:]]
+    def get_shape_infos_table(self, shape_infos):
+        headers = [
+            "Filename",
+            "Label",
+            "Type",
+            "Linking",
+            "Group ID",
+            "Difficult",
+            "Description",
+            "Flags",
+            "Points",
+        ]
+        table_data = []
+        for shape in shape_infos:
+            row = [
+                shape["filename"],
+                shape["label"],
+                shape["shape_type"],
+                str(shape["kie_linking"]),
+                str(shape["group_id"]),
+                str(shape["difficult"]),
+                shape["description"],
+                str(shape["flags"]),
+                str(shape["points"]),              
+            ]
+            table_data.append(row)
+        return headers, table_data
 
-        for row, info in enumerate(data):
-            for col, value in enumerate(info):
-                item = QTableWidgetItem(value)
-                self.table.setItem(row, col, item)
+    def populate_table(self, start_index: int = -1, end_index: int = -1):
+        if self.showing_label_infos:
+            total_infos = self.get_total_infos(start_index, end_index)
+            rows = len(total_infos) - 1
+            cols = len(total_infos[0])
+            self.table.setRowCount(rows)
+            self.table.setColumnCount(cols)
+            self.table.setHorizontalHeaderLabels(total_infos[0])
+
+            data = [list(map(str, info)) for info in total_infos[1:]]
+
+            for row, info in enumerate(data):
+                for col, value in enumerate(info):
+                    item = QTableWidgetItem(value)
+                    self.table.setItem(row, col, item)
+        else:
+            _, shape_infos = self.get_label_infos(start_index, end_index)
+            headers, table_data = self.get_shape_infos_table(shape_infos)
+            self.table.setRowCount(len(table_data))
+            self.table.setColumnCount(len(headers))
+            self.table.setHorizontalHeaderLabels(headers)
+
+            for row, data in enumerate(table_data):
+                for col, value in enumerate(data):
+                    item = QTableWidgetItem(value)
+                    item.setToolTip(value)
+                    self.table.setItem(row, col, item)
+            self.table.horizontalHeader().setSectionResizeMode(QtWidgets.QHeaderView.Stretch)
 
     def update_range(self):
         from_value = int(self.from_input.text()) if self.from_input.text() else self.start_index
@@ -210,25 +280,36 @@ def update_range(self):
             self.populate_table()
 
     def export_to_csv(self):
-        path, ok = QFileDialog.getSaveFileName(
-            self, self.tr("Save CSV"),
-            "annotations_statistics", "CSV files (*.csv)"
+        directory = QFileDialog.getExistingDirectory(
+            self, self.tr("Select Directory"),
+            ""
         )
-        if not ok or not path:
+        if not directory:
             return
-        if not path.endswith(".csv"):
-            path += ".csv"
 
         try:
-            total_infos = self.get_total_infos(1, len(self.image_file_list))
-            with open(path, "w", newline="", encoding="utf-8") as csvfile:
+            # Export label_infos
+            label_infos = self.get_total_infos(1, len(self.image_file_list))
+            label_infos_path = os.path.join(directory, "label_infos.csv")
+            with open(label_infos_path, "w", newline="", encoding="utf-8") as csvfile:
                 writer = csv.writer(csvfile)
-                for row in total_infos:
+                for row in label_infos:
                     writer.writerow(row)
+
+            # Export shape_infos
+            _, shape_infos = self.get_label_infos(1, len(self.image_file_list))
+            headers, shape_infos_data = self.get_shape_infos_table(shape_infos)
+            shape_infos_path = os.path.join(directory, "shape_infos.csv")
+            with open(shape_infos_path, "w", newline="", encoding="utf-8") as csvfile:
+                writer = csv.writer(csvfile)
+                writer.writerow(headers)
+                for row in shape_infos_data:
+                    writer.writerow(row)
+
             msg_box = QMessageBox()
             msg_box.setIcon(QMessageBox.Information)
             msg_box.setText(self.tr("Exporting successfully!"))
-            msg_box.setInformativeText(self.tr(f"Results have been saved to:\n{path}"))
+            msg_box.setInformativeText(self.tr(f"Results have been saved to:\n{label_infos_path}\nand\n{shape_infos_path}"))
             msg_box.setWindowTitle(self.tr("Success"))
             msg_box.exec_()
         except Exception as e:
@@ -238,3 +319,11 @@ def export_to_csv(self):
             error_dialog.setInformativeText(str(e))
             error_dialog.setWindowTitle(self.tr("Error"))
             error_dialog.exec_()
+
+    def toggle_info(self):
+        self.showing_label_infos = not self.showing_label_infos
+        if self.showing_label_infos:
+            self.toggle_button.setText(self.tr("Show Shape Infos"))
+        else:
+            self.toggle_button.setText(self.tr("Show Label Infos"))
+        self.populate_table(self.start_index, self.end_index)