Handle ICC profiles for PSDs

hydrus/client/ClientFiles.py

@@ -6,7 +6,7 @@ import threading
 import time
 import typing
 
-from hydrus.core import HydrusConstants as HC
+from hydrus.core import HydrusConstants as HC, HydrusPSDHandling
 from hydrus.core import HydrusData
 from hydrus.core import HydrusExceptions
 from hydrus.core import HydrusFileHandling
@@ -2006,22 +2006,34 @@ class FilesMaintenanceManager( object ):
         if mime not in HC.FILES_THAT_CAN_HAVE_ICC_PROFILE:
             
             return False
-            
         
         try:
             
             path = self._controller.client_files_manager.GetFilePath( hash, mime )
             
-            try:
+            if mime == HC.APPLICATION_PSD:
+
+                try:
+
+                    has_icc_profile = HydrusPSDHandling.PSDHasICCProfile(path)
+
+                except:
+
+                    return None
                 
-                pil_image = HydrusImageHandling.RawOpenPILImage( path )
-                
-            except:
-                
-                return None
+            else:
+
+                try:
+                    
+                    pil_image = HydrusImageHandling.RawOpenPILImage( path )
+                    
+                except:
+                    
+                    return None
                 
             
-            has_icc_profile = HydrusImageHandling.HasICCProfile( pil_image )
+                has_icc_profile = HydrusImageHandling.HasICCProfile( pil_image )
+            
             
             additional_data = has_icc_profile
             

hydrus/client/importing/ClientImportFiles.py

@@ -1,6 +1,6 @@
 import typing
 
-from hydrus.core import HydrusConstants as HC
+from hydrus.core import HydrusConstants as HC, HydrusPSDHandling
 from hydrus.core import HydrusData
 from hydrus.core import HydrusExceptions
 from hydrus.core import HydrusFileHandling
@@ -435,9 +435,15 @@ class FileImportJob( object ):
             
             try:
                 
-                pil_image = HydrusImageHandling.RawOpenPILImage( self._temp_path )
+                if mime == HC.APPLICATION_PSD:
+
+                    has_icc_profile = HydrusPSDHandling.PSDHasICCProfile( self._temp_path )
+
+                else:
                 
-                has_icc_profile = HydrusImageHandling.HasICCProfile( pil_image )
+                    pil_image = HydrusImageHandling.RawOpenPILImage( self._temp_path )
+                    
+                    has_icc_profile = HydrusImageHandling.HasICCProfile( pil_image )
                 
             except:
                 

hydrus/core/HydrusConstants.py

@@ -770,7 +770,7 @@ APPLICATIONS_WITH_THUMBNAILS = set({ IMAGE_SVG, APPLICATION_FLASH, APPLICATION_C
 
 MIMES_WITH_THUMBNAILS = set( IMAGES ).union( ANIMATIONS ).union( VIDEO ).union( APPLICATIONS_WITH_THUMBNAILS )
 
-FILES_THAT_CAN_HAVE_ICC_PROFILE = { IMAGE_JPEG, IMAGE_PNG, IMAGE_GIF, IMAGE_TIFF }
+FILES_THAT_CAN_HAVE_ICC_PROFILE = { IMAGE_JPEG, IMAGE_PNG, IMAGE_GIF, IMAGE_TIFF, APPLICATION_PSD }
 
 FILES_THAT_CAN_HAVE_EXIF = { IMAGE_JPEG, IMAGE_TIFF, IMAGE_PNG, IMAGE_WEBP }
 # images and animations that PIL can handle

hydrus/core/HydrusImageHandling.py

@@ -302,6 +302,8 @@ def GenerateNumPyImage( path, mime, force_pil = False ) -> numpy.array:
             HydrusData.ShowText( 'Loading PSD' )
             
         pil_image = HydrusPSDHandling.MergedPILImageFromPSD( path )
+ 
+        pil_image = DequantizePILImage( pil_image )
         
         numpy_image = GenerateNumPyImageFromPILImage( pil_image )
 

hydrus/core/HydrusPSDHandling.py

@@ -8,6 +8,7 @@ from hydrus.core import HydrusImageHandling
 try:
     
     from psd_tools import PSDImage
+    from psd_tools.constants import Resource
     
     PSD_TOOLS_OK = True
     
@@ -15,6 +16,16 @@ except:
     
     PSD_TOOLS_OK = False
 
+def PSDHasICCProfile(path: str):
+
+    if not PSD_TOOLS_OK:
+        
+        raise Exception( 'psd_tools unavailable' )
+
+    psd = PSDImage.open(path)
+
+    return Resource.ICC_PROFILE in psd.image_resources
+
 
 def MergedPILImageFromPSD(path: str) -> PILImage:
 
@@ -24,11 +35,12 @@ def MergedPILImageFromPSD(path: str) -> PILImage:
     
     psd = PSDImage.open(path)
 
-    pil_image = psd.topil()
+
+    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):
+    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
@@ -37,6 +49,13 @@ def MergedPILImageFromPSD(path: str) -> PILImage:
         # that has to happen for the thumbnail anyway.
         pil_image = pil_image.convert("RGB")
 
+
+    if Resource.ICC_PROFILE in psd.image_resources:
+        
+        icc = psd.image_resources.get_data(Resource.ICC_PROFILE)
+        
+        pil_image.info['icc_profile'] = icc
+
     return pil_image