Update app.py

app.py

@@ -72,26 +72,21 @@ def add_protein_features(cds_info):
     return cds_info
 
 # Function to generate genome diagram
-def create_genome_diagram(genbank_content, output_file_path, selected_features, color_1, color_2):
+def create_genome_diagram(genbank_content, output_file_path, display_features, colors):
     record = SeqIO.read(StringIO(genbank_content), "genbank")
     gd_diagram = GenomeDiagram.Diagram(record.id)
     gd_track_for_features = gd_diagram.new_track(1, name="Annotated Features")
     gd_feature_set = gd_track_for_features.new_set()
 
-    color_1 = colors.HexColor(color_1)
-    color_2 = colors.HexColor(color_2)
-
     for feature in record.features:
-        if feature.type in selected_features:
-            if len(gd_feature_set) % 2 == 0:
-                color = color_1
-            else:
-                color = color_2
-            gd_feature_set.add_feature(feature, color=color, label=True, label_size=10, label_angle=0)
+        if feature.type in display_features:
+            feature_color = colors.get(feature.type, colors['default'])  # Get color, default if not found
+            gd_feature_set.add_feature(feature, color=feature_color, label=True, label_size=10, label_angle=0)
 
     gd_diagram.draw(format="circular", circular=True, pagesize=(20*cm, 20*cm), start=0, end=len(record), circle_core=0.7)
     gd_diagram.write(output_file_path, "SVG")
 
+
 # Streamlit UI setup
 st.set_page_config(page_title="Genomic Data Dashboard", page_icon="🧬", layout="wide")
 uploaded_file = st.file_uploader("Upload a GenBank file", type=['gb', 'gbk'])
@@ -111,9 +106,13 @@ if uploaded_file is not None:
 
     # Sidebar options for diagram customization
     with st.sidebar:
-        diagram_color_1 = st.color_picker("Pick a color for features (odd positions)", '#ff9999')
-        diagram_color_2 = st.color_picker("Pick a color for features (even positions)", '#66b3ff')
-        feature_selection = st.multiselect("Select features to display:", ['gene', 'tRNA', 'CDS', 'rRNA'], default=['gene', 'CDS'])
+        # Sidebar options for color customization
+        color_gene = st.color_picker("Gene color", '#ff9999')
+        color_cds = st.color_picker("CDS color", '#66b3ff')
+        color_trna = st.color_picker("tRNA color", '#99ff99')
+        color_rrna = st.color_picker("rRNA color", '#f1c40f')
+        # Option to select what to display on the diagram
+        display_options = st.multiselect("Select features to display:", ['gene', 'tRNA', 'CDS', 'rRNA'], default=['gene', 'CDS'])
 
 
     # Main content
@@ -142,6 +141,15 @@ if uploaded_file is not None:
         kmers = calculate_kmers(sequence, k)  # Ensure this function is defined as per previous instructions
         st.bar_chart(pd.DataFrame.from_dict(kmers, orient='index', columns=['Frequency']).sort_values('Frequency', ascending=False).head(20))
 
+    # Construct the colors dictionary
+    feature_colors = {
+        'gene': color_gene,
+        'CDS': color_cds,
+        'tRNA': color_trna,
+        'rRNA': color_rrna,
+        'default': '#cccccc'  # Default color if needed
+    }
+    
     # Generate and display genome diagram with user-selected options
     output_file_path_svg = os.path.join(temp_dir, "genome_diagram.svg")
     create_genome_diagram(uploaded_file.getvalue().decode("utf-8"), output_file_path_svg, feature_selection, diagram_color_1, diagram_color_2)