File size: 3,000 Bytes
1fe6a2c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
import hashlib
import logging
from pathlib import Path
from typing import Callable, Optional
from urllib.request import urlretrieve
from zipfile import ZipFile

from tqdm.auto import tqdm

FILENAME = "craft_mlt_25k.pth"
URL = (
    "https://xc-models.oss-cn-zhangjiakou.aliyuncs.com/modelscope/studio/easyocr/craft_mlt_25k.zip"
)
MD5SUM = "2f8227d2def4037cdb3b34389dcf9ec1"
MD5MSG = "MD5 hash mismatch, possible file corruption"


logger = logging.getLogger(__name__)


def calculate_md5(path: Path) -> str:
    hash_md5 = hashlib.md5()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(4096), b""):
            hash_md5.update(chunk)
    return hash_md5.hexdigest()


def print_progress_bar(t: tqdm) -> Callable[[int, int, Optional[int]], None]:
    last = 0

    def update_to(

        count: int = 1, block_size: int = 1, total_size: Optional[int] = None

    ):
        nonlocal last
        if total_size is not None:
            t.total = total_size
        t.update((count - last) * block_size)
        last = count

    return update_to


def download_and_unzip(

    url: str, filename: str, model_storage_directory: Path, verbose: bool = True

):
    zip_path = model_storage_directory / "temp.zip"
    with tqdm(
        unit="B", unit_scale=True, unit_divisor=1024, miniters=1, disable=not verbose
    ) as t:
        reporthook = print_progress_bar(t)
        urlretrieve(url, str(zip_path), reporthook=reporthook)
    with ZipFile(zip_path, "r") as zipObj:
        zipObj.extract(filename, str(model_storage_directory))
    zip_path.unlink()


def prepare_model(model_storage_directory: Path, download=True, verbose: bool = True) -> bool:
    model_storage_directory.mkdir(parents=True, exist_ok=True)

    detector_path = model_storage_directory / FILENAME

    # try get model path
    model_available = False
    if not detector_path.is_file():
        if not download:
            raise FileNotFoundError(f"Missing {detector_path} and downloads disabled")
        logger.info(
            "Downloading detection model, please wait. "
            "This may take several minutes depending upon your network connection."
        )
    elif calculate_md5(detector_path) != MD5SUM:
        logger.warning(MD5MSG)
        if not download:
            raise FileNotFoundError(
                f"MD5 mismatch for {detector_path} and downloads disabled"
            )
        detector_path.unlink()
        logger.info(
            "Re-downloading the detection model, please wait. "
            "This may take several minutes depending upon your network connection."
        )
    else:
        model_available = True

    if not model_available:
        download_and_unzip(URL, FILENAME, model_storage_directory, verbose)
        if calculate_md5(detector_path) != MD5SUM:
            raise ValueError(MD5MSG)
        logger.info("Download complete")

    return detector_path