Skip to content

Fix mixed image channel sizes in image classifiers#678

Open
Creylay wants to merge 1 commit into
developfrom
fix/image-classifiers-mixed-channels
Open

Fix mixed image channel sizes in image classifiers#678
Creylay wants to merge 1 commit into
developfrom
fix/image-classifiers-mixed-channels

Conversation

@Creylay
Copy link
Copy Markdown
Collaborator

@Creylay Creylay commented Jun 4, 2026

Summary

Image classifiers (MLP, CNN, LeNet5) failed when training on datasets containing images with mixed channel formats and/or class-imbalanced splits. Two independent bugs were fixed:

  1. The transform pipeline did not normalize image channels before converting to tensors, causing torch.stack to fail when a batch mixed grayscale, RGB, and RGBA images.
  2. label_to_idx was built only from training split labels, so classes with no training examples (present only in validation) had no index. This caused prepare_output to produce a mixed int/string list that PyArrow could not convert.

Type of Change

  • Backend change
  • Frontend change
  • CI / Workflow change
  • Build / Packaging change
  • Bug fix
  • Documentation

Changes (by file)

  • DashAI/back/models/mlp_image_classifier.py: added convert("RGB") to transform pipeline; build label_to_idx from Categorical.categories (full dataset class list) instead of training split labels; changed prepare_output fallback from string to -1.
  • DashAI/back/models/cnn_image_classifier.py: same fixes.
  • DashAI/back/models/lenet5_image_classifier.py: same fixes.
  • DashAI/back/models/base_torchvision_image_classifier.py: label index fixes only (convert("RGB") was already present).

Testing

Train any image classification model on a dataset that:

  • Mixes grayscale, RGB, and/or RGBA images — training should complete without RuntimeError: stack expects each tensor to be equal size.
  • Has classes with few examples that may land entirely in the validation split — training should complete without ArrowInvalid: tried to convert to int64 or y_true contains values {-1}.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant