更新了选择target的功能

app.py

@@ -222,8 +222,10 @@ with gr.Blocks() as demo:
     # componets = []
     # for _ in range(TIME_PLOTS_NUM):
     with gr.Row():
-        with gr.Column(scale=3):
-            select_box = gr.Dropdown(choices=["items"], label="Select some items", multiselect=True, interactive=True)
+        with gr.Column(scale=2):
+            select_sample_box = gr.Dropdown(choices=["items"], label="Select some items", multiselect=True, interactive=True)
+        with gr.Column(scale=2):
+            select_subtarget_box = gr.Dropdown(choices=["subtargets"], label="Select some subtargets", multiselect=True, interactive=True)
         with gr.Column(scale=1):
             select_buttom = gr.Button("Show selected items")
     with gr.Row():
@@ -232,7 +234,7 @@ with gr.Blocks() as demo:
         with gr.Column(scale=3):
             plot = gr.Plot()
     user_input_text = gr.Textbox(placeholder="输入一些内容")
-    # componets.append({"select_box": select_box,
+    # componets.append({"select_sample_box": select_sample_box,
     #                   "statistics_textbox": statistics_textbox,
     #                   "user_input_text": user_input_text,
     #                   "plot": plot})
@@ -248,7 +250,7 @@ with gr.Blocks() as demo:
             cp_result: gr.update(visible=False, value=""),
         }
 
-    def show_dataset_at_config_and_split_and_page(dataset: str, config: str, split: str, page: str|List[str]) -> dict:
+    def show_dataset_at_config_and_split_and_page(dataset: str, config: str, split: str, page: str|List[str], sub_targets: List[int|str]) -> dict:
         try:
             ret = {}
             if dataset != 'Salesforce/lotsa_data':
@@ -261,15 +263,17 @@ with gr.Blocks() as demo:
                 df_list, id_list = [], []
                 for i, page in enumerate(page):
                     df, max_page, info = get_page(dataset, config, split, page)
-                    df = clean_up_df(df)
+                    global tot_samples, tot_targets
+                    tot_samples, tot_targets = max_page, len(df['target'][0]) if isinstance(df['target'][0], np.ndarray) else 1
+                    
+                    df = clean_up_df(df, sub_targets)
                     row = df.iloc[0]
                     id_list.append(row['item_id'])
                     # 将单行的DataFrame展开为新的DataFrame
                     df_without_index = row.drop('item_id').to_frame().T
                     df_expanded = df_without_index.apply(pd.Series.explode).reset_index(drop=True).fillna(0)
                     df_list.append(df_expanded)
-                    global tot_samples
-                    tot_samples = max_page
+
             return {
                 statistics_textbox: gr.update(value=create_statistic(df_list, id_list)),
                 plot: gr.update(value=create_plot(df_list, id_list)),
@@ -292,8 +296,9 @@ with gr.Blocks() as demo:
     def show_dataset_at_config_and_split(dataset: str, config: str, split: str) -> dict:
         try:
             return {
-                **show_dataset_at_config_and_split_and_page(dataset, config, split, "1"),
-                select_box: gr.update(choices=[f"{i+1}" for i in range(tot_samples)], value=["1"]),
+                **show_dataset_at_config_and_split_and_page(dataset, config, split, "1", [0]),
+                select_sample_box: gr.update(choices=[f"{i+1}" for i in range(tot_samples)], value=["1"]),
+                select_subtarget_box: gr.update(choices=[i for i in range(tot_targets)]+['all'], value=[0]),
                 # cp_page: gr.update(value="1", visible=True),
                 # cp_goto_page: gr.update(visible=True),
                 # cp_goto_next_page: gr.update(visible=True),
@@ -336,17 +341,18 @@ with gr.Blocks() as demo:
     all_outputs = [cp_config, cp_split,
                 #    cp_page, cp_goto_page, cp_goto_next_page,
                    cp_result, cp_info, cp_error,
-                   select_box, select_buttom, statistics_textbox, user_input_text, plot]
+                   select_sample_box, select_subtarget_box,
+                   select_buttom, statistics_textbox, user_input_text, plot]
     cp_go.click(show_dataset, inputs=[cp_dataset], outputs=all_outputs)
     cp_config.change(show_dataset_at_config, inputs=[cp_dataset, cp_config], outputs=all_outputs)
     cp_split.change(show_dataset_at_config_and_split, inputs=[cp_dataset, cp_config, cp_split], outputs=all_outputs)
     # cp_goto_page.click(show_dataset_at_config_and_split_and_page, inputs=[cp_dataset, cp_config, cp_split, cp_page], outputs=all_outputs)
     # cp_goto_next_page.click(show_dataset_at_config_and_split_and_next_page, inputs=[cp_dataset, cp_config, cp_split, cp_page], outputs=all_outputs)
     user_input_text.submit(save_to_file, inputs=user_input_text)
-    select_buttom.click(show_dataset_at_config_and_split_and_page, inputs=[cp_dataset, cp_config, cp_split, select_box], outputs=all_outputs)
+    select_buttom.click(show_dataset_at_config_and_split_and_page, inputs=[cp_dataset, cp_config, cp_split, select_sample_box, select_subtarget_box], outputs=all_outputs)
 
 
 if __name__ == "__main__":
 
     app = gr.mount_gradio_app(app, demo, path="/")
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="127.0.0.1", port=7860)

