Fix screenshot generation to handle colors correctly

We want to convert premult sRGB in output framebuffer to sRGB
non premult pixels.

Fixes #9525

kitty/fast_data_types.pyi

@@ -282,6 +282,7 @@ CELL_PROGRAM: int
 CELL_FG_PROGRAM: int
 CELL_BG_PROGRAM: int
 BLIT_PROGRAM: int
+SCREENSHOT_PROGRAM: int
 ROUNDED_RECT_PROGRAM: int
 DECORATION: int
 BLINK: int

kitty/screenshot_fragment.glsl

@@ -0,0 +1,54 @@
+#pragma kitty_include_shader <linear2srgb.glsl>
+
+uniform sampler2D image;
+uniform vec2 src_size;  // Source texture size in pixels
+
+in vec2 texcoord;
+out vec4 output_color;
+
+void main() {
+    // The input texture contains sRGB colors with premultiplied alpha.
+    // We need to output unpremultiplied sRGB colors with proper downscaling.
+    
+    // For proper downscaling, we need to:
+    // 1. Sample neighboring pixels
+    // 2. Convert from sRGB to linear (unpremultiplying first)
+    // 3. Average in linear space
+    // 4. Convert back to sRGB
+    // 5. Output unpremultiplied
+    
+    // Calculate the texel size
+    vec2 texel_size = 1.0 / src_size;
+    
+    // Sample a 2x2 grid for better quality downscaling
+    // This provides basic bilinear-like filtering in linear space
+    vec2 tc = texcoord;
+    
+    vec4 s00 = texture(image, tc + vec2(-0.25, -0.25) * texel_size);
+    vec4 s10 = texture(image, tc + vec2( 0.25, -0.25) * texel_size);
+    vec4 s01 = texture(image, tc + vec2(-0.25,  0.25) * texel_size);
+    vec4 s11 = texture(image, tc + vec2( 0.25,  0.25) * texel_size);
+    
+    // Unpremultiply and convert to linear for each sample
+    vec3 linear00 = s00.a > 0.0 ? srgb2linear(s00.rgb / s00.a) : vec3(0.0);
+    vec3 linear10 = s10.a > 0.0 ? srgb2linear(s10.rgb / s10.a) : vec3(0.0);
+    vec3 linear01 = s01.a > 0.0 ? srgb2linear(s01.rgb / s01.a) : vec3(0.0);
+    vec3 linear11 = s11.a > 0.0 ? srgb2linear(s11.rgb / s11.a) : vec3(0.0);
+    
+    // Average the alpha values
+    float avg_alpha = (s00.a + s10.a + s01.a + s11.a) * 0.25;
+    
+    // For proper downsampling with transparency, weight colors by their alpha
+    // This ensures partially transparent pixels contribute proportionally
+    vec3 weighted_sum = linear00 * s00.a + linear10 * s10.a + linear01 * s01.a + linear11 * s11.a;
+    float total_weight = s00.a + s10.a + s01.a + s11.a;
+    
+    // Calculate the weighted average color in linear space
+    vec3 avg_linear = total_weight > 0.0 ? weighted_sum / total_weight : vec3(0.0);
+    
+    // Convert back to sRGB
+    vec3 srgb_color = linear2srgb(avg_linear);
+    
+    // Output unpremultiplied sRGB color
+    output_color = vec4(srgb_color, avg_alpha);
+}

kitty/screenshot_vertex.glsl

@@ -0,0 +1,2 @@
+uniform vec4 src_rect, dest_rect;
+#pragma kitty_include_shader <blit_common.glsl>

kitty/shaders.c

@@ -25,6 +25,7 @@ enum {
     TINT_PROGRAM,
     TRAIL_PROGRAM,
     BLIT_PROGRAM,
+    SCREENSHOT_PROGRAM,
     ROUNDED_RECT_PROGRAM,
     NUM_PROGRAMS
 };
@@ -363,6 +364,10 @@ typedef struct {
 } BlitProgramLayout;
 static BlitProgramLayout blit_program_layout;
 
+typedef struct {
+    ScreenshotUniforms uniforms;
+} ScreenshotProgramLayout;
+static ScreenshotProgramLayout screenshot_program_layout;
 
 static void
 init_cell_program(void) {
@@ -390,6 +395,7 @@ init_cell_program(void) {
     get_uniform_locations_tint(TINT_PROGRAM, &tint_program_layout.uniforms);
     get_uniform_locations_trail(TRAIL_PROGRAM, &trail_program_layout.uniforms);
     get_uniform_locations_blit(BLIT_PROGRAM, &blit_program_layout.uniforms);
+    get_uniform_locations_screenshot(SCREENSHOT_PROGRAM, &screenshot_program_layout.uniforms);
     get_uniform_locations_rounded_rect(ROUNDED_RECT_PROGRAM, &rounded_rect_program_layout.uniforms);
 }
 
@@ -718,6 +724,7 @@ set_cell_uniforms(bool force) {
             glUniform1f(cu->text_gamma_adjustment, text_gamma_adjustment);
         }
         bind_program(BLIT_PROGRAM); glUniform1i(blit_program_layout.uniforms.image, GRAPHICS_UNIT);
+        bind_program(SCREENSHOT_PROGRAM); glUniform1i(screenshot_program_layout.uniforms.image, GRAPHICS_UNIT);
         constants_set = true;
     }
 }
