wehub-resource-sync
e904b667c6
PaddleOCR PR Tests GPU / test-pr-gpu (push) Blocked by required conditions
PaddleOCR PR Tests / test-pr (push) Blocked by required conditions
PaddleOCR PR Tests / test-pr-python (3.8) (push) Waiting to run
Build/Publish Develop Docs / deploy (push) Failing after 1s
PaddleOCR Code Style Check / check-code-style (push) Failing after 1s
PaddleOCR PR Tests GPU / detect-changes (push) Failing after 1s
PaddleOCR PR Tests GPU / test-pr-gpu-impl (push) Waiting to run
PaddleOCR PR Tests / detect-changes (push) Failing after 1s
PaddleOCR PR Tests / test-pr-python (3.13) (push) Waiting to run
PaddleOCR PR Tests / test-pr-python (3.9) (push) Waiting to run
545 lines
21 KiB
Python
545 lines
21 KiB
Python
# copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
|
||
#
|
||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||
# you may not use this file except in compliance with the License.
|
||
# You may obtain a copy of the License at
|
||
#
|
||
# http://www.apache.org/licenses/LICENSE-2.0
|
||
#
|
||
# Unless required by applicable law or agreed to in writing, software
|
||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
# See the License for the specific language governing permissions and
|
||
# limitations under the License.
|
||
import numpy as np
|
||
import cv2
|
||
import math
|
||
import os
|
||
import json
|
||
import random
|
||
import traceback
|
||
import multiprocessing
|
||
import urllib.request
|
||
import urllib.parse
|
||
import threading
|
||
import concurrent.futures
|
||
from collections import OrderedDict
|
||
from paddle.io import Dataset
|
||
from .imaug import transform, create_operators
|
||
from paddle import get_device
|
||
|
||
# ------------------------------------------------------------------ #
|
||
# Per-worker-process URL prefetch cache
|
||
#
|
||
# Each DataLoader worker is a forked process with its own copy of
|
||
# these globals. The thread pool and LRU cache are therefore
|
||
# completely independent across workers — no cross-process locking
|
||
# needed, and memory usage is bounded per worker.
|
||
#
|
||
# Memory budget (per worker):
|
||
# _URL_CACHE_MAX × avg_image_size ≈ 200 × 270 KB ≈ 54 MB
|
||
# With num_workers=4 the total extra footprint is ~216 MB.
|
||
#
|
||
# How prefetch works:
|
||
# _ensure_index_map() fires inside the worker when an epoch changes.
|
||
# It calls _prefetch_epoch_urls(), which scans the new _index_map,
|
||
# picks the first _URL_PREFETCH_SUBMIT URL entries (epoch order ≈
|
||
# access order), and submits them to the background thread pool.
|
||
# _load_image_bytes() checks the cache / in-flight future before
|
||
# falling back to a synchronous download.
|
||
# ------------------------------------------------------------------ #
|
||
|
||
_URL_CACHE_MAX = 200 # max cached images per worker
|
||
_URL_PREFETCH_SUBMIT = 200 # URL items submitted to thread pool per epoch
|
||
|
||
# LRU cache: url -> bytes
|
||
_url_cache: "OrderedDict[str, bytes]" = OrderedDict()
|
||
_url_cache_lock = threading.Lock()
|
||
|
||
# In-flight futures: url -> Future
|
||
_url_futures: "dict[str, concurrent.futures.Future]" = {}
|
||
_url_futures_lock = threading.Lock()
|
||
|
||
# Lazily created per-process thread pool
|
||
_url_executor: "concurrent.futures.ThreadPoolExecutor | None" = None
|
||
_url_executor_lock = threading.Lock()
|
||
|
||
|
||
def _get_url_executor():
|
||
global _url_executor
|
||
if _url_executor is None:
|
||
with _url_executor_lock:
|
||
if _url_executor is None:
|
||
_url_executor = concurrent.futures.ThreadPoolExecutor(
|
||
max_workers=4,
|
||
thread_name_prefix="url_prefetch",
|
||
)
|
||
return _url_executor
|
||
|
||
|
||
def _encode_url(url):
|
||
"""Percent-encode non-ASCII characters in the URL path so that
|
||
urllib can handle URLs containing CJK or other non-ASCII filenames.
|
||
Scheme, netloc, query and fragment are left untouched.
|
||
"""
|
||
parts = urllib.parse.urlparse(url)
|
||
encoded_path = urllib.parse.quote(parts.path, safe="/:@!$&'()*+,;=")
|
||
return urllib.parse.urlunparse(parts._replace(path=encoded_path))
|
||
|
||
|
||
def _download_url_bytes(url):
|
||
"""Download *url*, store in LRU cache, remove from futures dict.
|
||
The futures entry is always cleaned up (success or failure) so that
|
||
a failed URL can be retried on the next access.
|
||
"""
|
||
encoded = _encode_url(url)
|
||
try:
|
||
with urllib.request.urlopen(encoded, timeout=30) as resp:
|
||
data = resp.read()
|
||
except Exception:
|
||
with _url_futures_lock:
|
||
_url_futures.pop(url, None)
|
||
raise
|
||
with _url_cache_lock:
|
||
if url not in _url_cache:
|
||
if len(_url_cache) >= _URL_CACHE_MAX:
|
||
_url_cache.popitem(last=False) # evict LRU entry
|
||
_url_cache[url] = data
|
||
else:
|
||
_url_cache.move_to_end(url)
|
||
with _url_futures_lock:
|
||
_url_futures.pop(url, None)
|
||
return data
|
||
|
||
|
||
def _submit_url_prefetch(url):
|
||
"""Submit a background download for *url* if not already cached/in-flight."""
|
||
with _url_cache_lock:
|
||
if url in _url_cache:
|
||
return
|
||
with _url_futures_lock:
|
||
if url in _url_futures:
|
||
return
|
||
future = _get_url_executor().submit(_download_url_bytes, url)
|
||
_url_futures[url] = future
|
||
|
||
|
||
def _prefetch_epoch_urls(index_map, all_lines, delimiter):
|
||
"""Scan *index_map* and submit the first _URL_PREFETCH_SUBMIT URL
|
||
items for background download. Called inside worker processes."""
|
||
submitted = 0
|
||
for file_idx in index_map:
|
||
if submitted >= _URL_PREFETCH_SUBMIT:
|
||
break
|
||
try:
|
||
line = all_lines[file_idx].decode("utf-8")
|
||
fname = line.strip("\n").split(delimiter)[0]
|
||
if fname and fname[0] == "[": # JSON list — skip
|
||
continue
|
||
if fname.startswith("http://") or fname.startswith("https://"):
|
||
_submit_url_prefetch(fname)
|
||
submitted += 1
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
def _load_image_bytes(img_path):
|
||
"""Return raw image bytes. For URLs checks prefetch cache/future first."""
|
||
if img_path.startswith("http://") or img_path.startswith("https://"):
|
||
# 1. Cache hit — return immediately
|
||
with _url_cache_lock:
|
||
if img_path in _url_cache:
|
||
_url_cache.move_to_end(img_path)
|
||
return _url_cache[img_path]
|
||
# 2. In-flight future — wait for background download to finish
|
||
with _url_futures_lock:
|
||
future = _url_futures.get(img_path)
|
||
if future is not None:
|
||
return future.result(timeout=60)
|
||
# 3. Cold miss — download synchronously (also fills cache)
|
||
return _download_url_bytes(img_path)
|
||
with open(img_path, "rb") as f:
|
||
return f.read()
|
||
|
||
|
||
def _img_path_exists(img_path):
|
||
"""Return True if the image source is accessible (local file exists or URL)."""
|
||
if img_path.startswith("http://") or img_path.startswith("https://"):
|
||
return True
|
||
return os.path.exists(img_path)
|
||
|
||
|
||
class SimpleDataSet(Dataset):
|
||
def __init__(self, config, mode, logger, seed=None):
|
||
super(SimpleDataSet, self).__init__()
|
||
self.logger = logger
|
||
self.mode = mode.lower()
|
||
|
||
global_config = config["Global"]
|
||
dataset_config = config[mode]["dataset"]
|
||
loader_config = config[mode]["loader"]
|
||
|
||
self.delimiter = dataset_config.get("delimiter", "\t")
|
||
label_file_list = dataset_config.pop("label_file_list")
|
||
data_source_num = len(label_file_list)
|
||
ratio_list = dataset_config.get("ratio_list", 1.0)
|
||
if isinstance(ratio_list, (float, int)):
|
||
ratio_list = [float(ratio_list)] * int(data_source_num)
|
||
self.label_file_list = label_file_list
|
||
self.ratio_list = ratio_list
|
||
|
||
assert (
|
||
len(ratio_list) == data_source_num
|
||
), "The length of ratio_list should be the same as the file_list."
|
||
self.data_dir = dataset_config["data_dir"]
|
||
self.do_shuffle = loader_config["shuffle"]
|
||
self.seed = seed
|
||
self.need_reset = True in [x < 1 for x in ratio_list]
|
||
|
||
logger.info("Initialize indexs of datasets:%s" % label_file_list)
|
||
|
||
if self.need_reset:
|
||
# Pre-load all lines once (immutable, never re-read from disk).
|
||
# Per-epoch ratio sampling is done via _index_map (virtual idx -> global idx).
|
||
self._all_lines, self.file_boundaries = self._load_all_lines(
|
||
label_file_list
|
||
)
|
||
self._index_map = self._generate_index_map(seed)
|
||
self._cached_epoch = seed if seed is not None else 0
|
||
# data_lines / data_idx_order_list kept for API compat but NOT used in __getitem__
|
||
self.data_lines = self._all_lines
|
||
self.data_idx_order_list = list(range(len(self._index_map)))
|
||
else:
|
||
self._all_lines = None
|
||
self._index_map = None
|
||
self._cached_epoch = None
|
||
self.file_boundaries = None
|
||
self.data_lines = self.get_image_info_list(label_file_list, ratio_list)
|
||
self.data_idx_order_list = list(range(len(self.data_lines)))
|
||
if self.mode == "train" and self.do_shuffle:
|
||
self.shuffle_data_random()
|
||
|
||
# Shared epoch value: workers read this via shared memory to detect epoch changes
|
||
self._shared_epoch = multiprocessing.Value("i", seed if seed is not None else 0)
|
||
|
||
self.ops = create_operators(dataset_config["transforms"], global_config)
|
||
self.ext_op_transform_idx = dataset_config.get("ext_op_transform_idx", 2)
|
||
|
||
# ------------------------------------------------------------------ #
|
||
# Data loading helpers
|
||
# ------------------------------------------------------------------ #
|
||
|
||
def _load_all_lines(self, file_list):
|
||
"""Read all label files once. Returns (all_lines, file_boundaries)."""
|
||
if isinstance(file_list, str):
|
||
file_list = [file_list]
|
||
all_lines = []
|
||
boundaries = [0]
|
||
for f in file_list:
|
||
with open(f, "rb") as fh:
|
||
lines = fh.readlines()
|
||
all_lines.extend(lines)
|
||
boundaries.append(len(all_lines))
|
||
return all_lines, boundaries
|
||
|
||
def _generate_index_map(self, seed):
|
||
"""Generate virtual-index -> global-index mapping.
|
||
|
||
Replicates the EXACT sampling logic of original get_image_info_list +
|
||
shuffle_data_random: for each file, random.seed(seed) then
|
||
random.sample to pick indices, then random.seed(seed) + shuffle.
|
||
|
||
Since random.sample(population, k) with the same seed selects the
|
||
same POSITIONS regardless of population type, sampling from
|
||
range(start, end) yields the same positions as from lines[start:end].
|
||
"""
|
||
sampled = []
|
||
for i in range(len(self.ratio_list)):
|
||
start = self.file_boundaries[i]
|
||
end = self.file_boundaries[i + 1]
|
||
file_size = end - start
|
||
count = round(file_size * self.ratio_list[i])
|
||
if self.mode == "train" or self.ratio_list[i] < 1.0:
|
||
random.seed(seed)
|
||
sampled.extend(random.sample(range(start, end), count))
|
||
else:
|
||
sampled.extend(range(start, end))
|
||
if self.mode == "train" and self.do_shuffle:
|
||
random.seed(seed)
|
||
random.shuffle(sampled)
|
||
return sampled
|
||
|
||
def _ensure_index_map(self):
|
||
"""Lazily rebuild _index_map when worker detects epoch change via shared memory.
|
||
Also triggers URL prefetch on first call (epoch 0) and on every epoch change.
|
||
"""
|
||
if self._all_lines is None:
|
||
return
|
||
current_epoch = self._shared_epoch.value
|
||
epoch_changed = current_epoch != self._cached_epoch
|
||
first_call = not getattr(self, "_url_prefetch_initialized", False)
|
||
|
||
if epoch_changed:
|
||
self._index_map = self._generate_index_map(current_epoch)
|
||
self._cached_epoch = current_epoch
|
||
|
||
if epoch_changed or first_call:
|
||
self._url_prefetch_initialized = True
|
||
_prefetch_epoch_urls(self._index_map, self._all_lines, self.delimiter)
|
||
|
||
def get_image_info_list(self, file_list, ratio_list):
|
||
if isinstance(file_list, str):
|
||
file_list = [file_list]
|
||
data_lines = []
|
||
for idx, file in enumerate(file_list):
|
||
with open(file, "rb") as f:
|
||
lines = f.readlines()
|
||
if self.mode == "train" or ratio_list[idx] < 1.0:
|
||
random.seed(self.seed)
|
||
lines = random.sample(lines, round(len(lines) * ratio_list[idx]))
|
||
data_lines.extend(lines)
|
||
return data_lines
|
||
|
||
def shuffle_data_random(self):
|
||
random.seed(self.seed)
|
||
random.shuffle(self.data_lines)
|
||
return
|
||
|
||
# ------------------------------------------------------------------ #
|
||
# Epoch update (called from main process each epoch)
|
||
# ------------------------------------------------------------------ #
|
||
|
||
def reset_data_lines(self, seed=None, epoch=None):
|
||
"""Signal new epoch to persistent workers via shared memory.
|
||
|
||
Workers lazily rebuild their _index_map on next __getitem__ call.
|
||
No disk I/O, no dataloader reconstruction.
|
||
"""
|
||
self.seed = seed
|
||
epoch_val = epoch if epoch is not None else (seed if seed is not None else 0)
|
||
self._shared_epoch.value = int(epoch_val)
|
||
|
||
if self._all_lines is not None:
|
||
# Update main-process index_map (used by len() and batch_sampler)
|
||
self._index_map = self._generate_index_map(seed)
|
||
self._cached_epoch = int(epoch_val)
|
||
self.data_idx_order_list = list(range(len(self._index_map)))
|
||
else:
|
||
# Fallback for non-ratio cases
|
||
self.data_lines = self.get_image_info_list(
|
||
self.label_file_list, self.ratio_list
|
||
)
|
||
self.data_idx_order_list = list(range(len(self.data_lines)))
|
||
if self.mode == "train" and self.do_shuffle:
|
||
self.shuffle_data_random()
|
||
|
||
# ------------------------------------------------------------------ #
|
||
# Data access
|
||
# ------------------------------------------------------------------ #
|
||
|
||
def _try_parse_filename_list(self, file_name):
|
||
# multiple images -> one gt label
|
||
if len(file_name) > 0 and file_name[0] == "[":
|
||
try:
|
||
info = json.loads(file_name)
|
||
file_name = random.choice(info)
|
||
except:
|
||
pass
|
||
return file_name
|
||
|
||
def get_ext_data(self):
|
||
ext_data_num = 0
|
||
for op in self.ops:
|
||
if hasattr(op, "ext_data_num"):
|
||
ext_data_num = getattr(op, "ext_data_num")
|
||
break
|
||
load_data_ops = self.ops[: self.ext_op_transform_idx]
|
||
ext_data = []
|
||
|
||
while len(ext_data) < ext_data_num:
|
||
if self._index_map is not None:
|
||
# Sample from current epoch's subset (same as original)
|
||
self._ensure_index_map()
|
||
rand_virtual = np.random.randint(len(self._index_map))
|
||
file_idx = self._index_map[rand_virtual]
|
||
data_line = self._all_lines[file_idx]
|
||
else:
|
||
file_idx = self.data_idx_order_list[np.random.randint(self.__len__())]
|
||
data_line = self.data_lines[file_idx]
|
||
data_line = data_line.decode("utf-8")
|
||
substr = data_line.strip("\n").split(self.delimiter)
|
||
file_name = substr[0]
|
||
file_name = self._try_parse_filename_list(file_name)
|
||
label = substr[1]
|
||
img_path = (
|
||
file_name
|
||
if file_name.startswith("http://") or file_name.startswith("https://")
|
||
else os.path.join(self.data_dir, file_name)
|
||
)
|
||
data = {"img_path": img_path, "label": label}
|
||
if not _img_path_exists(img_path):
|
||
continue
|
||
try:
|
||
data["image"] = _load_image_bytes(img_path)
|
||
except Exception:
|
||
continue
|
||
data = transform(data, load_data_ops)
|
||
|
||
if data is None:
|
||
continue
|
||
if "polys" in data.keys():
|
||
if data["polys"].shape[1] != 4:
|
||
continue
|
||
ext_data.append(data)
|
||
return ext_data
|
||
|
||
def __getitem__(self, idx):
|
||
if self._index_map is not None:
|
||
self._ensure_index_map()
|
||
file_idx = self._index_map[idx]
|
||
data_line = self._all_lines[file_idx]
|
||
else:
|
||
file_idx = self.data_idx_order_list[idx]
|
||
data_line = self.data_lines[file_idx]
|
||
try:
|
||
data_line = data_line.decode("utf-8")
|
||
substr = data_line.strip("\n").split(self.delimiter)
|
||
file_name = substr[0]
|
||
file_name = self._try_parse_filename_list(file_name)
|
||
label = substr[1]
|
||
img_path = (
|
||
file_name
|
||
if file_name.startswith("http://") or file_name.startswith("https://")
|
||
else os.path.join(self.data_dir, file_name)
|
||
)
|
||
data = {"img_path": img_path, "label": label}
|
||
if not _img_path_exists(img_path):
|
||
raise Exception("{} does not exist!".format(img_path))
|
||
data["image"] = _load_image_bytes(img_path)
|
||
data["ext_data"] = self.get_ext_data()
|
||
data["filename"] = data["img_path"]
|
||
data["epoch"] = self._shared_epoch.value
|
||
outs = transform(data, self.ops)
|
||
except:
|
||
self.logger.error(
|
||
"When parsing line {}, error happened with msg: {}".format(
|
||
data_line, traceback.format_exc()
|
||
)
|
||
)
|
||
outs = None
|
||
if outs is None:
|
||
# during evaluation, we should fix the idx to get same results for many times of evaluation.
|
||
rnd_idx = (
|
||
np.random.randint(self.__len__())
|
||
if self.mode == "train"
|
||
else (idx + 1) % self.__len__()
|
||
)
|
||
return self.__getitem__(rnd_idx)
|
||
return outs
|
||
|
||
def __len__(self):
|
||
if self._index_map is not None:
|
||
return len(self._index_map)
|
||
return len(self.data_idx_order_list)
|
||
|
||
|
||
class MultiScaleDataSet(SimpleDataSet):
|
||
def __init__(self, config, mode, logger, seed=None):
|
||
super(MultiScaleDataSet, self).__init__(config, mode, logger, seed)
|
||
self.ds_width = config[mode]["dataset"].get("ds_width", False)
|
||
if self.ds_width:
|
||
self.wh_aware()
|
||
|
||
def wh_aware(self):
|
||
data_line_new = []
|
||
wh_ratio = []
|
||
for line in self.data_lines:
|
||
data_line_new.append(line)
|
||
line = line.decode("utf-8")
|
||
name, label, w, h = line.strip("\n").split(self.delimiter)
|
||
wh_ratio.append(float(w) / float(h))
|
||
|
||
self.data_lines = data_line_new
|
||
self.wh_ratio = np.array(wh_ratio)
|
||
self.wh_ratio_sort = np.argsort(self.wh_ratio)
|
||
self.data_idx_order_list = list(range(len(self.data_lines)))
|
||
|
||
def resize_norm_img(self, data, imgW, imgH, padding=True):
|
||
img = data["image"]
|
||
h = img.shape[0]
|
||
w = img.shape[1]
|
||
if not padding:
|
||
resized_image = cv2.resize(
|
||
img, (imgW, imgH), interpolation=cv2.INTER_LINEAR
|
||
)
|
||
resized_w = imgW
|
||
else:
|
||
ratio = w / float(h)
|
||
if math.ceil(imgH * ratio) > imgW:
|
||
resized_w = imgW
|
||
else:
|
||
resized_w = int(math.ceil(imgH * ratio))
|
||
resized_image = cv2.resize(img, (resized_w, imgH))
|
||
resized_image = resized_image.astype("float32")
|
||
|
||
resized_image = resized_image.transpose((2, 0, 1)) / 255
|
||
resized_image -= 0.5
|
||
resized_image /= 0.5
|
||
padding_im = np.zeros((3, imgH, imgW), dtype=np.float32)
|
||
padding_im[:, :, :resized_w] = resized_image
|
||
valid_ratio = min(1.0, float(resized_w / imgW))
|
||
data["image"] = padding_im
|
||
data["valid_ratio"] = valid_ratio
|
||
if "iluvatar_gpu" in get_device():
|
||
data["valid_ratio"] = np.float32(valid_ratio)
|
||
return data
|
||
|
||
def __getitem__(self, properties):
|
||
# properties is a tuple, contains (width, height, index)
|
||
img_height = properties[1]
|
||
idx = properties[2]
|
||
if self.ds_width and properties[3] is not None:
|
||
wh_ratio = properties[3]
|
||
img_width = img_height * (
|
||
1 if int(round(wh_ratio)) == 0 else int(round(wh_ratio))
|
||
)
|
||
file_idx = self.wh_ratio_sort[idx]
|
||
else:
|
||
file_idx = self.data_idx_order_list[idx]
|
||
img_width = properties[0]
|
||
wh_ratio = None
|
||
|
||
data_line = self.data_lines[file_idx]
|
||
try:
|
||
data_line = data_line.decode("utf-8")
|
||
substr = data_line.strip("\n").split(self.delimiter)
|
||
file_name = substr[0]
|
||
file_name = self._try_parse_filename_list(file_name)
|
||
label = substr[1]
|
||
img_path = (
|
||
file_name
|
||
if file_name.startswith("http://") or file_name.startswith("https://")
|
||
else os.path.join(self.data_dir, file_name)
|
||
)
|
||
data = {"img_path": img_path, "label": label}
|
||
if not _img_path_exists(img_path):
|
||
raise Exception("{} does not exist!".format(img_path))
|
||
data["image"] = _load_image_bytes(img_path)
|
||
data["ext_data"] = self.get_ext_data()
|
||
outs = transform(data, self.ops[:-1])
|
||
if outs is not None:
|
||
outs = self.resize_norm_img(outs, img_width, img_height)
|
||
outs = transform(outs, self.ops[-1:])
|
||
except:
|
||
self.logger.error(
|
||
"When parsing line {}, error happened with msg: {}".format(
|
||
data_line, traceback.format_exc()
|
||
)
|
||
)
|
||
outs = None
|
||
if outs is None:
|
||
# during evaluation, we should fix the idx to get same results for many times of evaluation.
|
||
rnd_idx = (idx + 1) % self.__len__()
|
||
return self.__getitem__([img_width, img_height, rnd_idx, wh_ratio])
|
||
return outs
|