diff --git a/modules/fc_kks.py b/modules/fc_kks.py index 385ba9b..f1455a4 100644 --- a/modules/fc_kks.py +++ b/modules/fc_kks.py @@ -1,17 +1,10 @@ import shutil -from enum import Enum + from pathlib import Path -from typing import Literal +from util.classifier import CardType, get_card_type from util.logger import logger -class CardType(Enum): - UNKNOWN = "UNKNOWN" - KK = "KK" - KKSP = "KKSP" - KKS = "KKS" - - class FilterConvertKKS: def __init__(self, config, file_manager): """Initializes the FilterConvertKKS module. @@ -23,29 +16,19 @@ class FilterConvertKKS: self.file_manager = file_manager self.convert = self.config.fc_kks["Convert"] - def get_list(self, folder_path: Path) -> list[str]: + def get_list(self, folder_path: Path) -> list[Path]: """Get list of PNG files in the folder.""" - folder = Path(folder_path) - return [str(file) for file in folder.rglob("*.png")] + return [file for file in folder_path.rglob("*.png")] def check_png(self, card_path: Path) -> CardType: """Check the PNG file and return its type.""" - card_path = Path(card_path) - with card_path.open("rb") as card: - data = card.read() - card_type = CardType.UNKNOWN - if b"KoiKatuChara" in data: - card_type = CardType.KK - if b"KoiKatuCharaSP" in data: - card_type = CardType.KKSP - elif b"KoiKatuCharaSun" in data: - card_type = CardType.KKS - logger.info(f"{card_type.value}", f"{card_path.name}") + card_type = get_card_type(card_path) + logger.info(f"{card_type.value}", f"{card_path.name}") return card_type def convert_kk(self, card_name: str, card_path: Path, destination_path: Path): """Convert KKS card to KK.""" - card_path = Path(card_path) # Convert to Path object + card_path = Path(card_path) with card_path.open(mode="rb") as card: data = card.read() @@ -74,7 +57,7 @@ class FilterConvertKKS: count = len(png_list) if count > 0: - logger.info("SCRIPT", "kk: Koikatsu / KKSP: Koikatsu Special / KKS: Koikatsu Sunshine") + logger.info("SCRIPT", "KK: Koikatsu / KKSP: Koikatsu Special / KKS: Koikatsu Sunshine") logger.line() logger.info("FOLDER", str(path)) for png in png_list: @@ -87,7 +70,6 @@ class FilterConvertKKS: count = len(kks_card_list) if count > 0: - print(kks_card_list) # Create target directories if they don't exist kks_folder.mkdir(exist_ok=True) diff --git a/modules/install_chara.py b/modules/install_chara.py index 21bcc8c..94caada 100644 --- a/modules/install_chara.py +++ b/modules/install_chara.py @@ -1,4 +1,5 @@ from pathlib import Path +from util.classifier import CardType, get_card_type, is_male, is_coordinate from util.logger import logger from typing import Optional @@ -16,18 +17,24 @@ class InstallChara: self.input_path = Path(self.config.install_chara["InputPath"]) def resolve_png(self, image_path: Path): - with image_path.open("rb") as card: - data = card.read() - if b"KoiKatuChara" in data: - if b"KoiKatuCharaSP" in data or b"KoiKatuCharaSun" in data: - basename = Path(image_path).name - logger.error("CHARA", f"{basename} is a KKS card") - return - self.file_manager.copy_and_paste("CHARA", image_path, self.game_path["chara"]) - elif b"KoiKatuClothes" in data: - self.file_manager.copy_and_paste("COORD", image_path, self.game_path["coordinate"]) - else: - self.file_manager.copy_and_paste("OVERLAYS", image_path, self.game_path["Overlays"]) + image_bytes = image_path.read_bytes() + card_type = get_card_type(image_bytes) + + match card_type: + case CardType.KK: + if is_male(image_bytes): + self.file_manager.copy_and_paste("CHARA M", image_path, self.game_path["charaMale"]) + else: + self.file_manager.copy_and_paste("CHARA F", image_path, self.game_path["charaFemale"]) + + case CardType.KKS | CardType.KKSP: + logger.error("CHARA", f"{image_path.name} is a {card_type.value} card") + + case CardType.UNKNOWN: + if is_coordinate(image_bytes): + self.file_manager.copy_and_paste("COORD", image_path, self.game_path["coordinate"]) + else: + self.file_manager.copy_and_paste("OVERLAYS", image_path, self.game_path["Overlays"]) def run(self, folder_path: Optional[Path] = None): if folder_path is None: diff --git a/modules/remove_chara.py b/modules/remove_chara.py index d9ff49f..4ff864e 100644 --- a/modules/remove_chara.py +++ b/modules/remove_chara.py @@ -1,4 +1,5 @@ from pathlib import Path +from util.classifier import CardType, get_card_type, is_male, is_coordinate from util.logger import logger @@ -15,16 +16,21 @@ class RemoveChara: self.input_path = self.config.remove_chara["InputPath"] def resolve_png(self, image_path: Path): - with image_path.open("rb") as card: - data = card.read() - if b"KoiKatuChara" in data: - if b"KoiKatuCharaSP" in data or b"KoiKatuCharaSun" in data: - return - self.file_manager.find_and_remove("CHARA", image_path, self.game_path["chara"]) - elif b"KoiKatuClothes" in data: - self.file_manager.find_and_remove("COORD", image_path, self.game_path["coordinate"]) - else: - self.file_manager.find_and_remove("OVERLAYS", image_path, self.game_path["Overlays"]) + image_bytes = image_path.read_bytes() + card_type = get_card_type(image_bytes) + + match card_type: + case CardType.KK: + if is_male(image_bytes): + self.file_manager.find_and_remove("CHARA M", image_path, self.game_path["charaMale"]) + else: + self.file_manager.find_and_remove("CHARA F", image_path, self.game_path["charaFemale"]) + + case CardType.UNKNOWN: + if is_coordinate(image_bytes): + self.file_manager.find_and_remove("COORD", image_path, self.game_path["coordinate"]) + else: + self.file_manager.find_and_remove("OVERLAYS", image_path, self.game_path["Overlays"]) def run(self): foldername = self.input_path.name diff --git a/test.py b/test.py index f7559b6..e638369 100644 --- a/test.py +++ b/test.py @@ -34,6 +34,7 @@ paths_to_check = [ GAMEPATH / "UserData", GAMEPATH / "BepInEx", GAMEPATH / "mods", + GAMEPATH / "UserData" / "chara" / "male", GAMEPATH / "UserData" / "chara" / "female", GAMEPATH / "UserData" / "coordinate", GAMEPATH / "UserData" / "Overlays" @@ -192,8 +193,8 @@ def test_remove(): def main(): """Main function to run all tests.""" - kk_cards = download('https://db.bepis.moe/koikatsu?type=base') - kks_cards = download('https://db.bepis.moe/koikatsu?type=sunshine') + kk_cards = download('https://db.bepis.moe/koikatsu?type=steam&orderby=popularity') + kks_cards = download('https://db.bepis.moe/koikatsu?type=sunshine&orderby=popularity') if kk_cards and kks_cards: test_fckss(kk_cards, kks_cards) diff --git a/util/classifier.py b/util/classifier.py new file mode 100644 index 0000000..e420fb4 --- /dev/null +++ b/util/classifier.py @@ -0,0 +1,32 @@ +from enum import Enum +from pathlib import Path +from typing import Union + + +class CardType(Enum): + UNKNOWN = "UNKNOWN" + KK = "KK" + KKSP = "KKSP" + KKS = "KKS" + + +def get_card_type(card: Union[str, Path, bytes]): + if isinstance(card, (str, Path)): + card = Path(card).read_bytes() + + card_type = CardType.UNKNOWN + if b"KoiKatuChara" in card: + card_type = CardType.KK + if b"KoiKatuCharaSP" in card: + card_type = CardType.KKSP + elif b"KoiKatuCharaSun" in card: + card_type = CardType.KKS + return card_type + + +def is_male(image_bytes: bytes): + return b'sex\x00' in image_bytes + + +def is_coordinate(image_bytes: bytes): + return 'b"KoiKatuClothes"' in image_bytes \ No newline at end of file diff --git a/util/config.py b/util/config.py index 96e894d..721942f 100644 --- a/util/config.py +++ b/util/config.py @@ -46,7 +46,8 @@ class Config: "UserData": base / "UserData", "BepInEx": base / "BepInEx", "mods": base / "mods", - "chara": base / "UserData" / "chara" / "female", + "charaMale": base / "UserData" / "chara" / "male", + "charaFemale": base / "UserData" / "chara" / "female", "coordinate": base / "UserData" / "coordinate", "Overlays": base / "UserData" / "Overlays" }