Repository: docarray/docarray

Issue:
`plot_embeddings` with `image_sprites` from uri
**To produce:**

```python
from docarray import DocumentArray, Document
import torchvision

def preproc(d: Document):
    return (d.load_uri_to_image_tensor()  
             .set_image_tensor_shape((200, 200))  
             .set_image_tensor_normalization() 
             .set_image_tensor_channel_axis(-1, 0)) 

left_da = DocumentArray.from_files('left/*.jpg')

left_da.apply(preproc)

model = torchvision.models.resnet50(pretrained=True)  
left_da.embed(model)  

left_da.plot_embeddings(image_sprites=True)
```

Error:
```
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Input In [1], in <module>
     14 model = torchvision.models.resnet50(pretrained=True)  
     15 left_da.embed(model)  
---> 17 left_da.plot_embeddings(image_sprites=True)

File ~/workplaces/docarray-env/docarray/docarray/array/mixins/plot.py:152, in PlotMixin.plot_embeddings(self, title, path, image_sprites, min_image_size, channel_axis, start_server, host, port)
    140     if len(self) > max_docs:
    141         warnings.warn(
    142             f'''
    143             {self!r} has more than {max_docs} elements, which is the maximum number of image sprites can support. 
   (...)
    149             '''
    150         )
--> 152     self.plot_image_sprites(
    153         os.path.join(path, sprite_fn),
    154         canvas_size=canvas_size,
    155         min_size=min_image_size,
    156         channel_axis=channel_axis,
    157     )
    159 self.save_embeddings_csv(os.path.join(path, emb_fn), delimiter='\t')
    161 _exclude_fields = ('embedding', 'tensor', 'scores')

File ~/workplaces/docarray-env/docarray/docarray/array/mixins/plot.py:331, in PlotMixin.plot_image_sprites(self, output, canvas_size, min_size, channel_axis)
    329 row_id = floor(img_id / img_per_row)
    330 col_id = img_id % img_per_row
--> 331 sprite_img[
    332     (row_id * img_size) : ((row_id + 1) * img_size),
    333     (col_id * img_size) : ((col_id + 1) * img_size),
    334 ] = _d.tensor
    336 img_id += 1
    337 if img_id >= max_num_img:

ValueError: could not broadcast input array from shape (16,16,200) into shape (16,16,3)
```

**Suggested fix:** 
passing `use_uri` flag to `plot_embeddings` to allow visualizing the dots using `uri` attribute. I would be happy to raise a PR for this. 


Discussion:
the broadcast issue can be resolved by changing the channel axis using `channel_axis=0` but this will not visualize the original images. 
> the broadcast issue can be resolved by checking the channel axis using `channel_axis=0` but this will not visualize the original images.

maybe u can transform the axis back before plotting? I guess the `stacking` is only needed at `embedding` time

> maybe u can transform the axis back before plotting? I guess the `stacking` is only needed at `embedding` time

Changing the axis will work , however , it would be nice to have an option to show the original images when projecting embeddings without processing them again. 
More than the original, you expect plot_embeddings to force some axis as the channel one?
> More than the original, you expect plot_embeddings to force some axis as the channel one?


just a flag to force loading from uri and setting the channel axis:
https://github.com/jina-ai/docarray/blob/79c13a09980c97c9b4fcf9857566ff2a46c95c83/docarray/array/mixins/plot.py#L321

```python
if use_uri or _d.content_type != 'tensor':
    _d.load_uri_to_image_tensor()
    channel_axis = -1
```
Another possible approach for this problem.  As the applied `preproc` function is only for model inference. Can we make it possible to pass the `prepoc` function to  `. embed()`:

```python
def preproc(d: Document)
    ....
    return d

da.embed(model, preproc_fn=preproc)
```
Then, the output (copy of the origin DA) of `preproc` can then be fed into the model for inference, and will not change the value of origin DA. 

With this change, the user can choose which method to use regarding context. If the user wants to store the updates, then apply `preproc` before `.embed()`. If the user don't want keep the updates, then apply the `preproc` within the `.embed()`


Resolving patch:
diff --git a/docarray/array/mixins/plot.py b/docarray/array/mixins/plot.py
--- a/docarray/array/mixins/plot.py
+++ b/docarray/array/mixins/plot.py
@@ -108,6 +108,7 @@ def plot_embeddings(
         start_server: bool = True,
         host: str = '127.0.0.1',
         port: Optional[int] = None,
+        image_source: str = 'tensor',
     ) -> str:
         """Interactively visualize :attr:`.embeddings` using the Embedding Projector.
 
@@ -121,6 +122,7 @@ def plot_embeddings(
         :param min_image_size: only used when `image_sprites=True`. the minimum size of the image
         :param channel_axis: only used when `image_sprites=True`. the axis id of the color channel, ``-1`` indicates the color channel info at the last axis
         :param start_server: if set, start a HTTP server and open the frontend directly. Otherwise, you need to rely on ``return`` path and serve by yourself.
+        :param image_source: specify where the image comes from, can be ``uri`` or ``tensor``. empty tensor will fallback to uri
         :return: the path to the embeddings visualization info.
         """
         from ...helper import random_port, __resources_path__
