From 6098eaf282fa17534e52105f5aa66de94595a868 Mon Sep 17 00:00:00 2001
From: Paul Friederichsen <floogulinc@gmail.com>
Date: Sat, 9 Sep 2023 14:39:09 -0500
Subject: [PATCH] Improve PSD handling and reduce memory usage (#1434)

* Implement our own PSD to PIL conversion

Without inefficient remove_white_background step.

* Stop trying to load PSDs directly with PIL

* Formatting

* Remove unused imports
---
 hydrus/core/HydrusImageHandling.py | 31 +++++++------
 hydrus/core/HydrusPSDHandling.py   | 70 ++++++++++++++++++++++++------
 2 files changed, 74 insertions(+), 27 deletions(-)

diff --git a/hydrus/core/HydrusImageHandling.py b/hydrus/core/HydrusImageHandling.py
index 235e739b..f728df1e 100644
--- a/hydrus/core/HydrusImageHandling.py
+++ b/hydrus/core/HydrusImageHandling.py
@@ -270,7 +270,22 @@ def GenerateNumPyImage( path, mime, force_pil = False ) -> numpy.array:
         
         HydrusData.ShowText( 'Loading media: ' + path )
         
-    
+
+    if mime == HC.APPLICATION_PSD:
+        
+        if HG.media_load_report_mode:
+            
+            HydrusData.ShowText( 'Loading PSD' )
+            
+        pil_image = HydrusPSDHandling.MergedPILImageFromPSD( path )
+        
+        pil_image = DequantizePILImage( pil_image )
+        
+        numpy_image = GenerateNumPyImageFromPILImage( pil_image )
+
+        return StripOutAnyUselessAlphaChannel( numpy_image )
+        
+
     if not OPENCV_OK:
         
         force_pil = True
@@ -316,19 +331,7 @@ def GenerateNumPyImage( path, mime, force_pil = False ) -> numpy.array:
             pass
             
         
-    if mime == HC.APPLICATION_PSD:
-        
-        if HG.media_load_report_mode:
-            
-            HydrusData.ShowText( 'Loading PSD' )
-            
-        pil_image = HydrusPSDHandling.MergedPILImageFromPSD( path )
-        
-        pil_image = DequantizePILImage( pil_image )
-        
-        numpy_image = GenerateNumPyImageFromPILImage( pil_image )
-        
-    elif mime in PIL_ONLY_MIMETYPES or force_pil:
+    if mime in PIL_ONLY_MIMETYPES or force_pil:
         
         if HG.media_load_report_mode:
             
diff --git a/hydrus/core/HydrusPSDHandling.py b/hydrus/core/HydrusPSDHandling.py
index 01efe55a..49c573e0 100644
--- a/hydrus/core/HydrusPSDHandling.py
+++ b/hydrus/core/HydrusPSDHandling.py
@@ -8,7 +8,9 @@ from hydrus.core import HydrusExceptions, HydrusImageHandling
 try:
     
     from psd_tools import PSDImage
-    from psd_tools.constants import Resource
+    from psd_tools.constants import Resource, ColorMode, Resource
+    from psd_tools.api.numpy_io import has_transparency, get_transparency_index
+    from psd_tools.api.pil_io import get_pil_mode, get_pil_channels, _create_image
     
     PSD_TOOLS_OK = True
     
@@ -38,18 +40,15 @@ def MergedPILImageFromPSD( path: str ) -> PILImage:
     
     psd = PSDImage.open( path )
     
-    pil_image = psd.topil( apply_icc = False )
-    
-    no_alpha = psd._record.layer_and_mask_information.layer_info is not None and psd._record.layer_and_mask_information.layer_info.layer_count > 0
-    
-    if HydrusImageHandling.PILImageHasTransparency( pil_image ) and no_alpha:
-        # merged image from psd-tools has transparency when it shouldn't
-        # see https://github.com/psd-tools/psd-tools/issues/369
-        # and https://github.com/psd-tools/psd-tools/pull/370
-        
-        # I think it's fine to convert to RGB in all cases since eventually
-        # that has to happen for the thumbnail anyway.
-        pil_image = pil_image.convert( 'RGB' )
+    #pil_image = psd.topil( apply_icc = False )
+
+    if psd.has_preview():
+
+        pil_image = convert_image_data_to_pil(psd)
+
+    else:
+
+        raise HydrusExceptions.UnsupportedFileException('PSD file has no embedded preview!')
         
     
     if Resource.ICC_PROFILE in psd.image_resources:
@@ -104,3 +103,48 @@ def GetPSDResolutionFallback( path: str ):
     
     return ( width, height )
     
+
+# modified from psd-tools source:
+# https://github.com/psd-tools/psd-tools/blob/main/src/psd_tools/api/pil_io.py
+
+def convert_image_data_to_pil(psd: PSDImage):
+    alpha = None
+
+    channel_data = psd._record.image_data.get_data(psd._record.header)
+    size = (psd.width, psd.height)
+    
+    channels = [_create_image(size, c, psd.depth) for c in channel_data]
+
+    # has_transparency not quite correct
+    # see https://github.com/psd-tools/psd-tools/issues/369
+    # and https://github.com/psd-tools/psd-tools/pull/370
+    no_alpha = psd._record.layer_and_mask_information.layer_info is not None and psd._record.layer_and_mask_information.layer_info.layer_count > 0
+
+    if has_transparency(psd) and not no_alpha:
+        alpha = channels[get_transparency_index(psd)]
+
+    if psd.color_mode == ColorMode.INDEXED:
+        image = channels[0]
+        image.putpalette(psd._record.color_mode_data.interleave())
+    elif psd.color_mode == ColorMode.MULTICHANNEL:
+        image = channels[0]  # Multi-channel mode is a collection of alpha.
+    else:
+        mode = get_pil_mode(psd.color_mode)
+        image = PILImage.merge(mode, channels[:get_pil_channels(mode)])
+
+    if not image:
+        return None
+
+    return post_process(image, alpha)
+
+
+def post_process(image, alpha):
+    # Fix inverted CMYK.
+    if image.mode == 'CMYK':
+        from PIL import ImageChops
+        image = ImageChops.invert(image)
+
+    # In Pillow, alpha channel is only available in RGB or L.
+    if alpha and image.mode in ('RGB', 'L'):
+        image.putalpha(alpha)
+    return image