Merge pull request #18 from Jordan-Pierce/update-tileprogress

Update TileProgress class and related code

tests/test_yolo_tiler.py

@@ -6,8 +6,9 @@
 
 
 def progress_callback(progress: TileProgress):
-    print(f"Processing {progress.current_image} in {progress.current_set} set: "
-          f"tile {progress.current_tile}/{progress.total_tiles}")
+    print(f"Processing {progress.current_image_name} in {progress.current_set_name} set: "
+          f"tile {progress.current_tile_idx}/{progress.total_tiles}, "
+          f"image {progress.current_image_idx}/{progress.total_images}")
 
 
 src = "./tests/segmentation"
@@ -33,7 +34,7 @@ def progress_callback(progress: TileProgress):
     target=dst,
     config=config,
     num_viz_samples=100,
-    # progress_callback=progress_callback
+    progress_callback=progress_callback
 )
 
 # Run tiling process

yolo_tiler/yolo_tiler.py

@@ -102,10 +102,12 @@ def get_effective_area(self, image_width: int, image_height: int) -> Tuple[int,
 @dataclass
 class TileProgress:
     """Data class to track tiling progress"""
-    current_tile: int
+    current_set_name: str
+    current_image_name: str
+    current_image_idx: int
+    total_images: int
+    current_tile_idx: int
     total_tiles: int
-    current_set: str
-    current_image: str
 
 
 class YoloTiler:
@@ -164,21 +166,21 @@ def _tqdm_callback(self, progress: TileProgress):
             progress: TileProgress object containing current progress
             
         """
-        if progress.current_set not in self._progress_bars:
-            self._progress_bars[progress.current_set] = tqdm(
+        if progress.current_set_name not in self._progress_bars:
+            self._progress_bars[progress.current_set_name] = tqdm(
                 total=progress.total_tiles,
-                desc=progress.current_set,
+                desc=progress.current_set_name,
                 unit='items'
             )
 
         # Update progress
-        self._progress_bars[progress.current_set].n = progress.current_tile
-        self._progress_bars[progress.current_set].refresh()
+        self._progress_bars[progress.current_set_name].n = progress.current_tile_idx
+        self._progress_bars[progress.current_set_name].refresh()
 
         # Close and cleanup if task is complete
-        if progress.current_tile >= progress.total_tiles:
-            self._progress_bars[progress.current_set].close()
-            del self._progress_bars[progress.current_set]
+        if progress.current_tile_idx >= progress.total_tiles:
+            self._progress_bars[progress.current_set_name].close()
+            del self._progress_bars[progress.current_set_name]
 
     def _setup_logger(self) -> logging.Logger:
         """Configure logging for the tiler"""
@@ -412,7 +414,7 @@ def _save_labels(self, labels: List, path: Path, is_segmentation: bool) -> None:
             df = pd.DataFrame(labels, columns=['class', 'x1', 'y1', 'w', 'h'])
             df.to_csv(path, sep=' ', index=False, header=False, float_format='%.6f')
 
-    def tile_image(self, image_path: Path, label_path: Path, folder: str) -> None:
+    def tile_image(self, image_path: Path, label_path: Path, folder: str, current_image_idx: int, total_images: int) -> None:
         """
         Tile an image and its corresponding labels, properly handling margins.
         """
@@ -524,10 +526,12 @@ def clean_geometry(geom: Polygon) -> Polygon:
 
                 if self.progress_callback:
                     progress = TileProgress(
-                        current_tile=tile_idx + 1,
+                        current_tile_idx=tile_idx + 1,
                         total_tiles=total_tiles,
-                        current_set=folder.rstrip('/'),
-                        current_image=image_path.name
+                        current_set_name=folder.rstrip('/'),
+                        current_image_name=image_path.name,
+                        current_image_idx=current_image_idx,
+                        total_images=total_images
                     )
                     self.progress_callback(progress)
 
@@ -668,10 +672,12 @@ def split_data(self) -> None:
 
             if self.progress_callback:
                 progress = TileProgress(
-                    current_tile=tile_idx + 1,
+                    current_tile_idx=tile_idx + 1,
                     total_tiles=num_valid,
-                    current_set='valid',
-                    current_image=image_path.name
+                    current_set_name='valid',
+                    current_image_name=image_path.name,
+                    current_image_idx=0,  # Placeholder, update as needed
+                    total_images=0  # Placeholder, update as needed
                 )
                 self.progress_callback(progress)
 
@@ -680,10 +686,12 @@ def split_data(self) -> None:
             self._move_split_data(image_path, label_path, 'test')
             if self.progress_callback:
                 progress = TileProgress(
-                    current_tile=tile_idx + 1,
+                    current_tile_idx=tile_idx + 1,
                     total_tiles=num_test,
-                    current_set='test',
-                    current_image=image_path.name
+                    current_set_name='test',
+                    current_image_name=image_path.name,
+                    current_image_idx=0,  # Placeholder, update as needed
+                    total_images=0  # Placeholder, update as needed
                 )
                 self.progress_callback(progress)
 
@@ -724,11 +732,13 @@ def _process_subfolder(self, subfolder: str) -> None:
             self.logger.error(f"Number of images and labels do not match in {subfolder} directory, skipping")
             return
 
+        total_images = len(image_paths)
+
         # Process each image
-        for image_path, label_path in list(zip(image_paths, label_paths)):
+        for current_image_idx, (image_path, label_path) in enumerate(zip(image_paths, label_paths)):
             assert image_path.stem == label_path.stem, "Image and label filenames do not match"
             self.logger.info(f'Processing {image_path}')
-            self.tile_image(image_path, label_path, subfolder)
+            self.tile_image(image_path, label_path, subfolder, current_image_idx + 1, total_images)
 
     def _check_and_split_data(self) -> None:
         """Check if valid or test folders are empty and split data if necessary."""
@@ -778,10 +788,12 @@ def visualize_random_samples(self) -> None:
 
             if self.progress_callback:
                 progress = TileProgress(
-                    current_tile=tile_idx + 1,
+                    current_tile_idx=tile_idx + 1,
                     total_tiles=num_samples,
-                    current_set='rendered',
-                    current_image=image_path.name
+                    current_set_name='rendered',
+                    current_image_name=image_path.name,
+                    current_image_idx=0,  # Placeholder, update as needed
+                    total_images=0  # Placeholder, update as needed
                 )
                 self.progress_callback(progress)