Update app.py

app.py

@@ -78,6 +78,7 @@ def add_protein_features(cds_info):
 
 # Updated Function to generate genome diagram
 def create_genome_diagram(genbank_content, output_file_path, colors_dict, diagram_type="linear", diagram_size=(30, 10)):
+    from Bio.Graphics.GenomeDiagram import Diagram, Track, FeatureSet
     record = SeqIO.read(StringIO(genbank_content), "genbank")
     gd_diagram = Diagram(record.id)
     
@@ -85,18 +86,20 @@ def create_genome_diagram(genbank_content, output_file_path, colors_dict, diagra
     max_tracks = len(colors_dict)
     track_indices = {feature_type: idx+1 for idx, feature_type in enumerate(colors_dict.keys())}
     feature_tracks = {}
+    feature_sets = {}  # Dictionary to store FeatureSets
     for feature_type, idx in track_indices.items():
         feature_tracks[feature_type] = gd_diagram.new_track(
             idx, name=feature_type, scale=False, greytrack=False, height=0.5
         )
-        feature_tracks[feature_type].new_set()
+        # Store the FeatureSet
+        feature_sets[feature_type] = feature_tracks[feature_type].new_set()
     
     for feature in record.features:
         feature_type = feature.type
         if feature_type in colors_dict:
             color = colors.HexColor(colors_dict[feature_type])
-            # Adjust label parameters to minimize overlap
-            feature_set = feature_tracks[feature_type].sets[0]
+            # Retrieve the FeatureSet from the dictionary
+            feature_set = feature_sets[feature_type]
             feature_set.add_feature(
                 feature,
                 color=color,