Type annotation

dataset.py

@@ -18,7 +18,7 @@ class ImageDataset:
             )
         return ImageDataset(new_images)
 
-    def remove_image(self, image_id):
+    def remove_image(self, image_id: int):
         """Return new instance with image removed by ID."""
         new_images = [img for img in self.images if img["id"] != image_id]
         # Reindex remaining images
@@ -27,7 +27,7 @@ class ImageDataset:
             img["name"] = f"Image {i + 1}"
         return ImageDataset(new_images)
 
-    def update_label(self, image_id, label):
+    def update_label(self, image_id: int, label: str):
         """Return new instance with updated label."""
         new_images = []
         for img in self.images:
@@ -37,7 +37,7 @@ class ImageDataset:
             new_images.append(new_img)
         return ImageDataset(new_images)
 
-    def update_all_labels(self, labels_dict):
+    def update_all_labels(self, labels_dict: dict):
         """Return new instance with all labels updated."""
         new_images = []
         for img in self.images: