Add ghost pixel thresholding support with thresholding flip button

handler.py

@@ -110,6 +110,7 @@ class EndpointHandler:
                 - 'inputs': A dictionary with the following keys:
                     - 'image_url' (str): The URL of the image to be upscaled.
                     - 'outscale' (float): The scaling factor for the upscaling process.
+                    - 'apply_thresholding' (bool): Whether to apply thresholding to the upscaled image.
 
         Returns:
             Dict[str, List[float]]: A dictionary containing the results of the upscaling process, which includes:
@@ -121,11 +122,13 @@ class EndpointHandler:
         ############################################################
         # get inputs and download image
         ############################################################
-        self.logger.info(">>> 1/7: GETTING INPUTS....")
+        self.logger.info(">>> 1/8: GETTING INPUTS....")
         try:
             inputs = data.pop("inputs", data)
             outscale = float(inputs.pop("outscale", 3))
+            apply_thresholding = inputs.pop("apply_thresholding", False)
             self.logger.info(f"outscale: {outscale}")
+            self.logger.info(f"apply_thresholding: {apply_thresholding}")
             image_url = inputs["image_url"]
         except Exception as e:
             self.logger.error(f"Error getting inputs: {e}")
@@ -143,7 +146,7 @@ class EndpointHandler:
         ############################################################
         # run assertions
         ############################################################
-        self.logger.info(">>> 2/7: RUNNING ASSERTIONS ON IMAGE....")
+        self.logger.info(">>> 2/8: RUNNING ASSERTIONS ON IMAGE....")
         
         # get image size and mode
         in_size, in_mode = image.size, image.mode
@@ -151,7 +154,7 @@ class EndpointHandler:
         
         # check image size and mode and return dict
         try:
-            assert in_mode in ["RGB", "RGBA", "L"], f"Unsupported image mode: {in_mode}"
+            assert in_mode in Image.MODES, f"Unsupported image mode: {in_mode}"
             if self.tiling_size == 0:
                 assert in_size[0] * in_size[1] <  self.max_image_size, f"Image is too large: {in_size}: {in_size[0] * in_size[1]} is greater than {self.max_image_size}"
             assert outscale > 1 and outscale <= 10, f"Outscale must be between 1 and 10: {outscale}"
@@ -163,7 +166,7 @@ class EndpointHandler:
         ############################################################
         # Convert RGB to BGR (PIL uses RGB, OpenCV expects BGR)
         ############################################################
-        self.logger.info(f">>> 3/7: CONVERTING IMAGE TO OPENCV BGR/BGRA FORMAT....")
+        self.logger.info(f">>> 3/8: CONVERTING IMAGE TO OPENCV BGR/BGRA FORMAT....")
         try:
             opencv_image = np.array(image)
         except Exception as e:
@@ -188,7 +191,7 @@ class EndpointHandler:
         ############################################################
         # upscale image
         ############################################################
-        self.logger.info(f">>> 4/7: UPSCALING IMAGE....")
+        self.logger.info(f">>> 4/8: UPSCALING IMAGE....")
         
         try:
             output, _ = self.model.enhance(opencv_image, outscale=outscale)
@@ -202,7 +205,7 @@ class EndpointHandler:
         ############################################################
         # convert to RGB/RGBA format
         ############################################################
-        self.logger.info(f">>> 5/7: CONVERTING IMAGE TO RGB/RGBA FORMAT....")
+        self.logger.info(f">>> 5/8: CONVERTING IMAGE TO RGB/RGBA FORMAT....")
         out_shape = output.shape
         if len(out_shape) == 3:
             if out_shape[2] == 3:
@@ -213,22 +216,39 @@ class EndpointHandler:
             output = cv2.cvtColor(output, cv2.COLOR_GRAY2RGB)
         
         
-        ############################################################
+        ###########################################################
         # convert to PIL image
         ############################################################
-        self.logger.info(f">>> 6/7: CONVERTING IMAGE TO PIL....")
-        try:
-            img_byte_arr = BytesIO()
+        self.logger.info(f">>> 6/8: CONVERTING IMAGE TO PIL....")
+        try:#
             output = Image.fromarray(output)
         except Exception as e:
             self.logger.error(f"Error converting upscaled image to PIL: {e}")
             return {"image_url": None, "image_key": None, "error": f"Failed to convert upscaled image to PIL: {e}"}
         
         
