Update app.py

app.py

@@ -1,225 +1,278 @@
-import gradio as gr
-import pandas as pd
-import os
-import shutil
-import matplotlib.pyplot as plt
-import numpy as np
-import japanize_matplotlib
-
-
-class DataProcessor:
-   def __init__(self, bodypart_names, x_max, y_max):
-      self.bodypart_names = bodypart_names.split(',')
-      self.x_max = x_max
-      self.y_max = y_max
-      self.output_folder = 'output_plots'
-
-   def process_csv(self, file_path):
-      df = pd.read_csv(file_path, header=[1, 2])
-      df_likelihood = self.extract_likelihood(df)
-      df = self.remove_first_column_and_likelihood(df)
-      df = self.rename_bodyparts(df)
-      return df, df_likelihood
-
-   def remove_first_column_and_likelihood(self, df):
-      df = df.drop(df.columns[0], axis=1)
-      df = df[df.columns.drop(list(df.filter(regex='likelihood')))]
-      return df
-
-   def rename_bodyparts(self, df):
-      current_names = df.columns.get_level_values(0).unique()
-      if len(self.bodypart_names) != len(current_names):
-         raise ValueError(
-             "The length of bodypart_names must be equal to the number of bodyparts.")
-      mapping = dict(zip(current_names, self.bodypart_names))
-      new_columns = [(mapping[col[0]], col[1]) if col[0]
-                     in mapping else col for col in df.columns]
-      df.columns = pd.MultiIndex.from_tuples(new_columns)
-      return df
-
-   def extract_likelihood(self, df):
-      # likelihood列のみを抽出する
-      df = df[df.columns[df.columns.get_level_values(1) == 'likelihood']]
-      df.drop(df.columns[0], axis=1)
-      current_names = df.columns.get_level_values(0).unique()
-      mapping = dict(zip(current_names, self.bodypart_names))
-      new_columns = [(mapping[col[0]], col[1]) if col[0]
-                     in mapping else col for col in df.columns]
-      df.columns = pd.MultiIndex.from_tuples(new_columns)
-
-      return df
-
-   def plot_scatter(self, df):
-      if not os.path.exists(self.output_folder):
-         os.makedirs(self.output_folder)
-      return self.plot_scatter_fixed(df, self.output_folder, self.x_max, self.y_max)
-
-   def plot_scatter_fixed(self, df, output_folder, x_max, y_max):
-      image_paths = []
-      bodyparts = df.columns.get_level_values(0).unique()
-      colors = plt.cm.rainbow(np.linspace(0, 1, len(bodyparts)))
-      for i, bodypart in enumerate(bodyparts):
-         x = df[bodypart]['x'].values
-         y = df[bodypart]['y'].values
-         plt.figure(figsize=(8, 6))
-         plt.scatter(x, y, color=colors[i], label=bodypart)
-         plt.scatter(x[0], y[0], color='black', marker='o', s=100)
-         plt.text(x[0], y[0], ' Start', color='black',
-                  fontsize=12, verticalalignment='bottom')
-         plt.xlim(0, x_max)
-         plt.ylim(0, y_max)
-         plt.gca().invert_yaxis()
-         plt.title(f'トラッキングの座標({bodypart})')
-         plt.xlabel('X Coordinate(pixel)')
-         plt.ylabel('Y Coordinate(pixel)')
-         plt.legend(loc='upper right')
-         plt.savefig(f'{output_folder}/{bodypart}.png')
-         image_paths.append(f'{output_folder}/{bodypart}.png')
-         plt.close()
-
-      plt.figure(figsize=(8, 6))
-      for i, bodypart in enumerate(bodyparts):
-         x = df[bodypart]['x'].values
-         y = df[bodypart]['y'].values
-         plt.scatter(x, y, color=colors[i], label=bodypart)
-      plt.xlim(0, x_max)
-      plt.ylim(0, y_max)
-      plt.gca().invert_yaxis()
-      plt.title('トラッキングの座標(全付属肢)')
-      plt.xlabel('X Coordinate(pixel)')
-      plt.ylabel('Y Coordinate(pixel)')
-      plt.legend(loc='upper right')
-      plt.savefig(f'{output_folder}/all_plot.png')
-      image_paths.append(f'{output_folder}/all_plot.png')
-      plt.close()
-
-      return image_paths
-
-   def plot_trajectories(self, df):
-      image_paths = []
-      bodyparts = df.columns.get_level_values(0).unique()
-      colors = plt.cm.rainbow(np.linspace(0, 1, len(bodyparts)))
-
-      for i, bodypart in enumerate(bodyparts):
-         x = df[bodypart]['x'].values
-         y = df[bodypart]['y'].values
-         plt.figure(figsize=(8, 6))
-         plt.plot(x, color=colors[i], label=bodypart +
-                  "(x座標)", linestyle='dashed')
-         plt.plot(y, color=colors[i], label=bodypart + "(y座標)")
-         plt.title(f'トラッキングの座標({bodypart})')
-         plt.xlabel('Frames')
-         plt.ylabel('Coordinate(pixel)')
-         plt.legend(loc='upper right')
-         plt.savefig(f'{self.output_folder}/{bodypart}_trajectories.png')
-         image_paths.append(
-             f'{self.output_folder}/{bodypart}_trajectories.png')
-         plt.close()
-
-      plt.figure(figsize=(8, 6))
-      for i, bodypart in enumerate(bodyparts):
-         x = df[bodypart]['x'].values
-         y = df[bodypart]['y'].values
-         plt.plot(x, color=colors[i], label=bodypart +
-                  "(x座標)", linestyle='dashed')
-         plt.plot(y, color=colors[i], label=bodypart + "(y座標)")
-      plt.title(f'トラッキングの座標({bodypart})')
-      plt.xlabel('Frames')
-      plt.ylabel('Coordinate(pixel)')
-      plt.legend(loc='upper right')
-      plt.savefig(f'{self.output_folder}/all_trajectories.png')
-      image_paths.append(f'{self.output_folder}/all_trajectories.png')
-      plt.close()
-      return image_paths
-
-   def plot_likelihood(self, likelihood_df):
-      image_paths = []
-      plt.ylim(0, 1.0)
-      bodyparts = likelihood_df.columns.get_level_values(0).unique()
-      colors = plt.cm.rainbow(np.linspace(0, 1, len(bodyparts)))
-
-      # 付属肢ごとに尤度をプロット
-      for i, bodypart in enumerate(bodyparts):
-         plt.figure(figsize=(8, 6))
-         plt.ylim(0, 1.0)
-         plt.plot(likelihood_df[bodypart], color=colors[i], label=bodypart)
-         plt.xlabel('Frames')
-         plt.ylabel('尤度')
-         plt.title('フレーム別の尤度')
-         # 凡例を右上の外側に表示
-         plt.legend(bbox_to_anchor=(1.05, 1),
-                    loc='upper left', borderaxespad=0)
-         # 凡例がはみ出さないようにレイアウトを調整
-         plt.tight_layout()
-         plt.savefig(f'{self.output_folder}/{bodypart}_likelihood.png')
-         image_paths.append(f'{self.output_folder}/{bodypart}_likelihood.png')
-         plt.close()
-
-      # 全ての付属肢の尤度をプロット
-      plt.figure(figsize=(8, 6))
-      plt.ylim(0, 1.0)
-      for i, column in enumerate(likelihood_df.columns):
-         plt.plot(likelihood_df[column], color=colors[i], label=column[0])
-      plt.xlabel('Frames')
-      plt.ylabel('尤度')
-      plt.title('フレーム別の尤度')
-      # 凡例を右上の外側に表示
-      plt.legend(bbox_to_anchor=(1.05, 1), loc='upper left', borderaxespad=0)
-      # 凡例がはみ出さないようにレイアウトを調整
-      plt.tight_layout()
-      plt.savefig(f'{self.output_folder}/likelihood_plot.png')
-      plt.close()
-      image_paths.append(f'{self.output_folder}/likelihood_plot.png')
-      return image_paths
-
-# 以下のGradioInterfaceクラスとメイン実行部分は変更なし
-
-
-class GradioInterface:
-   def __init__(self):
-      self.interface = gr.Interface(
-          fn=self.process_and_plot,
-          inputs=[
-              gr.File(label="CSVファイルをドラッグ＆ドロップ"),
-              gr.Textbox(label="付属肢の名前（カンマ区切り）",
-                         value="指節1, 指節2, 指節3, 指節4, 指節5, 指節6, 指節7, 指節8, 指節9,指節10, 指節11, 指節12, 指節13, 指節14, 触角(左), 触角(右), 頭部, 腹尾節"),
-              gr.Number(label="X軸の最大値", value=1920),
-              gr.Number(label="Y軸の最大値", value=1080),
-              gr.CheckboxGroup(
-                  label="プロットするグラフを選択",
-                  choices=["散布図", "軌跡図", "尤度グラフ"],
-                  value=["散布図", "軌跡図", "尤度グラフ"],
-                  type="value"
-              )
-          ],
-          outputs=[
-              gr.Gallery(label="散布図"),
-              gr.File(label="ZIPダウンロード")
-          ],
-          title="DeepLabCutグラフ出力ツール",
-          description="CSVファイルからグラフを作成します。"
-      )
-
-   def process_and_plot(self, file, bodypart_names, x_max, y_max, graph_choices):
-      processor = DataProcessor(bodypart_names, x_max, y_max)
-      df, df_likelihood = processor.process_csv(file.name)
-
-      all_image_paths = []
-      if "散布図" in graph_choices:
-         all_image_paths += processor.plot_scatter(df)
-      if "軌跡図" in graph_choices:
-         all_image_paths += processor.plot_trajectories(df)
-      if "尤度グラフ" in graph_choices:
-         all_image_paths += processor.plot_likelihood(df_likelihood)
-
-      shutil.make_archive(processor.output_folder,
-                          'zip', processor.output_folder)
-      return all_image_paths, processor.output_folder + '.zip'
-
-   def launch(self):
-      self.interface.launch()
-
-
-if __name__ == "__main__":
-   gradio_app = GradioInterface()
-   gradio_app.launch()
+   import gradio as gr
+   import polars as pl
+   import os
+   import shutil
+   import numpy as np
+   from PyQt6.QtWidgets import QApplication
+   from PyQt6 import QtCore
+   import pyqtgraph as pg
+   from pyqtgraph.exporters import ImageExporter
+
+   class DataProcessor:
+      def __init__(self, bodypart_names, x_max, y_max):
+         # 余分な空白を除去してリスト化
+         self.bodypart_names = [name.strip() for name in bodypart_names.split(',')]
+         self.x_max = x_max
+         self.y_max = y_max
+         self.output_folder = 'output_plots'
+         if not os.path.exists(self.output_folder):
+               os.makedirs(self.output_folder)
+
+      def process_csv(self, file_path):
+         # CSVをpolarsで読み込み（ヘッダーはなし）
+         df_raw = pl.read_csv(file_path, has_header=False)
+         # 2行分のヘッダー（インデックス1,2）を取得し、最初の列は除外する
+         header1 = df_raw.row(1)[1:]
+         header2 = df_raw.row(2)[1:]
+         new_columns = [f"{h1}|{h2}" for h1, h2 in zip(header1, header2)]
+         # データ部分はインデックス3以降（0-indexed）とし、先頭列を削除
+         df_data = df_raw.slice(3)
+         first_col = df_data.columns[0]
+         df_data = df_data.drop(first_col)
+         df_data.columns = new_columns
+
+         # likelihood列のみ抽出
+         df_likelihood = self.extract_likelihood(df_data)
+         # likelihood列を除去したデータ
+         df_no_likelihood = self.remove_likelihood(df_data)
+         # 付属肢名の置換（左側の名前を mapping で変更）
+         df_renamed = self.rename_bodyparts(df_no_likelihood)
+         return df_renamed, df_likelihood
+
+      def remove_likelihood(self, df):
+         # 列名が "bodypart|likelihood" となっている列を除外
+         new_cols = [col for col in df.columns if col.split("|")[1] != "likelihood"]
+         return df.select(new_cols)
+
+      def rename_bodyparts(self, df):
+         cols = df.columns
+         current_names = []
+         for col in cols:
+               bp = col.split("|")[0]
+               if bp not in current_names:
+                  current_names.append(bp)
+         if len(self.bodypart_names) != len(current_names):
+               raise ValueError("The length of bodypart_names must be equal to the number of bodyparts.")
+         mapping = dict(zip(current_names, self.bodypart_names))
+         new_cols = {col: f"{mapping[col.split('|')[0]]}|{col.split('|')[1]}" for col in cols}
+         return df.rename(new_cols)
+
+      def extract_likelihood(self, df):
+         # likelihood列のみを抽出
+         likelihood_cols = [col for col in df.columns if col.split("|")[1] == "likelihood"]
+         df_likelihood = df.select(likelihood_cols)
+         current_names = []
+         for col in likelihood_cols:
+               bp = col.split("|")[0]
+               if bp not in current_names:
+                  current_names.append(bp)
+         if len(self.bodypart_names) != len(current_names):
+               raise ValueError("The length of bodypart_names must be equal to the number of bodyparts.")
+         mapping = dict(zip(current_names, self.bodypart_names))
+         new_cols = {col: f"{mapping[col.split('|')[0]]}|{col.split('|')[1]}" for col in likelihood_cols}
+         return df_likelihood.rename(new_cols)
+
+      def get_bodyparts(self, df):
+         bodyparts = []
+         for col in df.columns:
+               bp = col.split("|")[0]
+               if bp not in bodyparts:
+                  bodyparts.append(bp)
+         return bodyparts
+
+      def plot_scatter(self, df):
+         image_paths = []
+         bodyparts = self.get_bodyparts(df)
+         app = QApplication.instance()
+         if app is None:
+               app = QApplication([])
+
+         # 個別の散布図を作成
+         for i, bodypart in enumerate(bodyparts):
+               try:
+                  x = np.array(df[f"{bodypart}|x"].to_list(), dtype=float)
+                  y = np.array(df[f"{bodypart}|y"].to_list(), dtype=float)
+               except Exception as e:
+                  continue
+               pw = pg.PlotWidget(title=f'トラッキングの座標({bodypart})')
+               pw.setLabel('bottom', 'X Coordinate(pixel)')
+               pw.setLabel('left', 'Y Coordinate(pixel)')
+               pw.setXRange(0, self.x_max)
+               pw.setYRange(0, self.y_max)
+               pw.invertY(True)
+               color = pg.intColor(i, len(bodyparts))
+               # 散布図アイテムの追加
+               scatter = pg.ScatterPlotItem(x=x, y=y, pen=pg.mkPen(color=color), symbol='o', brush=color)
+               pw.addItem(scatter)
+               # 始点を黒丸でハイライトし、"Start"テキストを追加
+               if len(x) > 0:
+                  scatter_start = pg.ScatterPlotItem(x=[x[0]], y=[y[0]], pen=pg.mkPen(color='k'), symbol='o', size=10, brush='k')
+                  pw.addItem(scatter_start)
+                  text = pg.TextItem("Start", color='k')
+                  text.setPos(x[0], y[0])
+                  pw.addItem(text)
+               # PNGにエクスポート
+               exporter = ImageExporter(pw.plotItem)
+               filename = os.path.join(self.output_folder, f"{bodypart}.png")
+               exporter.export(filename)
+               image_paths.append(filename)
+
+         # 全付属肢の散布図を作成
+         pw_all = pg.PlotWidget(title='トラッキングの座標(全付属肢)')
+         pw_all.setLabel('bottom', 'X Coordinate(pixel)')
+         pw_all.setLabel('left', 'Y Coordinate(pixel)')
+         pw_all.setXRange(0, self.x_max)
+         pw_all.setYRange(0, self.y_max)
+         pw_all.invertY(True)
+         for i, bodypart in enumerate(bodyparts):
+               try:
+                  x = np.array(df[f"{bodypart}|x"].to_list(), dtype=float)
+                  y = np.array(df[f"{bodypart}|y"].to_list(), dtype=float)
+               except Exception as e:
+                  continue
+               color = pg.intColor(i, len(bodyparts))
+               scatter = pg.ScatterPlotItem(x=x, y=y, pen=pg.mkPen(color=color), symbol='o', brush=color)
+               pw_all.addItem(scatter)
+         exporter_all = ImageExporter(pw_all.plotItem)
+         filename_all = os.path.join(self.output_folder, "all_plot.png")
+         exporter_all.export(filename_all)
+         image_paths.append(filename_all)
+         return image_paths
+
+      def plot_trajectories(self, df):
+         image_paths = []
+         bodyparts = self.get_bodyparts(df)
+         app = QApplication.instance()
+         if app is None:
+               app = QApplication([])
+
+         # 個別の軌跡図を作成
+         for i, bodypart in enumerate(bodyparts):
+               try:
+                  x = np.array(df[f"{bodypart}|x"].to_list(), dtype=float)
+                  y = np.array(df[f"{bodypart}|y"].to_list(), dtype=float)
+               except Exception as e:
+                  continue
+               pw = pg.PlotWidget(title=f'トラッキングの座標({bodypart})')
+               pw.setLabel('bottom', 'Frames')
+               pw.setLabel('left', 'Coordinate(pixel)')
+               pen_x = pg.mkPen(color=pg.intColor(i, len(bodyparts)), style=QtCore.Qt.PenStyle.DashLine)
+               pen_y = pg.mkPen(color=pg.intColor(i, len(bodyparts)))
+               pw.plot(x, pen=pen_x, name=f"{bodypart}(x座標)")
+               pw.plot(y, pen=pen_y, name=f"{bodypart}(y座標)")
+               exporter = ImageExporter(pw.plotItem)
+               filename = os.path.join(self.output_folder, f"{bodypart}_trajectories.png")
+               exporter.export(filename)
+               image_paths.append(filename)
+
+         # 全付属肢の軌跡図を作成
+         pw_all = pg.PlotWidget(title='トラッキングの座標(全付属肢)')
+         pw_all.setLabel('bottom', 'Frames')
+         pw_all.setLabel('left', 'Coordinate(pixel)')
+         for i, bodypart in enumerate(bodyparts):
+               try:
+                  x = np.array(df[f"{bodypart}|x"].to_list(), dtype=float)
+                  y = np.array(df[f"{bodypart}|y"].to_list(), dtype=float)
+               except Exception as e:
+                  continue
+               pen_x = pg.mkPen(color=pg.intColor(i, len(bodyparts)), style=QtCore.Qt.PenStyle.DashLine)
+               pen_y = pg.mkPen(color=pg.intColor(i, len(bodyparts)))
+               pw_all.plot(x, pen=pen_x, name=f"{bodypart}(x座標)")
+               pw_all.plot(y, pen=pen_y, name=f"{bodypart}(y座標)")
+         exporter_all = ImageExporter(pw_all.plotItem)
+         filename_all = os.path.join(self.output_folder, "all_trajectories.png")
+         exporter_all.export(filename_all)
+         image_paths.append(filename_all)
+         return image_paths
+
+      def plot_likelihood(self, df_likelihood):
+         image_paths = []
+         bodyparts = self.get_bodyparts(df_likelihood)
+         app = QApplication.instance()
+         if app is None:
+               app = QApplication([])
+
+         # 付属肢ごとの尤度グラフを作成
+         for i, bodypart in enumerate(bodyparts):
+               try:
+                  likelihood = np.array(df_likelihood[f"{bodypart}|likelihood"].to_list(), dtype=float)
+               except Exception as e:
+                  continue
+               pw = pg.PlotWidget(title=f'フレーム別の尤度 ({bodypart})')
+               pw.setLabel('bottom', 'Frames')
+               pw.setLabel('left', '尤度')
+               pw.setYRange(0, 1.0)
+               color = pg.intColor(i, len(bodyparts))
+               pw.plot(likelihood, pen=pg.mkPen(color=color), name=bodypart)
+               exporter = ImageExporter(pw.plotItem)
+               filename = os.path.join(self.output_folder, f"{bodypart}_likelihood.png")
+               exporter.export(filename)
+               image_paths.append(filename)
+
+         # 全付属肢の尤度グラフを作成
+         pw_all = pg.PlotWidget(title='フレーム別の尤度 (全付属肢)')
+         pw_all.setLabel('bottom', 'Frames')
+         pw_all.setLabel('left', '尤度')
+         pw_all.setYRange(0, 1.0)
+         for i, bodypart in enumerate(bodyparts):
+               try:
+                  likelihood = np.array(df_likelihood[f"{bodypart}|likelihood"].to_list(), dtype=float)
+               except Exception as e:
+                  continue
+               color = pg.intColor(i, len(bodyparts))
+               pw_all.plot(likelihood, pen=pg.mkPen(color=color), name=bodypart)
+         exporter_all = ImageExporter(pw_all.plotItem)
+         filename_all = os.path.join(self.output_folder, "likelihood_plot.png")
+         exporter_all.export(filename_all)
+         image_paths.append(filename_all)
+         return image_paths
+
+   class GradioInterface:
+      def __init__(self):
+         self.interface = gr.Interface(
+               fn=self.process_and_plot,
+               inputs=[
+                  gr.File(label="CSVファイルをドラッグ＆ドロップ"),
+                  gr.Textbox(
+                     label="付属肢の名前（カンマ区切り）",
+                     value="指節1, 指節2, 指節3, 指節4, 指節5, 指節6, 指節7, 指節8, 指節9, 指節10, 指節11, 指節12, 指節13, 指節14, 触角(左), 触角(右), 頭部, 腹尾節"
+                  ),
+                  gr.Number(label="X軸の最大値", value=1920),
+                  gr.Number(label="Y軸の最大値", value=1080),
+                  gr.CheckboxGroup(
+                     label="プロットするグラフを選択",
+                     choices=["散布図", "軌跡図", "尤度グラフ"],
+                     value=["散布図", "軌跡図", "尤度グラフ"],
+                     type="value"
+                  )
+               ],
+               outputs=[
+                  gr.Gallery(label="散布図"),
+                  gr.File(label="ZIPダウンロード")
+               ],
+               title="DeepLabCutグラフ出力ツール",
+               description="CSVファイルからグラフを作成します。"
+         )
+
+      def process_and_plot(self, file, bodypart_names, x_max, y_max, graph_choices):
+         processor = DataProcessor(bodypart_names, x_max, y_max)
+         df, df_likelihood = processor.process_csv(file.name)
+
+         all_image_paths = []
+         if "散布図" in graph_choices:
+               all_image_paths += processor.plot_scatter(df)
+         if "軌跡図" in graph_choices:
+               all_image_paths += processor.plot_trajectories(df)
+         if "尤度グラフ" in graph_choices:
+               all_image_paths += processor.plot_likelihood(df_likelihood)
+
+         shutil.make_archive(processor.output_folder, 'zip', processor.output_folder)
+         return all_image_paths, processor.output_folder + '.zip'
+
+      def launch(self):
+         self.interface.launch()
+
+
+   if __name__ == "__main__":
+      gradio_app = GradioInterface()
+      gradio_app.launch()