dataset.py

class ImageDataset:
    def __init__(self, images=None):
        self.images = images if images is not None else []

    def add_images(self, files):
        """Return new instance with added images."""
        new_images = self.images.copy()
        if files is None:
            return self
        for _, file in enumerate(files):
            new_images.append(
                {
                    "id": len(new_images),
                    "path": file.name,
                    "label": "",
                    "name": f"Image {len(new_images) + 1}",
                }
            )
        return ImageDataset(new_images)

    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
        for i, img in enumerate(new_images):
            img["id"] = i
            img["name"] = f"Image {i + 1}"
        return ImageDataset(new_images)

    def update_label(self, image_id: int, label: str):
        """Return new instance with updated label."""
        new_images = []
        for img in self.images:
            new_img = img.copy()
            if new_img["id"] == image_id:
                new_img["label"] = label
            new_images.append(new_img)
        return ImageDataset(new_images)

    def update_all_labels(self, labels_dict: dict):
        """Return new instance with all labels updated."""
        new_images = []
        for img in self.images:
            new_img = img.copy()
            if img["id"] in labels_dict:
                new_img["label"] = labels_dict[img["id"]]
            new_images.append(new_img)
        return ImageDataset(new_images)