diff --git a/flash/data/batch.py b/flash/data/batch.py index f08be37d02..38929f6079 100644 --- a/flash/data/batch.py +++ b/flash/data/batch.py @@ -243,6 +243,8 @@ def default_uncollate(batch: Any): batch_type = type(batch) if isinstance(batch, Tensor): + if len(batch.shape) == 0: # 0 shape tensors + return batch return list(torch.unbind(batch, 0)) elif isinstance(batch, Mapping): diff --git a/flash_examples/predict/object_detection.py b/flash_examples/predict/object_detection.py new file mode 100644 index 0000000000..dce35f0c3c --- /dev/null +++ b/flash_examples/predict/object_detection.py @@ -0,0 +1,31 @@ +# Copyright The PyTorch Lightning team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +from flash import Trainer +from flash.data.utils import download_data +from flash.vision import ObjectDetector + +# 1. Download the data +# Dataset Credit: https://www.kaggle.com/ultralytics/coco128 +download_data("https://github.com/zhiqwang/yolov5-rt-stack/releases/download/v0.3.0/coco128.zip", "data/") + +# 2. Load the model from a checkpoint +model = ObjectDetector.load_from_checkpoint("https://flash-weights.s3.amazonaws.com/object_detection_model.pt") + +# 3. Detect the object on the images +predictions = model.predict([ + "data/coco128/images/train2017/000000000025.jpg", + "data/coco128/images/train2017/000000000520.jpg", + "data/coco128/images/train2017/000000000532.jpg", +]) +print(predictions)