util.py 1.2 KB

12345678910111213141516171819202122232425262728293031323334
  1. # Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. import platform
  15. import multiprocessing as mp
  16. def is_pic(img_name):
  17. valid_suffix = ["JPEG", "jpeg", "JPG", "jpg", "BMP", "bmp", "PNG", "png"]
  18. suffix = img_name.split(".")[-1]
  19. if suffix not in valid_suffix:
  20. return False
  21. return True
  22. def get_num_workers(num_workers):
  23. if not platform.system() == "Linux":
  24. # Dataloader with multi-process model is not supported
  25. # on MacOS and Windows currently.
  26. return 0
  27. if num_workers == "auto":
  28. num_workers = mp.cpu_count() // 2 if mp.cpu_count() // 2 < 2 else 2
  29. return num_workers