From 88142677254dcd24cb2d4ebad8e1b77e65d533be Mon Sep 17 00:00:00 2001
From: sml2h3 <sml2h3@gmail.com>
Date: Mon, 21 Feb 2022 02:53:47 +0800
Subject: [PATCH] beta0.1

---
 configs/base.py     |  4 ++--
 utils/cache_data.py |  8 ++++----
 utils/load_cache.py | 19 ++++++++++---------
 utils/train.py      |  4 +++-
 4 files changed, 19 insertions(+), 16 deletions(-)

diff --git a/configs/base.py b/configs/base.py
index 6752cd2..6728858 100644
--- a/configs/base.py
+++ b/configs/base.py
@@ -36,8 +36,8 @@ class Config(object):
                 "SAVE_CHECKPOINTS_STEP": 2000,
                 "TARGET": {
                     "Accuracy": 0.97,
-                    "Epoch": 200,
-                    "Cost": 0.005
+                    "Epoch": 20,
+                    "Cost": 0.05
                 },
                 "LR": 0.01
             }
diff --git a/utils/cache_data.py b/utils/cache_data.py
index 2ed059e..474c028 100644
--- a/utils/cache_data.py
+++ b/utils/cache_data.py
@@ -55,12 +55,10 @@ class CacheData:
         logger.info("\nChecking labels.txt ...")
         error_files = set(labels_filename_lines).difference(set(files))
         logger.info("\nCheck labels.txt end! {} errors!".format(len(error_files)))
-        for ef in error_files:
-            labels_lines.remove(ef)
         del files
-        self.__collect_data(labels_lines, images_path, is_file=True)
+        self.__collect_data(labels_lines, images_path, error_files, is_file=True)
 
-    def __collect_data(self, lines, base_path, is_file=False):
+    def __collect_data(self, lines, base_path, error_files, is_file=False):
         labels = []
         caches = []
 
@@ -72,6 +70,8 @@ class CacheData:
             else:
                 filename = file
                 label = "_".join(filename.split("_")[:-1])
+            if filename in error_files:
+                continue
             label = label.replace(" ", "")
             if filename.split('.')[-1] in self.allow_ext:
                 if " " in filename:
diff --git a/utils/load_cache.py b/utils/load_cache.py
index 18f86f6..ae7846e 100644
--- a/utils/load_cache.py
+++ b/utils/load_cache.py
@@ -3,14 +3,16 @@ import os
 
 import torch
 import tqdm
-import numpy as np
 
 from configs import Config
 from loguru import logger
 
 import torchvision
+from PIL import Image, ImageFile
 from torch.utils.data import DataLoader, Dataset, TensorDataset
 
+ImageFile.LOAD_TRUNCATED_IMAGES = True
+
 
 class LoadCache(Dataset):
     def __init__(self, cache_path: str, path: str, word: bool, image_channel: int, resize: list, charset: list):
@@ -48,20 +50,19 @@ class LoadCache(Dataset):
             else:
                 image_label = [image_label]
             if self.ImageChannel == 1:
-                mode = torchvision.io.ImageReadMode.GRAY
+                mode = "L"
             else:
-                mode = torchvision.io.ImageReadMode.RGB
-            image = torchvision.io.read_image(image_path, mode=mode)  # shape c, h, w
-            image_shape = image.shape
+                mode = "RGB"
+            image = Image.open(image_path).convert(mode)  # shape c, h, w
+            image_shape = image.size
             image_height = image_shape[1]
-            image_width = image_shape[2]
+            image_width = image_shape[0]
             width = self.resize[0]
             height = self.resize[1]
             if self.resize[0] == -1:
-                image = torchvision.transforms.Resize((height, int(image_width * (height / image_height))))(image)
+                image = image.resize((int(image_width * (height / image_height)), height))
             else:
-                image = torchvision.transforms.Resize((height, width))(image)
-            image = torchvision.transforms.ToPILImage()(image)
+                image = image.resize((width, height))
             label = [int(self.charset.index(item)) for item in list(image_label)]
             return image, label
 
diff --git a/utils/train.py b/utils/train.py
index 776b40b..795bbbd 100644
--- a/utils/train.py
+++ b/utils/train.py
@@ -71,8 +71,10 @@ class Train:
         loaders = load_cache.GetLoader(project_name)
         self.train = loaders.loaders['train']
         self.val = loaders.loaders['val']
+        del loaders
         logger.info("\nGet Data Loader End!")
 
+
         self.loss = 0
         self.avg_loss = 0
         self.start_time = time.time()
@@ -139,7 +141,7 @@ class Train:
                         self.net.export_onnx(self.net, dummy_input,
                                              os.path.join(self.models_path, "{}_{}_{}_{}_{}.onnx".format(
                                                  self.project_name, str(accuracy), self.epoch, self.step,
-                                                 time.localtime(self.now_time)))
+                                                 time.strftime("%Y-%m-%d-%H-%M-%S", time.localtime(self.now_time))))
                                              , input_names, output_names, dynamic_ax)
                         with open(os.path.join(self.models_path, "charset.json"), 'w', encoding="utf-8") as f:
                             f.write(json.dumps(self.net.charset, ensure_ascii=False))