From f320720bf298842b693ee65306f6759e3a3bfeee Mon Sep 17 00:00:00 2001
From: Edgar Riba <edgar.riba@gmail.com>
Date: Wed, 12 May 2021 16:39:12 +0200
Subject: [PATCH] add object detection prediction example and fix batch (#283)

* add object detection prediction example and fix batch

* add tests for default_unciollate

* Fix a test

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Co-authored-by: Ethan Harris <ethanwharris@gmail.com>
---
 flash/data/batch.py                        |  2 +
 flash_examples/predict/image_embedder.py   |  2 +-
 flash_examples/predict/object_detection.py | 31 ++++++++++++
 tests/data/test_batch.py                   | 57 ++++++++++++++++++++++
 4 files changed, 91 insertions(+), 1 deletion(-)
 create mode 100644 flash_examples/predict/object_detection.py
 create mode 100644 tests/data/test_batch.py

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/image_embedder.py b/flash_examples/predict/image_embedder.py
index 54df44a736..db4042383d 100644
--- a/flash_examples/predict/image_embedder.py
+++ b/flash_examples/predict/image_embedder.py
@@ -33,7 +33,7 @@
 random_image = torch.randn(1, 3, 244, 244)
 
 # 6. Generate an embedding from this random image.
-embeddings = embedder.predict(random_image, data_source="tensor")
+embeddings = embedder.predict(random_image, data_source="tensors")
 
 # 7. Print embeddings shape
 print(embeddings[0].shape)
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)
diff --git a/tests/data/test_batch.py b/tests/data/test_batch.py
new file mode 100644
index 0000000000..a57c941045
--- /dev/null
+++ b/tests/data/test_batch.py
@@ -0,0 +1,57 @@
+# 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.
+import torch
+from torch.tensor import Tensor
+from torch.testing import assert_allclose
+
+from flash.data.batch import default_uncollate
+
+
+class TestDefaultUncollate:
+
+    def test_smoke(self):
+        batch = torch.rand(2, 1)
+        assert default_uncollate(batch) is not None
+
+    def test_tensor_zero(self):
+        batch = torch.tensor(1)
+        output = default_uncollate(batch)
+        assert_allclose(batch, output)
+
+    def test_tensor_batch(self):
+        batch = torch.rand(2, 1)
+        output = default_uncollate(batch)
+        assert isinstance(output, list)
+        assert all([isinstance(x, torch.Tensor) for x in output])
+
+    def test_sequence(self):
+        B = 3  # batch_size
+
+        batch = {}
+        batch['a'] = torch.rand(B, 4)
+        batch['b'] = torch.rand(B, 2)
+        batch['c'] = torch.rand(B)
+
+        output = default_uncollate(batch)
+        assert isinstance(output, list)
+        assert len(batch) == B
+
+        for sample in output:
+            assert list(sample.keys()) == ['a', 'b', 'c']
+            assert isinstance(sample['a'], list)
+            assert len(sample['a']) == 4
+            assert isinstance(sample['b'], list)
+            assert len(sample['b']) == 2
+            assert isinstance(sample['c'], torch.Tensor)
+            assert len(sample['c'].shape) == 0