update

Lightning-Universe · May 12, 2021 · c512f75 · c512f75
1 parent 4e5d044
commit c512f75
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 4 deletions.
diff --git a/flash/vision/segmentation/data.py b/flash/vision/segmentation/data.py
@@ -52,7 +52,18 @@
 class SemanticSegmentationNumpyDataSource(NumpyDataSource):
 
     def load_sample(self, sample: Dict[str, Any], dataset: Optional[Any] = None) -> Dict[str, Any]:
-        sample[DefaultDataKeys.INPUT] = torch.from_numpy(sample[DefaultDataKeys.INPUT]).float()
+        img = torch.from_numpy(sample[DefaultDataKeys.INPUT]).float()
+        sample[DefaultDataKeys.INPUT] = img
+        sample[DefaultDataKeys.METADATA] = img.shape
+        return sample
+
+
+class SemanticSegmentationTensorDataSource(TensorDataSource):
+
+    def load_sample(self, sample: Dict[str, Any], dataset: Optional[Any] = None) -> Dict[str, Any]:
+        img = sample[DefaultDataKeys.INPUT].float()
+        sample[DefaultDataKeys.INPUT] = img
+        sample[DefaultDataKeys.METADATA] = img.shape
         return sample
 
 
@@ -162,7 +173,7 @@ def __init__(
             data_sources={
                 DefaultDataSources.FILES: SemanticSegmentationPathsDataSource(),
                 DefaultDataSources.FOLDERS: SemanticSegmentationPathsDataSource(),
-                DefaultDataSources.TENSORS: TensorDataSource(),
+                DefaultDataSources.TENSORS: SemanticSegmentationTensorDataSource(),
                 DefaultDataSources.NUMPY: SemanticSegmentationNumpyDataSource(),
             },
             default_data_source=DefaultDataSources.FILES,

diff --git a/tests/vision/segmentation/test_model.py b/tests/vision/segmentation/test_model.py
@@ -89,7 +89,7 @@ def test_predict_tensor():
     data_pipe = DataPipeline(preprocess=SemanticSegmentationPreprocess(num_classes=1))
     out = model.predict(img, data_source="tensors", data_pipeline=data_pipe)
     assert isinstance(out[0], torch.Tensor)
-    assert out[0].shape == (196, 196)
+    assert out[0].shape == (10, 20)
 
 
 def test_predict_numpy():
@@ -98,4 +98,4 @@ def test_predict_numpy():
     data_pipe = DataPipeline(preprocess=SemanticSegmentationPreprocess(num_classes=1))
     out = model.predict(img, data_source="numpy", data_pipeline=data_pipe)
     assert isinstance(out[0], torch.Tensor)
-    assert out[0].shape == (196, 196)
+    assert out[0].shape == (10, 20)