+        ############################################################
+        # apply thresholding
+        ############################################################
+        self.logger.info(f">>> 7/8: APPLYING THRESHOLDING....")
+        if apply_thresholding:
+            try:
+                if self.has_alpha_channel(image):
+                    self.logger.info(f"input image mode: {image.mode}, it has alpha channel, applying thresholding")
+                    output = self.apply_thresholding(image, output)
+                else:
+                    self.logger.info(f"input image mode: {image.mode}, it does not have alpha channel, skipping thresholding")
+            except Exception as e:
+                self.logger.error(f"Error applying thresholding: {e}")
+                return {"image_url": None, "image_key": None, "error": f"Failed to apply thresholding: {e}"}
+        else:
+            self.logger.info(f"thresholding is not enabled, skipping thresholding")
+        
+        
         ############################################################
         # upload to s3
         ############################################################
-        self.logger.info(f">>> 7/7: UPLOADING IMAGE TO S3....")
+        self.logger.info(f">>> 8/8: UPLOADING IMAGE TO S3....")
         try:
             image_url, key = self.upload_to_s3(output)
             self.logger.info(f"image uploaded to s3: {image_url}")
@@ -280,4 +300,99 @@ class EndpointHandler:
         """
         response = requests.get(image_url)
         image = Image.open(BytesIO(response.content))
-        return image
+        return image
+    
+    def has_alpha_channel(self, image: Image.Image) -> bool:
+        """
+        Check if the image has an alpha channel.
+        """
+        return image.mode in ("RGBA", "LA") or (image.mode == "P" and "transparency" in image.info)
+
+    def replace_mask_with_edge_aware_clamping(
+        self,
+        original_img: Image.Image,
+        esrgan_img: Image.Image,
+        edge_buffer: int = 2,
+        smooth_buffer: int = 10,
+        default_lower: int = 10,
+        default_upper: int = 245
+    ) -> Image.Image:
+        """
+        Apply edge-aware alpha channel clamping to merge original and ESRGAN images.
+        
+        Args:
+            original_img: Original image as PIL Image
+            esrgan_img: ESRGAN upscaled image as PIL Image
+            edge_buffer: Buffer for edge regions
+            smooth_buffer: Buffer for smooth regions
+            default_lower: Default lower threshold
+            default_upper: Default upper threshold
+            
+        Returns:
+            PIL Image with merged alpha channel
+        """
+        # Convert images to RGBA if needed
+        original_img = original_img.convert("RGBA")
+        esrgan_img = esrgan_img.convert("RGBA")
+        esr_w, esr_h = esrgan_img.size
+
+        # Upscale original alpha with bicubic resampling
+        orig_alpha = original_img.getchannel("A")
+        upscaled_alpha = orig_alpha.resize((esr_w, esr_h), resample=Image.Resampling.BICUBIC)
+        alpha_arr = np.array(upscaled_alpha, dtype=np.uint8)
+
+        # Edge detection on original alpha (for alignment)
+        edge_mask = cv2.Canny(alpha_arr, threshold1=50, threshold2=150)
+        edge_mask = (edge_mask > 0).astype(np.uint8) * 255
+
+        # Determine dynamic cutoffs for edge vs. smooth regions
+        semi_transparent_mask = (alpha_arr > 0) & (alpha_arr < 255)
+        if np.any(semi_transparent_mask):
+            min_alpha = int(np.min(alpha_arr[semi_transparent_mask]))
+            max_alpha = int(np.max(alpha_arr[semi_transparent_mask]))
+            lower_edge = max(0, min_alpha - edge_buffer)
+            upper_edge = min(255, max_alpha + edge_buffer)
+            lower_smooth = max(0, min_alpha - smooth_buffer)
+            upper_smooth = min(255, max_alpha + smooth_buffer)
+        else:
+            lower_edge = lower_smooth = default_lower
+            upper_edge = upper_smooth = default_upper
+
+        # Apply selective clamping based on edge regions
+        clamped_alpha = alpha_arr.copy()
+        clamped_alpha[edge_mask == 255] = np.clip(clamped_alpha[edge_mask == 255], lower_edge, upper_edge)
+        clamped_alpha[edge_mask == 0] = np.clip(clamped_alpha[edge_mask == 0], lower_smooth, upper_smooth)
+
+        # Blend with ESRGAN's alpha (50% weight to original in edges)
+        esrgan_alpha = np.array(esrgan_img.getchannel("A"), dtype=np.uint8)
+        blended_alpha = np.where(edge_mask == 255,
+                                (clamped_alpha * 0.5 + esrgan_alpha * 0.5).astype(np.uint8),
+                                clamped_alpha)
+
+        # Replace and return
+        esrgan_array = np.array(esrgan_img, dtype=np.uint8)
+        esrgan_array[..., 3] = blended_alpha
+        return Image.fromarray(esrgan_array, mode="RGBA")
+
+    def apply_thresholding(self, input_image: Image.Image, esrgan_image: Image.Image):
+        """
+        Apply thresholding to the image to account for ghost pixels
+        
+        Args:
+            input_image (Image.Image): The input image.
+            esrgan_image (Image.Image): The ESRGAN image.
+        """
+        
+        self.logger.info(f"Applying edge-aware clamping")
+        
+        # Using method 3: Edge-aware clamping
+        self.logger.info(f"Input image size: {input_image.size}")
+        edge_aware_image = self.replace_mask_with_edge_aware_clamping(
+            input_image,
+            esrgan_image,
+            edge_buffer=2,
+            smooth_buffer=10
+        )
+        self.logger.info(f"Edge-aware clamping applied, image size: {edge_aware_image.size}")
+        
+        return edge_aware_image

test_api_endpoint.ipynb

@@ -55,7 +55,7 @@
     "out_scales = [3, 3, 2.49]\n",
     "for img_url, outscale in zip(img_urls, out_scales):\n",
     "\t# create payload\n",
-    "\tpayload = {\"inputs\": {\"image_url\": img_url, \"outscale\": outscale}}\n",
+    "\tpayload = {\"inputs\": {\"image_url\": img_url, \"outscale\": outscale, \"apply_thresholding\": True}}\n",
     "\t\t\n",
     "\tresponse = requests.post(API_URL, headers=headers, json=payload)\n",
     "\toutput_payload = response.json()\t\n",

test_handler.ipynb

@@ -9,7 +9,7 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "/usr/local/lib/python3.10/dist-packages/torchvision/transforms/functional_tensor.py:5: UserWarning: The torchvision.transforms.functional_tensor module is deprecated in 0.15 and will be **removed in 0.17**. Please don't rely on it. You probably just need to use APIs in torchvision.transforms.functional or in torchvision.transforms.v2.functional.\n",
+      "/workspace/miniconda3/envs/real-esrgan-hf/lib/python3.10/site-packages/torchvision/transforms/functional_tensor.py:5: UserWarning: The torchvision.transforms.functional_tensor module is deprecated in 0.15 and will be **removed in 0.17**. Please don't rely on it. You probably just need to use APIs in torchvision.transforms.functional or in torchvision.transforms.v2.functional.\n",
       "  warnings.warn(\n"
      ]
     }
@@ -48,7 +48,7 @@
       "INFO - model_path: /workspace/real-esrgan/weights/Real-ESRGAN-x4plus.pth\n",
       "INFO - TILING_SIZE is 0, not using tiling\n",
       "INFO - initializing model\n",
-      "INFO - model initialized in 1.977891206741333 seconds\n"
+      "INFO - model initialized in 90.91657543182373 seconds\n"
      ]
     }
    ],
@@ -59,49 +59,79 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": 6,
    "metadata": {},
    "outputs": [
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "INFO - >>> 1/7: GETTING INPUTS....\n",
-      "INFO - outscale: 4.0\n",
-      "INFO - downloading image from URL: https://jiffy-transfers.imgix.net/2/attachments/r267odvvfmkp6c5lccj1y6f9trb0?ixlib=rb-0.3.5\n",
-      "INFO - >>> 2/7: RUNNING ASSERTIONS ON IMAGE....\n",
-      "INFO - image.size: (1024, 1024), image.mode: RGB\n",
-      "INFO - >>> 3/7: CONVERTING IMAGE TO OPENCV BGR/BGRA FORMAT....\n",
-      "INFO - converting RGB image to BGR\n",
-      "INFO - >>> 4/7: UPSCALING IMAGE....\n",
-      "INFO - output.shape: (4096, 4096, 3)\n",
-      "INFO - >>> 5/7: CONVERTING IMAGE TO RGB/RGBA FORMAT....\n",
-      "INFO - >>> 6/7: CONVERTING IMAGE TO PIL....\n",
-      "INFO - >>> 7/7: UPLOADING IMAGE TO S3....\n",
-      "INFO - image uploaded to s3: https://jiffy-staging-upscaled-images.s3.amazonaws.com/25b91e15-b785-47ca-81a6-ad5fbdf8b92a.png\n"
+      "INFO - >>> 1/8: GETTING INPUTS....\n",
+      "INFO - outscale: 2.3\n",
+      "INFO - apply_thresholding: True\n",
+      "INFO - downloading image from URL: https://jiffy-staging-transfers.imgix.net/2/ai-research/test-ghost-pixels/10816322.png\n",
+      "INFO - >>> 2/8: RUNNING ASSERTIONS ON IMAGE....\n",
+      "INFO - image.size: (1080, 1080), image.mode: RGBA\n",
+      "INFO - >>> 3/8: CONVERTING IMAGE TO OPENCV BGR/BGRA FORMAT....\n",
+      "INFO - converting RGBA image to BGRA\n",
+      "INFO - >>> 4/8: UPSCALING IMAGE....\n"
      ]
     },
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "https://jiffy-staging-upscaled-images.s3.amazonaws.com/25b91e15-b785-47ca-81a6-ad5fbdf8b92a.png 25b91e15-b785-47ca-81a6-ad5fbdf8b92a.png\n"
+      "https://jiffy-staging-transfers.imgix.net/2/ai-research/test-ghost-pixels/10816322.png 2.3\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "INFO - output.shape: (2484, 2484, 4)\n",
+      "INFO - >>> 5/8: CONVERTING IMAGE TO RGB/RGBA FORMAT....\n",
+      "INFO - >>> 6/8: CONVERTING IMAGE TO PIL....\n",
+      "INFO - >>> 7/8: APPLYING THRESHOLDING....\n",
+      "INFO - input image mode: RGBA, it has alpha channel, applying thresholding\n",
+      "INFO - Applying edge-aware clamping\n",
+      "INFO - Input image size: (1080, 1080)\n",
+      "INFO - Edge-aware clamping applied, image size: (2484, 2484)\n",
+      "INFO - >>> 8/8: UPLOADING IMAGE TO S3....\n",
+      "INFO - image uploaded to s3: https://jiffy-staging-upscaled-images.s3.amazonaws.com/0cd7ea5b-a6fc-4cbb-a536-57d9dbcc12b4.png\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "https://jiffy-staging-upscaled-images.s3.amazonaws.com/0cd7ea5b-a6fc-4cbb-a536-57d9dbcc12b4.png 0cd7ea5b-a6fc-4cbb-a536-57d9dbcc12b4.png\n"
      ]
     }
    ],
    "source": [
     "img_dir = \"test_data/\"\n",
-    "img_urls = [\"https://jiffy-transfers.imgix.net/2/attachments/r267odvvfmkp6c5lccj1y6f9trb0?ixlib=rb-0.3.5\",\n",
-    "            # \"https://jiffy-staging-transfers.imgix.net/2/development/attachments/zo31eau0ykhbwoddrjtlbyz6w9mp?ixlib=rb-0.3.5\", # larger than > 1.96M pixels\n",
-    "            # \"https://jiffy-staging-transfers.imgix.net/2/development/attachments/b8ecchms9rr9wk3g71kfpfprqg1v?ixlib=rb-0.3.5\" # larger than > 1.96M pixels\n",
+    "img_urls = [#\"https://jiffy-transfers.imgix.net/2/attachments/r267odvvfmkp6c5lccj1y6f9trb0\",\n",
+    "#             \"https://jiffy-staging-transfers.imgix.net/2/development/attachments/zo31eau0ykhbwoddrjtlbyz6w9mp\", # larger than > 1.96M pixels\n",
+    "#             \"https://jiffy-staging-transfers.imgix.net/2/development/attachments/b8ecchms9rr9wk3g71kfpfprqg1v\", # larger than > 1.96M pixels,\n",
+    "#             \"https://jiffy-staging-transfers.imgix.net/2/ai-research/test-ghost-pixels/9B500B09.png\", # thresholding example 1\n",
+    "            \"https://jiffy-staging-transfers.imgix.net/2/ai-research/test-ghost-pixels/10816322.png\"  # thresholding example 2\n",
+    "\n",
     "           ]\n",
     "\n",
-    "out_scales = [4]#, 3, 2]\n",
+    "out_scales = [#4, \n",
+    "              #3, \n",
+    "              #2, \n",
+    "              #2,\n",
+    "              2.3\n",
+    "]\n",
+    "\n",
     "for img_url, outscale in zip(img_urls, out_scales):\n",
+    "    print(img_url, outscale)\n",
     "    # create payload\n",
     "    payload = {\n",
     "        \"inputs\": {\"image_url\": img_url, \n",
-    "                    \"outscale\": outscale\n",
+    "                    \"outscale\": outscale,\n",
+    "                   \"apply_thresholding\": True,\n",
     "                    }\n",
     "            }\n",
     "    \n",
@@ -134,7 +164,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.10.12"
+   "version": "3.10.16"
   }
  },
  "nbformat": 4,