clean up

pyproject.toml

@@ -76,7 +76,6 @@ docstring-code-format = true
 [tool.ruff.lint.isort]
 combine-as-imports = true
 force-wrap-aliases = true
-known-local-folder = ["pi-tagger"]
 known-first-party = ["pi-tagger"]
 
 

tagger/common.py

@@ -1,9 +1,7 @@
-import json
-from dataclasses import asdict, dataclass
+from dataclasses import dataclass
 from functools import lru_cache
-from os import PathLike
 from pathlib import Path
-from typing import Any, Optional
+from typing import Optional
 
 import numpy as np
 import pandas as pd
@@ -12,16 +10,8 @@ from huggingface_hub.utils import HfHubHTTPError
 from PIL import Image
 
 
-class DictJsonMixin:
-    def asdict(self, *args, **kwargs) -> dict[str, Any]:
-        return asdict(self, *args, **kwargs)
-
-    def asjson(self, *args, **kwargs):
-        return json.dumps(asdict(self, *args, **kwargs))
-
-
 @dataclass
-class LabelData(DictJsonMixin):
+class LabelData:
     names: list[str]
     rating: list[np.int64]
     general: list[np.int64]
@@ -29,7 +19,7 @@ class LabelData(DictJsonMixin):
 
 
 @dataclass
-class ImageLabels(DictJsonMixin):
+class ImageLabels:
     caption: str
     booru: str
     rating: dict[str, float]
@@ -37,24 +27,6 @@ class ImageLabels(DictJsonMixin):
     character: dict[str, float]
 
 
-@lru_cache(maxsize=5)
-def load_labels(version: str = "v3", data_dir: PathLike = "./data") -> LabelData:
-    data_dir = Path(data_dir).resolve()
-    csv_path = data_dir.joinpath(f"selected_tags_{version}.csv")
-    if not csv_path.is_file():
-        raise FileNotFoundError(f"{csv_path.name} not found in {data_dir}")
-
-    df: pd.DataFrame = pd.read_csv(csv_path, usecols=["name", "category"])
-    tag_data = LabelData(
-        names=df["name"].tolist(),
-        rating=list(np.where(df["category"] == 9)[0]),
-        general=list(np.where(df["category"] == 0)[0]),
-        character=list(np.where(df["category"] == 4)[0]),
-    )
-
-    return tag_data
-
-
 @lru_cache(maxsize=5)
 def load_labels_hf(
     repo_id: str,