io.md 2.93 KB
Newer Older
Kai Chen's avatar
Kai Chen committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
## File IO

This module provides two universal API to load and dump files of different formats.

### Load and dump data
`mmcv` provides a universal api for loading and dumping data, currently
supported formats are json, yaml and pickle.

```python
import mmcv

# load data from a file
data = mmcv.load('test.json')
data = mmcv.load('test.yaml')
Kai Chen's avatar
Kai Chen committed
15
data = mmcv.load('test.pkl')
Kai Chen's avatar
Kai Chen committed
16
17
18
19
20
21
# load data from a file-like object
with open('test.json', 'r') as f:
    data = mmcv.load(f)

# dump data to a string
json_str = mmcv.dump(data, format='json')
Kai Chen's avatar
Kai Chen committed
22

Kai Chen's avatar
Kai Chen committed
23
# dump data to a file with a filename (infer format from file extension)
Kai Chen's avatar
Kai Chen committed
24
25
mmcv.dump(data, 'out.pkl')

Kai Chen's avatar
Kai Chen committed
26
27
28
29
30
# dump data to a file with a file-like object
with open('test.yaml', 'w') as f:
    data = mmcv.dump(data, f, format='yaml')
```

Kai Chen's avatar
Kai Chen committed
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
It is also very convenient to extend the api to support more file formats.
All you need to do is to write a file handler inherited from `BaseFileHandler`
and register it with one or several file formats.

You need to implement at least 3 methods.

```python
import mmcv

# To register multiple file formats, a list can be used as the argument.
# @mmcv.register_handler(['txt', 'log'])
@mmcv.register_handler('txt')
class TxtHandler1(mmcv.BaseFileHandler):

    def load_from_fileobj(self, file):
        return file.read()

    def dump_to_fileobj(self, obj, file):
        file.write(str(obj))

    def dump_to_str(self, obj, **kwargs):
        return str(obj)
```

Here is an example of `PickleHandler`.

```python
from six.moves import cPickle as pickle

class PickleHandler(mmcv.BaseFileHandler):

    def load_from_fileobj(self, file, **kwargs):
        return pickle.load(file, **kwargs)

    def load_from_path(self, filepath, **kwargs):
        return super(PickleHandler, self).load_from_path(
            filepath, mode='rb', **kwargs)

    def dump_to_str(self, obj, **kwargs):
        kwargs.setdefault('protocol', 2)
        return pickle.dumps(obj, **kwargs)

    def dump_to_fileobj(self, obj, file, **kwargs):
        kwargs.setdefault('protocol', 2)
        pickle.dump(obj, file, **kwargs)

    def dump_to_path(self, obj, filepath, **kwargs):
        super(PickleHandler, self).dump_to_path(
            obj, filepath, mode='wb', **kwargs)
```

Kai Chen's avatar
Kai Chen committed
82
83
84
85
86
87
88
89
90
91
92
93
94
95
### Load a text file as a list or dict

For example `a.txt` is a text file with 5 lines.
```
a
b
c
d
e
```

Then use `list_from_file` to load the list from a.txt.

```python
Kai Chen's avatar
Kai Chen committed
96
97
98
99
100
101
102
103
>>> mmcv.list_from_file('a.txt')
['a', 'b', 'c', 'd', 'e']
>>> mmcv.list_from_file('a.txt', offset=2)
['c', 'd', 'e']
>>> mmcv.list_from_file('a.txt', max_num=2)
['a', 'b']
>>> mmcv.list_from_file('a.txt', prefix='/mnt/')
['/mnt/a', '/mnt/b', '/mnt/c', '/mnt/d', '/mnt/e']
Kai Chen's avatar
Kai Chen committed
104
105
106
107
108
109
110
111
112
113
114
115
```

For example `b.txt` is a text file with 5 lines.
```
1 cat
2 dog cow
3 panda
```

Then use `dict_from_file` to load the list from a.txt.

```python
Kai Chen's avatar
Kai Chen committed
116
117
118
119
>>> mmcv.dict_from_file('b.txt')
{'1': 'cat', '2': ['dog', 'cow'], '3': 'panda'}
>>> mmcv.dict_from_file('b.txt', key_type=int)
{1: 'cat', 2: ['dog', 'cow'], 3: 'panda'}
Kai Chen's avatar
Kai Chen committed
120
```