Lightning-Universe / lightning-flash

Your PyTorch AI Factory - Flash enables you to easily configure and run complex AI recipes for over 15 tasks across 7 data domains
https://lightning-flash.readthedocs.io
Apache License 2.0
1.74k stars 212 forks source link

SemanticSegmentation has no available heads or backbones in FlashRegistry #1673

Open dframirez-usmc opened 1 year ago

dframirez-usmc commented 1 year ago

🐛 Bug

flash.image.SemanticSegmentation class has no available heads or backbones.

To Reproduce

# using latest Miniforge3-Linux-x86_64 for conda
conda create -n flash-test python=3.9 -y -c conda-forge
# alternatively:
# conda create -n flash-test python=3.10 -y -c conda-forge
conda activate flash-test

pip install 'lightning-flash[image]'

Code sample

From https://lightning-flash.readthedocs.io/en/stable/general/backbones.html#backbones-heads

from flash.image import SemanticSegmentation
# get the heads available for SemanticSegmentation
heads = SemanticSegmentation.available_heads()
# print the heads
print(heads)

# [ ] # EMPTY!
from flash.image import SemanticSegmentation
# get the backbones available for SemanticSegmentation
backbones = SemanticSegmentation.available_backbones()
# print the backbones
print(backbones)

# [ ] # EMPTY!

From https://lightning-flash.readthedocs.io/en/stable/reference/semantic_segmentation.html

import flash
import torch
from flash.core.data.utils import download_data
from flash.image import SemanticSegmentation, SemanticSegmentationData

download_data(
    "https://github.com/ongchinkiat/LyftPerceptionChallenge/releases/download/v0.1/carla-capture-20180513A.zip",
    "./data",
)
datamodule = SemanticSegmentationData.from_folders(
    train_folder="data/CameraRGB",
    train_target_folder="data/CameraSeg",
    val_split=0.1,
    transform_kwargs={"image_size": (256, 256)},
    num_classes=21,
    batch_size=4,
)
model = SemanticSegmentation(
    backbone="mobilenetv3_large_100",
    head="fpn",
    num_classes=datamodule.num_classes,
)

# KeyError: 'Key: mobilenetv3_large_100 is not in FlashRegistry. Available keys: []'

From https://github.com/Lightning-Universe/lightning-flash/blob/master/README.md

from flash.image import SemanticSegmentationData

dm = SemanticSegmentationData.from_folders(
    train_folder="data/CameraRGB",
    train_target_folder="data/CameraSeg",
    val_split=0.1,
    image_size=(256, 256),
    num_classes=21,
)

from flash.image import SemanticSegmentation

print(SemanticSegmentation.available_heads())

# [ ] # EMPTY!

print(SemanticSegmentation.available_backbones('fpn'))

# KeyError: 'Key: fpn is not in FlashRegistry. Available keys: []'

Expected behavior

SemanticSegmentation.available_heads() should return a list of heads. SemanticSegmentation.available_backbones() should return a list of backbones.

Environment

Issue exists in lightning-flash==0.8.1.post0 also

Additional context

ImageClassifier.available_heads() works as expected. ImageClassifier.available_backbones() works as expected.