@@ -154,6 +156,7 @@ def plot_embeddings(
                 canvas_size=canvas_size,
                 min_size=min_image_size,
                 channel_axis=channel_axis,
+                image_source=image_source,
             )
 
         self.save_embeddings_csv(os.path.join(path, emb_fn), delimiter='\t')
@@ -287,6 +290,7 @@ def plot_image_sprites(
         canvas_size: int = 512,
         min_size: int = 16,
         channel_axis: int = -1,
+        image_source: str = 'tensor',
     ) -> None:
         """Generate a sprite image for all image tensors in this DocumentArray-like object.
 
@@ -297,6 +301,7 @@ def plot_image_sprites(
         :param canvas_size: the size of the canvas
         :param min_size: the minimum size of the image
         :param channel_axis: the axis id of the color channel, ``-1`` indicates the color channel info at the last axis
+        :param image_source: specify where the image comes from, can be ``uri`` or ``tensor``. empty tensor will fallback to uri
         """
         if not self:
             raise ValueError(f'{self!r} is empty')
@@ -316,26 +321,36 @@ def plot_image_sprites(
             [img_size * img_per_row, img_size * img_per_row, 3], dtype='uint8'
         )
         img_id = 0
-        for d in self:
-            _d = copy.deepcopy(d)
-            if _d.content_type != 'tensor':
-                _d.load_uri_to_image_tensor()
-                channel_axis = -1
-
-            _d.set_image_tensor_channel_axis(channel_axis, -1).set_image_tensor_shape(
-                shape=(img_size, img_size)
-            )
-
-            row_id = floor(img_id / img_per_row)
-            col_id = img_id % img_per_row
-            sprite_img[
-                (row_id * img_size) : ((row_id + 1) * img_size),
-                (col_id * img_size) : ((col_id + 1) * img_size),
-            ] = _d.tensor
-
-            img_id += 1
-            if img_id >= max_num_img:
-                break
+        try:
+            for d in self:
+                _d = copy.deepcopy(d)
+
+                if image_source == 'uri' or (
+                    image_source == 'tensor' and _d.content_type != 'tensor'
+                ):
+                    _d.load_uri_to_image_tensor()
+                    channel_axis = -1
+                elif image_source not in ('uri', 'tensor'):
+                    raise ValueError(f'image_source can be only `uri` or `tensor`')
+
+                _d.set_image_tensor_channel_axis(
+                    channel_axis, -1
+                ).set_image_tensor_shape(shape=(img_size, img_size))
+
+                row_id = floor(img_id / img_per_row)
+                col_id = img_id % img_per_row
+                sprite_img[
+                    (row_id * img_size) : ((row_id + 1) * img_size),
+                    (col_id * img_size) : ((col_id + 1) * img_size),
+                ] = _d.tensor
+
+                img_id += 1
+                if img_id >= max_num_img:
+                    break
+        except Exception as ex:
+            raise ValueError(
+                'Bad image tensor. Try different `image_source` or `channel_axis`'
+            ) from ex
 
         from PIL import Image
 


Test patch:
diff --git a/tests/unit/array/mixins/test_plot.py b/tests/unit/array/mixins/test_plot.py
--- a/tests/unit/array/mixins/test_plot.py
+++ b/tests/unit/array/mixins/test_plot.py
@@ -8,7 +8,7 @@
 from docarray import DocumentArray, Document
 
 
-def test_sprite_image_generator(pytestconfig, tmpdir):
+def test_sprite_fail_tensor_success_uri(pytestconfig, tmpdir):
     da = DocumentArray.from_files(
         [
             f'{pytestconfig.rootdir}/**/*.png',
@@ -16,7 +16,26 @@ def test_sprite_image_generator(pytestconfig, tmpdir):
             f'{pytestconfig.rootdir}/**/*.jpeg',
         ]
     )
-    da.plot_image_sprites(tmpdir / 'sprint_da.png')
+    da.apply(
+        lambda d: d.load_uri_to_image_tensor().set_image_tensor_channel_axis(-1, 0)
+    )
+    with pytest.raises(ValueError):
+        da.plot_image_sprites()
+    da.plot_image_sprites(tmpdir / 'sprint_da.png', image_source='uri')
+    assert os.path.exists(tmpdir / 'sprint_da.png')
+
+
+@pytest.mark.parametrize('image_source', ['tensor', 'uri'])
+def test_sprite_image_generator(pytestconfig, tmpdir, image_source):
+    da = DocumentArray.from_files(
+        [
+            f'{pytestconfig.rootdir}/**/*.png',
+            f'{pytestconfig.rootdir}/**/*.jpg',
+            f'{pytestconfig.rootdir}/**/*.jpeg',
+        ]
+    )
+    da.apply(lambda d: d.load_uri_to_image_tensor())
+    da.plot_image_sprites(tmpdir / 'sprint_da.png', image_source=image_source)
     assert os.path.exists(tmpdir / 'sprint_da.png')