Update image_enhancer.oy

image_enhancer.oy

@@ -3,6 +3,7 @@ import torch
 from gfpgan import GFPGANer
 from tqdm import tqdm
 import cv2
+import warnings
 from enum import Enum
 
 class EnhancementMethod(str, Enum):
@@ -13,112 +14,108 @@ class EnhancementMethod(str, Enum):
 
 
 class Enhancer:
-    def __init__(self, method=EnhancementMethod, background_enhancement=True, upscale=2):
-        # Set up RealESRGAN for background enhancement
-        if background_enhancement:
-            if upscale == 2:
-                if not torch.cuda.is_available(): # CPU
-                    import warnings
-                    warnings.warn('The unoptimized RealESRGAN is slow on CPU. We do not use it. '
-                                'If you really want to use it, please modify the corresponding codes.')
-                    self.bg_upsampler = None
-                else:
-                    from basicsr.archs.rrdbnet_arch import RRDBNet
-                    from realesrgan import RealESRGANer
-                    model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2)
-                    self.bg_upsampler = RealESRGANer(
-                        scale=2,
-                        model_path='https://huggingface.co/dtarnow/UPscaler/resolve/main/RealESRGAN_x2plus.pth',
-                        model=model,
-                        tile=400,
-                        tile_pad=10,
-                        pre_pad=0,
-                        half=True)  # need to set False in CPU mode
-            elif upscale == 4:
-                if not torch.cuda.is_available(): # CPU
-                    import warnings
-                    warnings.warn('The unoptimized RealESRGAN is slow on CPU. We do not use it. '
-                                'If you really want to use it, please modify the corresponding codes.')
-                    self.bg_upsampler = None
-                else:
-                    from basicsr.archs.rrdbnet_arch import RRDBNet
-                    from realesrgan import RealESRGANer
-                    model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=4)
-                    self.bg_upsampler = RealESRGANer(
-                        scale=4,
-                        model_path='https://huggingface.co/lllyasviel/Annotators/resolve/main/RealESRGAN_x4plus.pth',
-                        model=model,
-                        tile=400,
-                        tile_pad=10,
-                        pre_pad=0,
-                        half=True)  # need to set False in CPU mode
-            else:
-                raise ValueError(f'Wrong upscale constant {upscale}.')
-        else:
-            self.bg_upsampler = None
+    def __init__(self, method: EnhancementMethod, background_enhancement=True, upscale=2):
+        self.method = method
+        self.background_enhancement = background_enhancement
+        self.upscale = upscale
+        self.bg_upsampler = None
+        self.realesrgan_enhancer = None
 
-        # Set up GPFGAN for face enhancement
-        if  method == 'gfpgan':
-            self.arch = 'clean'
-            self.channel_multiplier = 2
-            self.model_name = 'GFPGANv1.4'
-            self.url = 'https://huggingface.co/gmk123/GFPGAN/resolve/main/GFPGANv1.4.pth'
-        elif method == 'RestoreFormer':
-            self.arch = 'RestoreFormer'
-            self.channel_multiplier = 2
-            self.model_name = 'RestoreFormer'
-            self.url = 'https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/RestoreFormer.pth'
-        elif method == 'codeformer': # TODO:
-            self.arch = 'CodeFormer'
-            self.channel_multiplier = 2
-            self.model_name = 'CodeFormer'
-            self.url = 'https://huggingface.co/sinadi/aar/resolve/main/codeformer.pth'
+        if self.method != EnhancementMethod.realesrgan:
+            self.setup_face_enhancer()
+            if self.background_enhancement:
+                self.setup_background_enhancer()
         else:
