dev(narugo): add resource bg (7d1b655f) · Commits · git-mirror / Imgutils

docs/source/api_doc/resource/background.rst

0 → 100644

+51 −0

Original line number	Diff line number	Diff line
		imgutils.resource.background
		=======================================

		.. currentmodule:: imgutils.resource.background

		.. automodule:: imgutils.resource.background


		BackgroundImageSet
		------------------------------------------

		.. autoclass:: BackgroundImageSet
		:members: list_image_files, get_image_file, get_image, random_image_file, random_image



		list_bg_image_files
		------------------------------------------

		.. autofunction:: list_bg_image_files



		get_bg_image_file
		------------------------------------------

		.. autofunction:: get_bg_image_file



		get_bg_image
		------------------------------------------

		.. autofunction:: get_bg_image



		random_bg_image_file
		------------------------------------------

		.. autofunction:: random_bg_image_file



		random_bg_image
		------------------------------------------

		.. autofunction:: random_bg_image

docs/source/api_doc/resource/index.rst

0 → 100644

+14 −0

Original line number	Diff line number	Diff line
		imgutils.resource
		=====================

		.. currentmodule:: imgutils.resource

		.. automodule:: imgutils.resource


		.. toctree::
		:maxdepth: 3

		background

docs/source/index.rst

+1 −0

Original line number	Diff line number	Diff line
		@@ -37,6 +37,7 @@ configuration file's structure and their versions.
		api_doc/ocr/index
		api_doc/operate/index
		api_doc/pose/index
		api_doc/resource/index
		api_doc/restore/index
		api_doc/sd/index
		api_doc/segment/index

imgutils/resource/background.py

+136 −17

