Skip to content
This repository has been archived by the owner on Oct 9, 2023. It is now read-only.

add object detection prediction example and fix batch #283

Merged
merged 5 commits into from
May 12, 2021
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions flash/data/batch.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Any reason not to return a list here? (like [batch])

return list(torch.unbind(batch, 0))

elif isinstance(batch, Mapping):
Expand Down
31 changes: 31 additions & 0 deletions flash_examples/predict/object_detection.py
Original file line number Diff line number Diff line change
@@ -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)
57 changes: 57 additions & 0 deletions tests/data/test_batch.py
Original file line number Diff line number Diff line change
@@ -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