Update app.py

app.py

@@ -97,18 +97,27 @@ if upload_file is not None:
             else:
                 new_columns = [f'column_{i+1}' for i in range(len(df.columns))]
                 df.columns = new_columns
-                
                 st.dataframe(df, key="csv_dataframe")
-                fig = px.treemap(df, path=[px.Constant("all"), 'column_1', 'column_3'],
-                                  color='column_2')
-                fig.update_layout(margin = dict(t=50, l=25, r=25, b=25))
-                st.subheader("Tree map", divider = "red")
-                st.plotly_chart(fig)
+                all_columns = df.columns.tolist()
+                st.subheader("Select columns for the Tree Map", divider="blue")
+                parent_column = st.selectbox("Select the parent column:", all_columns)
+                value_column = st.selectbox("Select the value column:", all_columns)
+                color_column = st.selectbox("Select the color column (optional):", ["None"] + all_columns)
+                if parent_column and value_column:
+                    path_columns = [px.Constant("all"), parent_column]
+                    color_param = value_column if color_column == "None" else color_column
+                    fig = px.treemap(df,
+                             path=path_columns,
+                             values=value_column,
+                             color=color_param)
+                    fig.update_layout(margin=dict(t=50, l=25, r=25, b=25))
+                    st.subheader("Tree map", divider="red")
+                    st.plotly_chart(fig)
                 
                 
-                st.write("_number of rows_", df.shape[0])
-                st.write("_number of columns_", df.shape[1])
-                st.session_state.df = df  
+                    st.write("_number of rows_", df.shape[0])
+                    st.write("_number of columns_", df.shape[1])
+                    st.session_state.df = df  
         except pd.errors.ParserError:
             st.error("Error: The CSV file is not readable or is incorrectly formatted.")
             st.stop()