From 8e408640a4f2c9988c677cd3c370a667c40413fa Mon Sep 17 00:00:00 2001
From: Anders Haglund <andershaglund@gmail.com>
Date: Wed, 23 Nov 2022 14:50:58 -0800
Subject: [PATCH] Match stable diffusion result's histogram to image's

---
 lama_cleaner/app/src/adapters/inpainting.ts   |  1 +
 .../src/components/SidePanel/SidePanel.tsx    | 16 ++++++++
 lama_cleaner/app/src/store/Atoms.tsx          |  2 +
 lama_cleaner/model/base.py                    | 41 +++++++++++++++++++
 lama_cleaner/schema.py                        |  1 +
 lama_cleaner/server.py                        |  1 +
 6 files changed, 62 insertions(+)

diff --git a/lama_cleaner/app/src/adapters/inpainting.ts b/lama_cleaner/app/src/adapters/inpainting.ts
index 61056be..11b9ebe 100644
--- a/lama_cleaner/app/src/adapters/inpainting.ts
+++ b/lama_cleaner/app/src/adapters/inpainting.ts
@@ -54,6 +54,7 @@ export default async function inpaint(
   fd.append('sdGuidanceScale', settings.sdGuidanceScale.toString())
   fd.append('sdSampler', settings.sdSampler.toString())
   fd.append('sdSeed', seed ? seed.toString() : '-1')
+  fd.append('sdMatchHistograms', settings.sdMatchHistograms ? 'true' : 'false')
 
   fd.append('cv2Radius', settings.cv2Radius.toString())
   fd.append('cv2Flag', settings.cv2Flag.toString())
diff --git a/lama_cleaner/app/src/components/SidePanel/SidePanel.tsx b/lama_cleaner/app/src/components/SidePanel/SidePanel.tsx
index c2fd1a5..ca6dc43 100644
--- a/lama_cleaner/app/src/components/SidePanel/SidePanel.tsx
+++ b/lama_cleaner/app/src/components/SidePanel/SidePanel.tsx
@@ -120,6 +120,22 @@ const SidePanel = () => {
               }}
             />
 
+            <SettingBlock
+              title="Match Histograms"
+              input={
+                <Switch
+                  checked={setting.sdMatchHistograms}
+                  onCheckedChange={value => {
+                    setSettingState(old => {
+                      return { ...old, sdMatchHistograms: value }
+                    })
+                  }}
+                >
+                  <SwitchThumb />
+                </Switch>
+              }
+            />
+
             <SettingBlock
               className="sub-setting-block"
               title="Sampler"
diff --git a/lama_cleaner/app/src/store/Atoms.tsx b/lama_cleaner/app/src/store/Atoms.tsx
index 14a57f3..54d1d78 100644
--- a/lama_cleaner/app/src/store/Atoms.tsx
+++ b/lama_cleaner/app/src/store/Atoms.tsx
@@ -175,6 +175,7 @@ export interface Settings {
   sdSeed: number
   sdSeedFixed: boolean // true: use sdSeed, false: random generate seed on backend
   sdNumSamples: number
+  sdMatchHistograms: boolean
 
   // For OpenCV2
   cv2Radius: number
@@ -278,6 +279,7 @@ export const settingStateDefault: Settings = {
   sdSeed: 42,
   sdSeedFixed: true,
   sdNumSamples: 1,
+  sdMatchHistograms: false,
 
   // CV2
   cv2Radius: 5,
diff --git a/lama_cleaner/model/base.py b/lama_cleaner/model/base.py
index 00b0e30..45f468f 100644
--- a/lama_cleaner/model/base.py
+++ b/lama_cleaner/model/base.py
@@ -56,6 +56,9 @@ class InpaintModel:
         result = self.forward(pad_image, pad_mask, config)
         result = result[0:origin_height, 0:origin_width, :]
 
+        if config.sd_match_histograms:
+            result = self._match_histograms(result, image[:, :, ::-1], mask)
+
         if config.sd_mask_blur != 0:
             k = 2 * config.sd_mask_blur + 1
             mask = cv2.GaussianBlur(mask, (k, k), 0)
@@ -172,6 +175,44 @@ class InpaintModel:
 
         return crop_img, crop_mask, [l, t, r, b]
 
+    def _calculate_cdf(self, histogram):
+        cdf = histogram.cumsum()
+        normalized_cdf = cdf / float(cdf.max())
+        return normalized_cdf
+    
+    def _calculate_lookup(self, source_cdf, reference_cdf):
+        lookup_table = np.zeros(256)
+        lookup_val = 0
+        for source_index, source_val in enumerate(source_cdf):
+            for reference_index, reference_val in enumerate(reference_cdf):
+                if reference_val >= source_val:
+                    lookup_val = reference_index
+                    break
+            lookup_table[source_index] = lookup_val
+        return lookup_table
+    
+    def _match_histograms(self, source, reference, mask):
+        transformed_channels = []
+        for channel in range(source.shape[-1]):
+            source_channel = source[:, :, channel]
+            reference_channel = reference[:, :, channel]
+    
+            # only calculate histograms for non-masked parts
+            source_histogram, _ = np.histogram(source_channel[mask == 0], 256, [0,256])
+            reference_histogram, _ = np.histogram(reference_channel[mask == 0], 256, [0,256])    
+        
+            source_cdf = self._calculate_cdf(source_histogram)
+            reference_cdf = self._calculate_cdf(reference_histogram)
+        
+            lookup = self._calculate_lookup(source_cdf, reference_cdf)
+        
+            transformed_channels.append(cv2.LUT(source_channel, lookup))
+    
+        result = cv2.merge(transformed_channels)
+        result = cv2.convertScaleAbs(result)
+    
+        return result
+
     def _run_box(self, image, mask, box, config: Config):
         """
 
diff --git a/lama_cleaner/schema.py b/lama_cleaner/schema.py
index 4a374c5..21b0c69 100644
--- a/lama_cleaner/schema.py
+++ b/lama_cleaner/schema.py
@@ -48,6 +48,7 @@ class Config(BaseModel):
     sd_sampler: str = SDSampler.ddim
     # -1 mean random seed
     sd_seed: int = 42
+    sd_match_histograms: bool = False
 
     # cv2
     cv2_flag: str = 'INPAINT_NS'
diff --git a/lama_cleaner/server.py b/lama_cleaner/server.py
index 33b8079..71677de 100644
--- a/lama_cleaner/server.py
+++ b/lama_cleaner/server.py
@@ -131,6 +131,7 @@ def process():
         sd_guidance_scale=form["sdGuidanceScale"],
         sd_sampler=form["sdSampler"],
         sd_seed=form["sdSeed"],
+        sd_match_histograms=form["sdMatchHistograms"],
         cv2_flag=form["cv2Flag"],
         cv2_radius=form['cv2Radius']
     )