Original line number	Diff line number	Diff line
		@@ -25,6 +25,12 @@ _BG_REPO = 'deepghs/anime-bg'

		@lru_cache()
		def _global_df() -> pd.DataFrame:
		"""
		Load the global dataframe containing information about background images.

		:return: The global dataframe containing information about background images.
		:rtype: pd.DataFrame
		"""
		return pd.read_csv(hf_hub_download(
		repo_id=_BG_REPO,
		repo_type='dataset',
		@@ -34,69 +40,140 @@ def _global_df() -> pd.DataFrame:

		@lru_cache()
		def _bg_root_dir() -> str:
		"""
		Get the root directory for storing background images.

		:return: The root directory for storing background images.
		:rtype: str
		"""
		root = os.path.join(get_storage_dir(), 'bg')
		os.makedirs(root, exist_ok=True)
		return root


		class BackgroundImageSet:
		def __init__(self, width: Optional[int] = None, height: Optional[int] = None,
		scale_only: bool = False, strict_level: float = 1.0, min_range: int = 5,
		min_width: Optional[int] = None, min_height: Optional[int] = None):
		def __init__(self, width: Optional[float] = None, height: Optional[float] = None,
		strict_level: float = 1.5, min_selected: int = 5,
		min_width: Optional[int] = None, min_height: Optional[int] = None,
		min_resolution: Optional[int] = None):
		"""
		Initialize a BackgroundImageSet instance.

		:param width: The desired width of background images. (default: None)
		:type width: Optional[float]

		:param height: The desired height of background images. (default: None)
		:type height: Optional[float]

		:param strict_level: The strictness level for selecting images. (default: 1.0)
		:type strict_level: float

		:param min_selected: The minimum number of images to consider for selection. (default: 5)
		:type min_selected: int

		:param min_width: The minimum width of background images to consider. (default: None)
		:type min_width: Optional[int]

		:param min_height: The minimum height of background images to consider. (default: None)
		:type min_height: Optional[int]

		:param min_resolution: The minimum resolution of background images to consider. (default: None)
		:type min_resolution: Optional[int]
		"""
		df = _global_df().copy()
		if min_width:
		df = df[df['width'] >= min_width]
		if min_height:
		df = df[df['height'] >= min_height]
		if min_resolution:
		df = df[(df['height'] * df['width']) >= min_resolution ** 2]

		if width and height:
		r1, r2 = width / height, height / width
		df['d1'] = (df['width'] / df['height'] - r1).abs()
		df['d2'] = (df['height'] / df['width'] - r2).abs()
		df['dx'] = df['d1'] + df['d2']
		if scale_only:
		df['dx'] = (df['width'] / df['height'] - r1).abs() + (df['height'] / df['width'] - r2).abs()
		df = df.sort_values(['dx'], ascending=[True])
		else:
		df['dy'] = ((df['width'] * df['height']) ** 0.5 - (width * height) ** 0.5).abs()
		df = df.sort_values(['dx', 'dy'], ascending=[True, True])
		df['i'] = range(len(df))
		idx = df['dx'] < (df['dx'].mean() - df['dx'].std() * strict_level)
		idx = idx \| (df['i'] < max(min_range, 1))
		idx = idx \| (df['i'] < max(min_selected, 1))
		idx = idx \| (df['dx'] <= df[idx]['dx'].max())
		df = df[idx]
		elif width and not height:
		df['dist'] = (df['width'] - width).abs()
		df = df.sort_values(['dist'], ascending=[True])
		df['i'] = range(len(df))
		idx = df['dist'] < (df['dist'].mean() - df['dist'].std() * strict_level)
		idx = idx \| (df['i'] < max(min_range, 1))
		idx = idx \| (df['i'] < max(min_selected, 1))
		idx = idx \| (df['dist'] <= df[idx]['dist'].max())
		df = df[idx]
		elif not width and height:
		df['dist'] = (df['height'] - height).abs()
		df = df.sort_values(['dist'], ascending=[True])
		df['i'] = range(len(df))
		idx = df['dist'] < (df['dist'].mean() - df['dist'].std() * strict_level)
		idx = idx \| (df['i'] < max(min_range, 1))
		idx = idx \| (df['i'] < max(min_selected, 1))
		idx = idx \| (df['dist'] <= df[idx]['dist'].max())
		df = df[idx]
		else:
		pass

		self.df = df
		if len(self.df) == 0:
		raise ValueError('No background images selected, please lower your settings.')
		self._map = {item['filename']: item for item in self.df.to_dict('records')}

		def list_image_files(self) -> List[str]:
		"""
		Get a list of filenames of background images.

		:return: A list of filenames of background images.
		:rtype: List[str]
		"""
		return self.df['filename'].tolist()

		def get_image_file(self, filename: str) -> str:
		"""
		Get the local file path of a background image by filename.

		:param filename: The filename of the background image.
		:type filename: str

		:return: The local file path of the background image.
		:rtype: str
		"""
		return self._load_local_image_file(filename)

		def get_image(self, filename: str) -> Image.Image:
		return load_image(self.get_image_file(filename))
		"""
		Get the PIL Image object of a background image by filename.

		:param filename: The filename of the background image.
		:type filename: str

		:return: The PIL Image object of the background image.
		:rtype: Image.Image
		"""
		return load_image(self.get_image_file(filename), mode='RGB')

		def _random_filename(self):
		return self.df['filename'].sample(n=1).tolist()[0]

		def random_image_file(self) -> str:
		return self.get_image_file(self.df['filename'].sample(n=1).tolist()[0])
		"""
		Get the filename of a randomly selected background image.

		:return: The filename of a randomly selected background image.
		:rtype: str
		"""
		return self.get_image_file(self._random_filename())

		def random_image(self) -> Image.Image:
		return load_image(self.random_image_file())
		"""
		Get the PIL Image object of a randomly selected background image.

		:return: The PIL Image object of a randomly selected background image.
		:rtype: Image.Image
		"""
		return self.get_image(self._random_filename())

		def _load_local_image_file(self, filename):
		filename = os.path.normcase(filename)
		@@ -120,24 +197,66 @@ class BackgroundImageSet:

		@lru_cache()
		def _get_default_set():
		"""
		Get the default BackgroundImageSet instance.

		:return: The default BackgroundImageSet instance.
		:rtype: BackgroundImageSet
		"""
		return BackgroundImageSet()


		def list_bg_image_files() -> List[str]:
		"""
		Get a list of filenames of background images.

		:return: A list of filenames of background images.
		:rtype: List[str]
		"""
		return _get_default_set().list_image_files()


		def get_bg_image_file(filename: str) -> str:
		"""
		Get the local file path of a background image by filename.

		:param filename: The filename of the background image.
		:type filename: str

		:return: The local file path of the background image.
		:rtype: str
		"""
		return _get_default_set().get_image_file(filename)


		def get_bg_image(filename) -> Image.Image:
		"""
		Get the PIL Image object of a background image by filename.

		:param filename: The filename of the background image.
		:type filename: str

		:return: The PIL Image object of the background image.
		:rtype: Image.Image
		"""
		return _get_default_set().get_image(filename)


		def random_bg_image_file() -> str:
		"""
		Get the filename of a randomly selected background image.

		:return: The filename of a randomly selected background image.
		:rtype: str
		"""
		return _get_default_set().random_image_file()


		def random_bg_image() -> Image.Image:
		"""
		Get the PIL Image object of a randomly selected background image.

		:return: The PIL Image object of a randomly selected background image.
		:rtype: Image.Image
		"""
		return _get_default_set().random_image()

test/resource/init.py

0 → 100644

+0 −0

Empty file added.