53 lines
1.3 KiB
Markdown
53 lines
1.3 KiB
Markdown
|
# Adding New Data Transforms
|
||
|
|
|
||
|
|
## Customization data transformation
|
||
|
|
|
||
|
|
The customized data transformation must inherited from `BaseTransform` and implement `transform` function.
|
||
|
|
Here we use a simple flipping transformation as example:
|
||
|
|
|
||
|
|
```python
|
||
|
|
import random
|
||
|
|
import mmcv
|
||
|
|
from mmcv.transforms import BaseTransform, TRANSFORMS
|
||
|
|
|
||
|
|
@TRANSFORMS.register_module()
|
||
|
|
class MyFlip(BaseTransform):
|
||
|
|
def __init__(self, direction: str):
|
||
|
|
super().__init__()
|
||
|
|
self.direction = direction
|
||
|
|
|
||
|
|
def transform(self, results: dict) -> dict:
|
||
|
|
img = results['img']
|
||
|
|
results['img'] = mmcv.imflip(img, direction=self.direction)
|
||
|
|
return results
|
||
|
|
```
|
||
|
|
|
||
|
|
Moreover, import the new class.
|
||
|
|
|
||
|
|
```python
|
||
|
|
from .my_pipeline import MyFlip
|
||
|
|
```
|
||
|
|
|
||
|
|
Thus, we can instantiate a `MyFlip` object and use it to process the data dict.
|
||
|
|
|
||
|
|
```python
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
transform = MyFlip(direction='horizontal')
|
||
|
|
data_dict = {'img': np.random.rand(224, 224, 3)}
|
||
|
|
data_dict = transform(data_dict)
|
||
|
|
processed_img = data_dict['img']
|
||
|
|
```
|
||
|
|
|
||
|
|
Or, we can use `MyFlip` transformation in data pipeline in our config file.
|
||
|
|
|
||
|
|
```python
|
||
|
|
pipeline = [
|
||
|
|
...
|
||
|
|
dict(type='MyFlip', direction='horizontal'),
|
||
|
|
...
|
||
|
|
]
|
||
|
|
```
|
||
|
|
|
||
|
|
Note that if you want to use `MyFlip` in config, you must ensure the file containing `MyFlip` is imported during runtime.
|