Update app.py

app.py

@@ -5,6 +5,7 @@ from Bio import SeqIO
 from Bio.SeqUtils.ProtParam import ProteinAnalysis
 from Bio.Graphics import GenomeDiagram
 from reportlab.lib.colors import Color, lightblue, blue
+from reportlab.lib import colors
 from reportlab.lib.units import cm
 from io import StringIO
 from collections import Counter
@@ -71,19 +72,23 @@ def add_protein_features(cds_info):
     return cds_info
 
 # Function to generate genome diagram
-def create_genome_diagram(genbank_content, output_file_path):
+def create_genome_diagram(genbank_content, output_file_path, selected_features, color_1, color_2):
     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 not in ["gene", "CDS"]:
-            continue
-        color = lightblue if len(gd_feature_set) % 2 == 0 else blue
-        gd_feature_set.add_feature(feature, color=color, label=True, label_size=10, label_angle=0)
+        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)
 
-    # Change the output format to SVG
     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")
 
@@ -104,6 +109,13 @@ if uploaded_file is not None:
         window_size = st.number_input('GC content sliding window size', min_value=100, max_value=10000, value=1000)
         k = st.number_input('k-mer size', min_value=1, max_value=10, value=6)
 
+    # 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'])
+
+
     # Main content
     col1, col2 = st.columns(2)
     with col1:
@@ -130,11 +142,11 @@ 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))
 
-    # Generate and display genome diagram
+    # 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)
+    create_genome_diagram(uploaded_file.getvalue().decode("utf-8"), output_file_path_svg, feature_selection, diagram_color_1, diagram_color_2)
     st.image(output_file_path_svg, caption='Genome Diagram')
-
+    
     # Additional Information
     with st.expander("View All Genes"):
         st.dataframe(gene_df)