Merge pull request #10 from shamspias/feat/rt-detr

Fixing label index
shamspias · Sep 11, 2024 · df50307 · df50307
2 parents 1f3c24e + 22a6ed4
commit df50307
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 3 deletions.
diff --git a/app/main.py b/app/main.py
@@ -65,7 +65,7 @@ def continue_ui(self):
                                                                                        "mmdet",
                                                                                        ]
                                                                )
-            self.config.sahi_device = st.sidebar.selectbox("Device:", ["cpu"])
+            self.config.sahi_device = st.sidebar.selectbox("Device:", ["cpu", "cuda:0"])
             self.config.sahi_slice_size = st.sidebar.slider("SAHI slice size:", 128, 512, (256, 256))
             self.config.sahi_overlap_ratio = st.sidebar.slider("SAHI overlap ratio:", 0.1, 0.5, (0.2, 0.2))
             self.sahi_config = {

diff --git a/formats/base_format.py b/formats/base_format.py
@@ -66,24 +66,29 @@ def process_results(self, results: Dict, img_dimensions, supported_classes) -> L
             for box in results['boxes']:  # Assuming SAHI results are formatted similarly
                 class_id = int(box['cls'][0])
                 if class_id in supported_classes:  # Check if class_id is in the list of supported classes
+                    class_id_index = supported_classes.index(
+                        class_id)  # Get index of class_id in supported_classes list
                     xmin, ymin, xmax, ymax = box['xyxy'][0]
                     x_center = ((xmin + xmax) / 2) / img_width
                     y_center = ((ymin + ymax) / 2) / img_height
                     width = (xmax - xmin) / img_width
                     height = (ymax - ymin) / img_height
-                    annotations.append(f"{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}")
+                    annotations.append(f"{class_id_index} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}")
         else:
             for result in results:
                 if hasattr(result, 'boxes') and result.boxes is not None:
                     for box in result.boxes:
                         class_id = int(box.cls[0])
                         if class_id in supported_classes:  # Check if class_id is in the list of supported classes
+                            class_id_index = supported_classes.index(
+                                class_id)  # Get index of class_id in supported_classes list
                             xmin, ymin, xmax, ymax = box.xyxy[0]
                             x_center = ((xmin + xmax) / 2) / img_width
                             y_center = ((ymin + ymax) / 2) / img_height
                             width = (xmax - xmin) / img_width
                             height = (ymax - ymin) / img_height
-                            annotations.append(f"{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}")
+                            annotations.append(
+                                f"{class_id_index} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}")
 
         return annotations