studies/code/my/vision/image_utils.py at master · yanwenheng/studies · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
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
#!/usr/bin/env python
# -*- coding:utf-8 -*-
"""
Time:
    2021-04-26 9:05 下午

Author:
    huayang

Subject:

"""
import os

__all__ = [
    'get_real_ext',
    'rename_to_real_ext'
]


def get_real_ext(image_path, return_is_same=False):
    """@Image Utils
    获取图像文件的真实后缀
    如果不是图片，返回后缀为 None
    该方法不能判断图片是否完整

    Args:
        image_path:
        return_is_same: 是否返回 `is_same`

    Returns:
        ext_real, is_same
        真实后缀，真实后缀与当前后缀是否相同
        如果当前文件不是图片，则 ext_real 为 None
    """
    import imghdr

    # 获取当前后缀
    ext_cur = os.path.splitext(image_path)[1]

    if ext_cur.startswith('.'):
        ext_cur = ext_cur[1:]

    # 获取真实后缀
    ext_real = imghdr.what(image_path)

    if return_is_same:
        # 是否相同
        is_same = ext_cur == ext_real or {ext_cur, ext_real} == {'jpg', 'jpeg'}

        return ext_real, is_same

    return ext_real


def rename_to_real_ext(image_path):
    """将图片重命名为真实后缀"""
    ext_real, is_same = get_real_ext(image_path, return_is_same=True)

    if is_same or ext_real is None:
        return

    prefix, _ = os.path.splitext(image_path)
    dst = '.'.join([prefix, ext_real])
    os.rename(image_path, dst)