utils.py

@@ -33,22 +33,22 @@ def ndarray_to_base64(ndarray):
     base64_str = base64.b64encode(buffer.getvalue()).decode('utf-8')
     return f"data:image/png;base64,{base64_str}"
 
-def flatten_ndarray_column(df, column_name):
+def flatten_ndarray_column(df, column_name, rows_to_include):
     """
-    将嵌套的np.ndarray列展平为多列。
+    将嵌套的np.ndarray列展平为多列，并只保留指定的行。
     """
-    def flatten_ndarray(ndarray):
+    def select_and_flatten(ndarray):
         if isinstance(ndarray, np.ndarray) and ndarray.dtype == 'O':
-            return np.concatenate([flatten_ndarray(subarray) for subarray in ndarray])
+            selected = [ndarray[i] for i in rows_to_include if i < len(ndarray)]
+            return np.concatenate([select_and_flatten(subarray) for subarray in selected])
         elif isinstance(ndarray, np.ndarray) and ndarray.ndim == 1:
             return np.expand_dims(ndarray, axis=0)
         return ndarray
 
-    flattened_data = df[column_name].apply(flatten_ndarray)
-    max_length = max(flattened_data.apply(len))
+    selected_data = df[column_name].apply(select_and_flatten)
 
-    for i in range(max_length):
-        df[f'{column_name}_{i}'] = flattened_data.apply(lambda x: x[i] if i < len(x) else np.nan)
+    for i in rows_to_include:
+        df[f'{column_name}_{i}'] = selected_data.apply(lambda x: x[i] if i < len(x) else np.nan)
 
     return df
 
@@ -110,16 +110,21 @@ def create_statistic(dfs: list[pd.DataFrame], ids: list[str]):
     combined_stats_df = pd.concat(stats_list, ignore_index=True)
     return combined_stats_df
 
-def clean_up_df(df: pd.DataFrame) -> pd.DataFrame:
+def clean_up_df(df: pd.DataFrame, rows_to_include: list[int]) -> pd.DataFrame:
     """
     清理数据集，将嵌套的np.ndarray列展平为多列。
     """
+    if 'all' in rows_to_include:
+        rows_to_include = list(range(len(df['target'][0]))) if isinstance(df['target'][0], np.ndarray) else 1
+    else:
+        rows_to_include = sorted(rows_to_include)
+
     df['timestamp'] = df.apply(lambda row: pd.date_range(
         start=row['start'], 
         periods=len(row['target'][0]) if isinstance(row['target'][0], np.ndarray) else len(row['target']), 
         freq=row['freq']
     ).to_pydatetime().tolist(), axis=1)
-    df = flatten_ndarray_column(df, 'target')
+    df = flatten_ndarray_column(df, 'target', rows_to_include)
     # 删除原始的start和freq列
     df.drop(columns=['start', 'freq', 'target'], inplace=True)
     if 'past_feat_dynamic_real' in df.columns: