From 00896c90b5acf133eec3f54564f025bda0a9344f Mon Sep 17 00:00:00 2001
From: kiwijuice56 <alfaroeric127@gmail.com>
Date: Wed, 31 Jan 2024 12:51:36 -0500
Subject: [PATCH] Added screenshot feature

---
 gui/user_interface.py     |  2 ++
 protein.py                |  1 +
 renderers/pdb_renderer.py | 21 +++++++++++++++------
 3 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/gui/user_interface.py b/gui/user_interface.py
index d86024b..1b534d7 100644
--- a/gui/user_interface.py
+++ b/gui/user_interface.py
@@ -62,6 +62,8 @@ def on_key_press(self, symbol, _modifiers):
             self.pdb_renderer.set_point_size(self.pdb_renderer.point_size - 1)
         if symbol == pyglet.window.key.O:
             self.pdb_renderer.outline = not self.pdb_renderer.outline
+        if symbol == pyglet.window.key.S:
+            self.pdb_renderer.transparent_save_image = True
 
     def on_mouse_scroll(self, x, y, _scroll_x, _scroll_y):
         self.on_mouse_motion(x, y, 0, 0)
diff --git a/protein.py b/protein.py
index 49fbc55..bc6b3f0 100644
--- a/protein.py
+++ b/protein.py
@@ -257,6 +257,7 @@ def get_color_from_palette(palette):
 
         luminance = residue.go_map[self.current_go_id]
         luminance *= luminance
+        luminance = min(1.0, luminance + 0.1)
 
         residue.outline_color = [int(luminance * self.MAX_OUTLINE_BRIGHTNESS)] * 3
 
diff --git a/renderers/pdb_renderer.py b/renderers/pdb_renderer.py
index b268ad5..60e46ca 100644
--- a/renderers/pdb_renderer.py
+++ b/renderers/pdb_renderer.py
@@ -82,10 +82,10 @@ class PDBRenderer:
     GRID_LINE_COUNT = 16
     BACKGROUND_COLOR = (0, 0, 0, 0)
 
-    SKY_COLOR1 = (114, 136, 142)
-    SKY_COLOR2 = (149, 169, 176)
-    GROUND_COLOR1 = (40, 50, 55)
-    GROUND_COLOR2 = (51, 64, 69)
+    SKY_COLOR1 = (64, 64, 64)
+    SKY_COLOR2 = (92, 92, 92)
+    GROUND_COLOR1 = (48, 48, 48)
+    GROUND_COLOR2 = (59, 59, 59)
     SKY_DISTANCE = 1000.0
     HORIZON_SIZE1 = 0.3
     HORIZON_SIZE2 = 0.1
@@ -108,6 +108,8 @@ def __init__(self, protein, window, bounding_box=None, point_size=8, outline=Tru
         """
         self.window = window
         self.protein = protein
+        self.save_image = False
+        self.transparent_save_image = False
 
         self.point_size = 0
         self.set_point_size(point_size)
@@ -241,14 +243,16 @@ def draw(self):
              self.SKY_COLOR2 * 2 + self.SKY_COLOR1 * 2 + self.SKY_COLOR1 * 4
              )
         )
-        quad_list.draw(GL_QUADS)
+        if not self.transparent_save_image:
+            quad_list.draw(GL_QUADS)
 
         glLoadIdentity()
         glMatrixMode(GL_PROJECTION)
         gluPerspective(self.FOV, self.window.width / float(self.window.height), self.Z_NEAR, self.Z_FAR)
         self.camera.draw()
 
-        self.grid_list.draw(GL_LINES)
+        if not self.transparent_save_image:
+            self.grid_list.draw(GL_LINES)
 
         glEnable(GL_POINT_SMOOTH)
         if self.outline:
@@ -260,6 +264,11 @@ def draw(self):
         glEnable(GL_DEPTH_TEST)
         self.atom_vertices.draw(GL_POINTS)
 
+        if self.save_image or self.transparent_save_image:
+            pyglet.image.get_buffer_manager().get_color_buffer().save('screenshot.png')
+            self.save_image = False
+            self.transparent_save_image = False
+
         # Clean up
         glDisable(GL_SCISSOR_TEST)
         glDisable(GL_LINE_SMOOTH)