-            raise ValueError(f'Wrong model version {method}.')
-        
-        # Determine the model path and if the model is not available, download it
-        model_path = os.path.join('gfpgan/weights', self.model_name + '.pth')
-        
+            self.setup_realesrgan_enhancer()
+
+    def setup_background_enhancer(self):
+        if not torch.cuda.is_available():
+            warnings.warn('The unoptimized RealESRGAN is slow on CPU. We do not use it.')
+            return
+
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=self.upscale)
+        model_path = f'https://huggingface.co/dtarnow/UPscaler/resolve/main/RealESRGAN_x{self.upscale}plus.pth'
+        self.bg_upsampler = RealESRGANer(
+            scale=self.upscale,
+            model_path=model_path,
+            model=model,
+            tile=400,
+            tile_pad=10,
+            pre_pad=0,
+            half=True)
+
+    def setup_realesrgan_enhancer(self):
+        if not torch.cuda.is_available():
+            raise ValueError('CUDA is not available for RealESRGAN')
+
+        model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=self.upscale)
+        model_path = f'https://huggingface.co/dtarnow/UPscaler/resolve/main/RealESRGAN_x{self.upscale}plus.pth'
+        self.realesrgan_enhancer = RealESRGANer(
+            scale=self.upscale,
+            model_path=model_path,
+            model=model,
+            tile=400,
+            tile_pad=10,
+            pre_pad=0,
+            half=True)
+
+    def setup_face_enhancer(self):
+        model_configs = {
+            EnhancementMethod.gfpgan: {
+                'arch': 'clean',
+                'channel_multiplier': 2,
+                'model_name': 'GFPGANv1.4',
+                'url': 'https://huggingface.co/gmk123/GFPGAN/resolve/main/GFPGANv1.4.pth'
+            },
+            EnhancementMethod.RestoreFormer: {
+                'arch': 'RestoreFormer',
+                'channel_multiplier': 2,
+                'model_name': 'RestoreFormer',
+                'url': 'https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/RestoreFormer.pth'
+            },
+            EnhancementMethod.codeformer: {
+                'arch': 'CodeFormer',
+                'channel_multiplier': 2,
+                'model_name': 'CodeFormer',
+                'url': 'https://huggingface.co/sinadi/aar/resolve/main/codeformer.pth'
+            }
+        }
+
+        config = model_configs.get(self.method)
+        if not config:
+            raise ValueError(f'Wrong model version {self.method}')
+
+        model_path = os.path.join('gfpgan/weights', config['model_name'] + '.pth')
         if not os.path.isfile(model_path):
-            model_path = os.path.join('checkpoints', self.model_name + '.pth')
-        
+            model_path = os.path.join('checkpoints', config['model_name'] + '.pth')
         if not os.path.isfile(model_path):
-            # Download pre-trained models from url
-            model_path = self.url
+            model_path = config['url']
 
-        self.restorer = GFPGANer(
+        self.face_enhancer = GFPGANer(
             model_path=model_path,
-            upscale=upscale,
-            arch=self.arch,
-            channel_multiplier=self.channel_multiplier,
+            upscale=self.upscale,
+            arch=config['arch'],
+            channel_multiplier=config['channel_multiplier'],
             bg_upsampler=self.bg_upsampler)
-        
 
-    def check_image_dimensions(self, image):
-        # Get the dimensions of the image
+    def check_image_resolution(self, image):
         height, width, _ = image.shape
-        return True
-
-        # Check if either dimension exceeds 2048 pixels :Todo
-        # if width > 2048 or height > 2048:
-        #     return True
+        return width, height
 
-        # else:
-        #     print("Image dimensions are within the limit.")
-        #     return True
-        
-
-    def enhance(self, image):
+    async def enhance(self, image):
         img = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
-        if self.check_image_dimensions(img):
-            cropped_faces, restored_faces, r_img = self.restorer.enhance(
+        width, height = self.check_image_resolution(img)
+        
+        if self.method == EnhancementMethod.realesrgan:
+            enhanced_img, _ = await asyncio.to_thread(self.realesrgan_enhancer.enhance, img, outscale=self.upscale)
+        else:
+            _, _, enhanced_img = await asyncio.to_thread(self.face_enhancer.enhance,
                 img,
                 has_aligned=False,
                 only_center_face=False,
                 paste_back=True)
-        else:
-            r_img = img
         
-        r_img = cv2.cvtColor(r_img, cv2.COLOR_BGR2RGB)
+        enhanced_img = cv2.cvtColor(enhanced_img, cv2.COLOR_BGR2RGB)
+        enhanced_width, enhanced_height = self.check_image_resolution(enhanced_img)
 
-        return r_img
+        return enhanced_img, (width, height), (enhanced_width, enhanced_height)