Refactor

main.py

@@ -193,7 +193,10 @@ def predict(pdb_code, pdb_file):
     view.setBackgroundColor('white')
     view.addModel(pdb, "pdb")
     view.setStyle({'stick': {'colorscheme': {'prop': 'resi', 'C': 'turquoise'}}})
-   
+    view.addLight([0, 0, 10], [1, 1, 1], 1)  # Add directional light from the z-axis
+    view.setSpecular(0.5)  # Adjust the specular lighting effect
+    view.setAmbient(0.5)  # Adjust the ambient lighting effect
+    
     for i in range(10):
         adaptability_value = adaptability[topN_ind[i]]
         color = 'orange'
@@ -207,12 +210,6 @@ def predict(pdb_code, pdb_file):
             'color': color,
             'alpha': 0.75
         })
-        # Add lighting and shading options to the view object:
-        view.setStyle({'stick': {'colorscheme': {'prop': 'resi', 'C': 'turquoise'}}})
-        view.setStyle({'sphere': {}})
-        view.addLight([0, 0, 10], [1, 1, 1], 1)  # Add directional light from the z-axis
-        view.setSpecular(0.5)  # Adjust the specular lighting effect
-        view.setAmbient(0.5)  # Adjust the ambient lighting effect
         
     view.zoomTo()