yinuozhang commited on
Commit
bb04c63
·
1 Parent(s): bb32177

bug fixing

Browse files
Files changed (1) hide show
  1. app.py +18 -16
app.py CHANGED
@@ -10,7 +10,7 @@ from PIL import Image, ImageDraw, ImageFont
10
  import matplotlib.pyplot as plt
11
  import matplotlib.patches as patches
12
  from io import BytesIO
13
-
14
  import re
15
  from rdkit import Chem
16
 
@@ -1038,7 +1038,8 @@ def process_input(smiles_input=None, file_obj=None, show_linear=False,
1038
  show_segment_details=False, generate_3d=False, use_uff=False):
1039
  """Process input and create visualizations using PeptideAnalyzer"""
1040
  analyzer = PeptideAnalyzer()
1041
- structure_files = None
 
1042
 
1043
  # Handle direct SMILES input
1044
  if smiles_input:
@@ -1059,24 +1060,25 @@ def process_input(smiles_input=None, file_obj=None, show_linear=False,
1059
  generator = PeptideStructureGenerator()
1060
 
1061
  try:
1062
- # Create a list to store file tuples (name, data)
1063
- files_list = []
1064
-
1065
  # Generate ETKDG structure
1066
  mol_etkdg = generator.generate_structure_etkdg(smiles)
1067
- etkdg_bytes = generator.mol_to_sdf_bytes(mol_etkdg)
1068
- files_list.append(("structure_etkdg.sdf", etkdg_bytes))
 
 
 
1069
 
1070
  # Generate UFF structure if requested
1071
  if use_uff:
1072
  mol_uff = generator.generate_structure_uff(smiles)
1073
- uff_bytes = generator.mol_to_sdf_bytes(mol_uff)
1074
- files_list.append(("structure_uff.sdf", uff_bytes))
1075
-
1076
- structure_files = files_list if files_list else None
 
1077
 
1078
  except Exception as e:
1079
- return f"Error generating 3D structures: {str(e)}", None, None, []
1080
 
1081
  # Use analyzer to get sequence
1082
  segments = analyzer.split_on_bonds(smiles)
@@ -1148,13 +1150,13 @@ def process_input(smiles_input=None, file_obj=None, show_linear=False,
1148
 
1149
  if structure_files:
1150
  summary += "\n3D Structures Generated:\n"
1151
- for filename, _ in structure_files:
1152
- summary += f"- {filename}\n"
1153
 
1154
- return summary + output_text, img_cyclic, img_linear, structure_files
1155
 
1156
  except Exception as e:
1157
- return f"Error processing SMILES: {str(e)}", None, None
1158
 
1159
  # Handle file input
1160
  if file_obj is not None:
 
10
  import matplotlib.pyplot as plt
11
  import matplotlib.patches as patches
12
  from io import BytesIO
13
+ import tempfile
14
  import re
15
  from rdkit import Chem
16
 
 
1038
  show_segment_details=False, generate_3d=False, use_uff=False):
1039
  """Process input and create visualizations using PeptideAnalyzer"""
1040
  analyzer = PeptideAnalyzer()
1041
+ temp_dir = tempfile.mkdtemp() if generate_3d else None
1042
+ structure_files = []
1043
 
1044
  # Handle direct SMILES input
1045
  if smiles_input:
 
1060
  generator = PeptideStructureGenerator()
1061
 
1062
  try:
 
 
 
1063
  # Generate ETKDG structure
1064
  mol_etkdg = generator.generate_structure_etkdg(smiles)
1065
+ etkdg_path = os.path.join(temp_dir, "structure_etkdg.sdf")
1066
+ writer = Chem.SDWriter(etkdg_path)
1067
+ writer.write(mol_etkdg)
1068
+ writer.close()
1069
+ structure_files.append(etkdg_path)
1070
 
1071
  # Generate UFF structure if requested
1072
  if use_uff:
1073
  mol_uff = generator.generate_structure_uff(smiles)
1074
+ uff_path = os.path.join(temp_dir, "structure_uff.sdf")
1075
+ writer = Chem.SDWriter(uff_path)
1076
+ writer.write(mol_uff)
1077
+ writer.close()
1078
+ structure_files.append(uff_path)
1079
 
1080
  except Exception as e:
1081
+ return f"Error generating 3D structures: {str(e)}", None, None, None
1082
 
1083
  # Use analyzer to get sequence
1084
  segments = analyzer.split_on_bonds(smiles)
 
1150
 
1151
  if structure_files:
1152
  summary += "\n3D Structures Generated:\n"
1153
+ for filepath in structure_files:
1154
+ summary += f"- {os.path.basename(filepath)}\n"
1155
 
1156
+ return summary + output_text, img_cyclic, img_linear, structure_files if structure_files else None
1157
 
1158
  except Exception as e:
1159
+ return f"Error processing SMILES: {str(e)}", None, None, None
1160
 
1161
  # Handle file input
1162
  if file_obj is not None: