Update app.py

app.py

@@ -8,45 +8,34 @@ import japanize_matplotlib
 
 
 class DataProcessor:
-   def __init__(self, bodypart_names, x_max, y_max):
-      self.bodypart_names = bodypart_names.split(',')
+   def __init__(self, x_max, y_max):
       self.x_max = x_max
       self.y_max = y_max
       self.output_folder = 'output_plots'
+      self.bodypart_names = None  # 初期化時にはNoneに設定
 
    def process_csv(self, file_path):
       df = pd.read_csv(file_path, header=[1, 2])
+      
+      # CSVから自動的に付属肢名を抽出
+      self.bodypart_names = df.columns.get_level_values(0).unique().tolist()
+      # 最初の列（通常はscorerなど）を除外
+      if len(self.bodypart_names) > 0:
+         self.bodypart_names = self.bodypart_names[1:]
+      
       df_likelihood = self.extract_likelihood(df)
       df = self.remove_first_column_and_likelihood(df)
-      df = self.rename_bodyparts(df)
-      return df, df_likelihood
+      return df, df_likelihood, self.bodypart_names  # 抽出した付属肢名も返す
 
    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):
@@ -172,16 +161,12 @@ class DataProcessor:
       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(
@@ -192,16 +177,17 @@ class GradioInterface:
               )
           ],
           outputs=[
-              gr.Gallery(label="散布図"),
-              gr.File(label="ZIPダウンロード")
+              gr.Gallery(label="グラフ"),
+              gr.File(label="ZIPダウンロード"),
+              gr.Textbox(label="検出された付属肢")  # 検出された付属肢を表示するための出力を追加
           ],
           title="DeepLabCutグラフ出力ツール",
-          description="CSVファイルからグラフを作成します。"
+          description="CSVファイルからグラフを作成します。付属肢は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)
+   def process_and_plot(self, file, x_max, y_max, graph_choices):
+      processor = DataProcessor(x_max, y_max)
+      df, df_likelihood, bodypart_names = processor.process_csv(file.name)
 
       all_image_paths = []
       if "散布図" in graph_choices:
@@ -211,9 +197,11 @@ class GradioInterface:
       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'
+      # 付属肢の名前を表示用に結合
+      bodyparts_text = ", ".join(bodypart_names)
+      
+      shutil.make_archive(processor.output_folder, 'zip', processor.output_folder)
+      return all_image_paths, processor.output_folder + '.zip', bodyparts_text
 
    def launch(self):
       self.interface.launch()