0.3.0更新 (#10)

* 更新README
* 新增AI放大，支持多种算法
* 更改启动方式，现在可以带参数启动了
* 新增install.py，现在可以作为库来引入你的项目
* 修复若干BUG

README.md

@@ -4,12 +4,14 @@
 
 ✅批量处理图片  
 包括且不限于：
+
 * 批量翻转
 * 批量随机裁切
 * 图片对比度增强
 
 ✅批量处理标签  
 包括且不限于：
+
 * 批量删除标签
 * 批量插入标签
 * 批量修改标签
@@ -20,14 +22,31 @@
 
 ✅自动打标（试验性）
 
+✅AI图片放大（试验性）
+
 ✅子处理（试验性）
 
 ## TODO
 
-🚧重构自动打标代码
+🚧智能裁切
 
-🚧图片放大功能
+🚧重构自动打标代码
 
 🚧图形化界面
 
-如果这个项目给您提供了帮助，不妨点一个⭐star，谢谢
+## CREADIT
+
+### Upscale
+
+[Real ESRGAN](https://github.com/xinntao/Real-ESRGAN/)：一种流行的AI放大方案
+
+[Real CUGAN](https://github.com/bilibili/ailab/tree/main/Real-CUGAN)：更加适合二次元的AI放大方案
+
+[Real CUGAN-ncnn](https://github.com/Tohrusky/realcugan-ncnn-py)：感谢这位作者提供的RealCUGAN工具包
+
+### Tagger
+
+[WD-1.4-Tagger From SmilingWolf](https://huggingface.co/SmilingWolf)：自动打标模型
+
+---
+如果这个项目为您提供了帮助，不妨点一个⭐star，万分感谢！

module/__init__.py → dataset_processor/__init__.py

@@ -1,4 +1,5 @@
 from .data import Data
 from .filter import Filter
-from .processor import Processor,ProcessorError
+from .processor import Processor, ProcessorError
 from .uitl import *
+from .tools import *

dataset_processor/data.py

@@ -0,0 +1,55 @@
+from PIL import Image
+import os
+
+
+class Data:
+
+    # 图片读取并初始化
+    def __init__(self, path: str, name: str, ext: str):
+        self.token: list[str] = []
+        self.conduct = ""
+        self.repeat = 0
+        self.id = 0
+        self.name = name
+        self.ext = ext
+        self.path = path
+        # 读取图片
+        self.img = Image.open(os.path.join(path, name + ext))
+        self.size = self.img.size
+
+    # 载入标签
+    def input_token(self, file_name: str, option=None):
+        clean_tag = False
+        NO_CHECK = [  # 清洗排除标签
+            ':)', ';)', ':(', '>:)', '>:(', '\\(^o^)/',  # 括号相关
+            '^_^', '@_@', '>_@', '+_+', '+_-', 'o_o', '0_0', '|_|', '._.', '>_<', '=_=', '<o>_<o>', '<|>_<|>'  # 下划线相关
+        ]
+        if option.clean_tag:
+            clean_tag = True
+        with open(os.path.join(self.path, file_name), "r") as f:
+            self.token = f.read(-1).split(",")
+        for tag in self.token:
+            tag = tag.strip()
+            if clean_tag:
+                if tag not in NO_CHECK:
+                    tag = tag.replace("_", " ")
+                    tag = tag.replace("(", "\\(")
+                    tag = tag.replace(")", "\\)")
+
+    # 保存的方法
+    def save(self, output_dir, option):
+        # 默认命名方式：id_conduct_repeat.ext 比如"000001_r_0.jpg"
+        save_name = str(self.id).zfill(6) + self.conduct
+        if option:
+            if option.save_source_name or option.save_conduct_id:
+                save_name = str(self.id).zfill(6)
+                if option.save_source_name:
+                    save_name = save_name.join('_' + self.name)
+                if option.save_conduct_id:
+                    save_name = save_name.join(self.conduct)
+        self.img.save(os.path.join(output_dir, save_name + self.ext))
+        # print(save_name)
+        with open(os.path.join(output_dir, save_name + ".txt"), mode="w") as f:
+            text = ",".join(self.token)
+            f.write(text)
+        self.img.close()

dataset_processor/filter.py

@@ -0,0 +1,48 @@
+from .data import Data
+
+
+class Filter:
+    """
+    这是一个过滤器类，包含所有有关数据过滤的函数
+    编写规范如下：
+    def 过滤器名(data:Data,arg)->bool:
+        #代码块
+        return bool
+    其中，True表示该数据会被过滤，False则会被保留
+    """
+
+    def img_size(data: Data, size: list) -> bool:
+        min, max = tuple(size)
+        x, y = data.size
+        if min != -1:
+            if data.size[0] <= min or data.size[1] <= min:
+                return True
+        if max != -1:
+            if data.size[0] > max and data.size[1] > max:
+                return True
+        else:
+            return False
+
+    def tag_filter(data: Data, tag) -> bool:
+        if tag in data.token:
+            return True
+        else:
+            return False
+
+    def tag_selector(data: Data, tag) -> bool:
+        if tag in data.token:
+            return False
+        else:
+            return True
+
+    def tag_is_not_none(data: Data) -> bool:
+        if data.token:
+            return False
+        else:
+            return True
+
+    def tag_is_none(data: Data) -> bool:
+        if data.token:
+            return True
+        else:
+            return False

module/processor.py → dataset_processor/processor.py

@@ -1,87 +1,86 @@
 from random import randint as random
 from .data import Data
+from .tools.tagger import Tagger
+from .tools.upscale import UpscaleModel
 from PIL import Image
 from PIL import ImageEnhance
 import numpy as np
 
+
 class Processor:
-    # 在这里定义处理方法
     """
+    这是一个处理器类，包含所有有关数据处理的函数
     编写规范如下:
-    def 处理名(data:Data,args):
+    def 处理名(data:Data,args)->Data:
         #代码块
         return data
     """
 
-    def random_crop(data, size):
+    def random_crop(data: Data, size) -> Data:
         if not (data.size[0] <= size or data.size[1] <= size):
             x = random(1, data.size[0] - size)
             y = random(1, data.size[1] - size)
             box = (x, y, x + size, y + size)
             data.img = data.img.crop(box)
-            data.conduct += "_rc"
+            data.conduct += f"_rc{data.repeat}"
             data.size = data.img.size
         else:
             raise ImageTooSmallError(data.name + data.ext)
         return data
 
-    def flip(data):
+    def flip(data: Data) -> Data:
         data.img = data.img.transpose(Image.FLIP_LEFT_RIGHT)
-        data.conduct += "_f"
+        data.conduct += f"_f{data.repeat}"
         return data
 
-    def resize(data: Data, proportion: float):
-        data.size = (int(data.size[0] * proportion), int(data.size[1] * proportion))
-        data.img = data.img.resize(data.size)
-        data.conduct += "_r"
+    def resize(data: Data, proportion: float) -> Data:
+        size = (int(data.size[0] * proportion), int(data.size[1] * proportion))
+        data.img = data.img.resize(size)
+        data.conduct += f"_r{data.repeat}"
+        data.size = data.img.size
         return data
 
-    def force_resize(data: Data, size: list):
-        data.size = (size[0], size[1])
-        data.img = data.img.resize(data.size)
-        data.conduct += "_fr"
+    def force_resize(data: Data, size: list) -> Data:
+        data.img = data.img.resize(size)
+        data.conduct += f"_fr{data.repeat}"
+        data.size = data.img.size
         return data
-    
-    def offset(data: Data,offset:int):
-        data.img = data.img.offset(offset,0)
-        data.conduct += "_off"
+
+    def offset(data: Data, offset: int) -> Data:
+        data.img = data.img.offset(offset, 0)
+        data.conduct += f"_off{data.repeat}"
         return data
-    
-    def rotation(data: Data, rot:int):
+
+    def rotation(data: Data, rot: int) -> Data:
         data.img = data.img.rotate(rot)
-        data.conduct += "_rot"
+        data.conduct += f"_rot{data.repeat}"
         return data
-    
-    def contrast_enhancement(data: Data): #对比度增强
+
+    def contrast_enhancement(data: Data) -> Data:  # 对比度增强
         image = data.img
         enh_con = ImageEnhance.Contrast(image)
         contrast = 1.5
         data.img = enh_con.enhance(contrast)
-        data.conduct += "_con_e"
+        data.conduct += f"_con_e{data.repeat}"
         return data
-    
-    def brightness_enhancement(data: Data):#亮度增强
+
+    def brightness_enhancement(data: Data) -> Data:  # 亮度增强
         image = data.img
         enh_bri = ImageEnhance.Brightness(image)
         brightness = 1.5
         data.img = enh_bri.enhance(brightness)
-        data.conduct += "_bri_e"
+        data.conduct += f"_bri_e{data.repeat}"
         return data
 
-    def color_enhancement(data: Data):#颜色增强
+    def color_enhancement(data: Data) -> Data:  # 颜色增强
         image = data.img
         enh_col = ImageEnhance.Color(image)
         color = 1.5
         data.img = enh_col.enhance(color)
         data.conduct += "_col_e"
         return data
-
-    def random_enhancement(data: Data): #随机抖动
-        """
-        对图像进行颜色抖动
-        :param image: PIL的图像image
-        :return: 有颜色色差的图像image
-        """
+
+    def random_enhancement(data: Data) -> Data:  # 随机抖动
         image = data.img
         random_factor = np.random.randint(8, 31) / 10.  # 随机因子
         color_image = ImageEnhance.Color(image).enhance(random_factor)  # 调整图像的饱和度
@@ -91,55 +90,64 @@ def random_enhancement(data: Data): #随机抖动
         contrast_image = ImageEnhance.Contrast(brightness_image).enhance(random_factor)  # 调整图像对比度
         random_factor = np.random.randint(8, 20) / 10.  # 随机因子
         data.img = ImageEnhance.Sharpness(contrast_image).enhance(random_factor)  # 调整图像锐度
-        data.conduct += "_ran_e"
+        data.conduct += f"_ran_e{data.repeat}"
         return data
 
-    def none(data: Data):
+    def none(data: Data) -> Data:
         """
         无操作，主要用于一些特殊场景
         """
         return data
 
-    def append_tag(data: Data, tag: str):
+    def append_tag(data: Data, tag: str) -> Data:
         data.token.append(tag)
         return data
 
-    def remove_tag(data: Data, tag: str):
+    def remove_tag(data: Data, tag: str) -> Data:
         if tag in data.token:
             data.token.remove(tag)
         else:
-            raise TagNotExistError(tag,data.name + data.ext)
+            raise TagNotExistError(tag, data.name + data.ext)
         return data
 
-    def insert_tag(data: Data, tag: str):
+    def insert_tag(data: Data, tag: str) -> Data:
         data.token.insert(0, tag)
         return data
 
-    def tag_move_forward(data: Data,tag:str):
+    def tag_move_forward(data: Data, tag: str) -> Data:
         """
         将匹配项放到开头
         """
         if tag in data.token:
             data.token.remove(tag)
         else:
-            raise TagNotExistError(tag,data.name + data.ext)
+            raise TagNotExistError(tag, data.name + data.ext)
         data.token.insert(0, tag)
         return data
-    
-    def rename_tag(data:Data,tags:list[str]):
+
+    def rename_tag(data: Data, tags: list[str]) -> Data:
         """
         将Atag改名为Btag
         """
         tag_a = tags[0]
         tag_b = tags[1]
         if tag_a in data.token:
             index = data.token.index(tag_a)
-            data.token.insert(index,tag_b)
+            data.token.insert(index, tag_b)
             data.token.remove(tag_a)
         else:
-            raise TagNotExistError(tag_a,data.name + data.ext)
+            raise TagNotExistError(tag_a, data.name + data.ext)
         return data
 
+    def tag_image(data: Data, tagger: Tagger):
+        return tagger.tag_data(data)
+
+    def upscale_image(data: Data, upscale: UpscaleModel):
+        data.img = upscale.upscale_data(data)
+        data.size = data.img.size
+        return data
+
+
 # 自定义异常
 class ProcessorError(RuntimeError):
     def __init__(self, *args: object) -> None:
@@ -150,6 +158,7 @@ class ImageTooSmallError(ProcessorError):
     def __init__(self, name: str):
         print("image " + name + " is too small!")
 
+
 class TagNotExistError(ProcessorError):
-    def __init__(self,tag,name: str):
-        print("Tag"+ tag + "not exist in"+name+"!")
+    def __init__(self, tag, name: str):
+        print("Tag" + tag + "not exist in" + name + "!")

dataset_processor/tools/__init__.py

@@ -0,0 +1,2 @@
+from .tagger import Tagger, TaggerOption
+from .upscale import UpscaleModel, UpcaleOption