@@ -1439,7 +1446,8 @@ setup_os_window_for_rendering(OSWindow *os_window, Tab *tab, Window *active_wind
 // dimension of the source region (viewport or central region without tab bar).
 // Takes a screenshot of a rectangular region of the OSWindow's framebuffer.
 // The region parameter specifies which part of the framebuffer to capture.
-// Scaling is performed on the GPU using the BLIT_PROGRAM shader for better performance.
+// Scaling is performed on the GPU using the SCREENSHOT_PROGRAM shader for better performance.
+// The shader properly handles sRGB color space conversion and downscaling.
 // Setting the thumbnail dimensions to zero disables scaling.
 void
 take_screenshot_of_rectangular_region(OSWindow *os_window, Region region, unsigned char *dst_buf, unsigned *thumb_w, unsigned *thumb_h) {
@@ -1479,8 +1487,8 @@ take_screenshot_of_rectangular_region(OSWindow *os_window, Region region, unsign
     bind_framebuffer_for_output(temp_framebuffer);
     save_viewport_using_bottom_left_origin(0, 0, *thumb_w, *thumb_h);
 
-    // Use the blit program to render the scaled framebuffer
-    bind_program(BLIT_PROGRAM);
+    // Use the screenshot program to render the scaled framebuffer with proper color space handling
+    bind_program(SCREENSHOT_PROGRAM);
 
     // Set source rectangle (normalized coordinates: 0 to 1)
     // Note: OpenGL texture origin is bottom-left, but Region uses top-left origin
@@ -1489,10 +1497,13 @@ take_screenshot_of_rectangular_region(OSWindow *os_window, Region region, unsign
     float src_right_norm = (float)region.right / (float)vw;
     float src_bottom_norm = (float)(vh - region.bottom) / (float)vh;
     float src_top_norm = (float)(vh - region.top) / (float)vh;
-    glUniform4f(blit_program_layout.uniforms.src_rect, src_left_norm, src_top_norm, src_right_norm, src_bottom_norm);
+    glUniform4f(screenshot_program_layout.uniforms.src_rect, src_left_norm, src_top_norm, src_right_norm, src_bottom_norm);
 
     // Set destination rectangle (NDC coordinates: -1 to 1)
-    glUniform4f(blit_program_layout.uniforms.dest_rect, -1.0f, -1.0f, 1.0f, 1.0f);
+    glUniform4f(screenshot_program_layout.uniforms.dest_rect, -1.0f, -1.0f, 1.0f, 1.0f);
+
+    // Set the source texture size for proper downscaling
+    glUniform2f(screenshot_program_layout.uniforms.src_size, (float)vw, (float)vh);
 
     // Bind the source texture
     glActiveTexture(GL_TEXTURE0 + GRAPHICS_UNIT);
@@ -1631,7 +1642,7 @@ init_shaders(PyObject *module) {
 #define C(x) if (PyModule_AddIntConstant(module, #x, x) != 0) { PyErr_NoMemory(); return false; }
     C(CELL_PROGRAM); C(CELL_FG_PROGRAM); C(CELL_BG_PROGRAM); C(BORDERS_PROGRAM);
     C(GRAPHICS_PROGRAM); C(GRAPHICS_PREMULT_PROGRAM); C(GRAPHICS_ALPHA_MASK_PROGRAM);
-    C(BGIMAGE_PROGRAM); C(TINT_PROGRAM); C(TRAIL_PROGRAM); C(BLIT_PROGRAM); C(ROUNDED_RECT_PROGRAM);
+    C(BGIMAGE_PROGRAM); C(TINT_PROGRAM); C(TRAIL_PROGRAM); C(BLIT_PROGRAM); C(SCREENSHOT_PROGRAM); C(ROUNDED_RECT_PROGRAM);
     C(GLSL_VERSION);
     C(GL_VERSION);
     C(GL_VENDOR);

kitty/shaders.py

@@ -30,6 +30,7 @@ from .fast_data_types import (
     MARK_MASK,
     REVERSE,
     ROUNDED_RECT_PROGRAM,
+    SCREENSHOT_PROGRAM,
     STRIKETHROUGH,
     TINT_PROGRAM,
     TRAIL_PROGRAM,
@@ -221,6 +222,7 @@ class LoadShaderPrograms:
         program_for('tint').compile(TINT_PROGRAM, allow_recompile)
         program_for('trail').compile(TRAIL_PROGRAM, allow_recompile)
         program_for('blit').compile(BLIT_PROGRAM, allow_recompile)
+        program_for('screenshot').compile(SCREENSHOT_PROGRAM, allow_recompile)
         program_for('rounded_rect').compile(ROUNDED_RECT_PROGRAM, allow_recompile)
         init_cell_program()