Update app.py

app.py

@@ -121,21 +121,65 @@ class Predictor:
             return
 
         csv_path, model_path = self.download_model(model_repo)
-
+        
+        # タグデータのロード最適化
         tags_df = pd.read_csv(csv_path)
-        sep_tags = load_labels(tags_df)
-
-        self.tag_names = sep_tags[0]
-        self.rating_indexes = sep_tags[1]
-        self.general_indexes = sep_tags[2]
-        self.character_indexes = sep_tags[3]
-
-        model = rt.InferenceSession(model_path)
-        _, height, width, _ = model.get_inputs()[0].shape
+        self.tag_names = tags_df["name"].tolist()
+        self.tag_names = [x.replace("_", " ") if x not in kaomojis else x for x in self.tag_names]
+        
+        # カテゴリインデックスの効率的な抽出
+        categories = tags_df["category"].to_numpy()
+        self.rating_indexes = np.where(categories == 9)[0].tolist()
+        self.general_indexes = np.where(categories == 0)[0].tolist()
+        self.character_indexes = np.where(categories == 4)[0].tolist()
+
+        # ONNX実行時の最適化オプションを設定
+        sess_options = rt.SessionOptions()
+        sess_options.graph_optimization_level = rt.GraphOptimizationLevel.ORT_ENABLE_ALL
+        sess_options.enable_mem_pattern = True
+        sess_options.enable_cpu_mem_arena = True
+        
+        # マルチスレッド設定
+        sess_options.intra_op_num_threads = 4
+        sess_options.inter_op_num_threads = 2
+
+        # 最適化されたモデルロード
+        self.model = rt.InferenceSession(model_path, sess_options)
+        _, height, width, _ = self.model.get_inputs()[0].shape
         self.model_target_size = height
 
         self.last_loaded_repo = model_repo
-        self.model = model
+
+    def prepare_image_optimized(self, image):
+        target_size = self.model_target_size
+
+        # メモリ効率を高めるためにRGBAからRGBへの変換を最適化
+        if image.mode == 'RGBA':
+            canvas = Image.new("RGB", image.size, (255, 255, 255))
+            canvas.paste(image, mask=image.split()[3])
+            image = canvas
+        elif image.mode != 'RGB':
+            image = image.convert('RGB')
+
+        # 正方形パディングの最適化
+        image_shape = image.size
+        max_dim = max(image_shape)
+        
+        # リサイズが必要な場合のみパディングを適用
+        if image_shape[0] != image_shape[1]:
+            pad_left = (max_dim - image_shape[0]) // 2
+            pad_top = (max_dim - image_shape[1]) // 2
+            padded_image = Image.new("RGB", (max_dim, max_dim), (255, 255, 255))
+            padded_image.paste(image, (pad_left, pad_top))
+            image = padded_image
+
+        # リサイズの最適化 - 必要な場合のみ実行
+        if max_dim != target_size:
+            image = image.resize((target_size, target_size), Image.BICUBIC)
+
+        # NumPy配列への変換とBGR変換を一度に行う
+        image_array = np.asarray(image, dtype=np.float32)[:, :, ::-1]
+        return np.expand_dims(image_array, axis=0)
 
     def prepare_image(self, image):
         target_size = self.model_target_size
@@ -179,7 +223,7 @@ class Predictor:
     ):
         self.load_model(model_repo)
 
-        image = self.prepare_image(image)
+        image = self.prepare_image_optimized(image)
 
         input_name = self.model.get_inputs()[0].name
         label_name = self.model.get_outputs()[0].name