add tab

app.py

@@ -4,6 +4,7 @@ from lut_processor import create_lut_tab
 from sharpen_processor import create_sharpen_tab
 from color_match_processor import create_color_match_tab
 from simple_effects_processor import create_effects_tab
+from histogram_processor import create_histogram_tab
 
 with gr.Blocks(title="Image Processing Suite") as demo:
     gr.Markdown("# Image Processing Suite")
@@ -13,6 +14,7 @@ with gr.Blocks(title="Image Processing Suite") as demo:
     create_sharpen_tab()
     create_color_match_tab()
     create_effects_tab()
+    create_histogram_tab()
 
 if __name__ == "__main__":
     demo.launch(share=True)

histogram_processor.py

@@ -0,0 +1,129 @@
+import gradio as gr
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from skimage.exposure import match_histograms
+
+class HistogramMatcher(nn.Module):
+    def __init__(self, differentiable=False):
+        super(HistogramMatcher, self).__init__()
+        self.differentiable = differentiable
+
+    def forward(self, dst, ref):
+        B, C, H, W = dst.size()
+        hist_dst = self.cal_hist(dst)
+        hist_ref = self.cal_hist(ref)
+        tables = self.cal_trans_batch(hist_dst, hist_ref)
+        
+        rst = dst.clone()
+        for b in range(B):
+            for c in range(C):
+                rst[b,c] = tables[b*c, (dst[b,c] * 255).long()]
+        
+        return rst / 255.
+
+    def cal_hist(self, img):
+        B, C, H, W = img.size()
+        if self.differentiable:
+            hists = self.soft_histc_batch(img * 255, bins=256, min=0, max=256, sigma=75)
+        else:
+            hists = torch.stack([torch.histc(img[b,c] * 255, bins=256, min=0, max=255) 
+                               for b in range(B) for c in range(C)])
+        
+        hists = hists.float()
+        hists = F.normalize(hists, p=1)
+        bc, n = hists.size()
+        triu = torch.ones(bc, n, n, device=hists.device).triu()
+        hists = torch.bmm(hists[:,None,:], triu)[:,0,:]
+        return hists
+
+    def soft_histc_batch(self, x, bins=256, min=0, max=256, sigma=75):
+        B, C, H, W = x.size()
+        x = x.view(B*C, -1)
+        delta = float(max - min) / float(bins)
+        centers = float(min) + delta * (torch.arange(bins, device=x.device) + 0.5)
+        
+        x = torch.unsqueeze(x, 1)
+        centers = centers[None,:,None]
+        x = x - centers
+        x = torch.sigmoid(sigma * (x + delta/2)) - torch.sigmoid(sigma * (x - delta/2))
+        x = x.sum(dim=2)
+        return x
+
+    def cal_trans_batch(self, hist_dst, hist_ref):
+        hist_dst = hist_dst[:,None,:].repeat(1,256,1)
+        hist_ref = hist_ref[:,:,None].repeat(1,1,256)
+        table = hist_dst - hist_ref
+        table = torch.where(table>=0, 1., 0.)
+        table = torch.sum(table, dim=1) - 1
+        table = torch.clamp(table, min=0, max=255)
+        return table
+
+def apply_histogram_matching(image, reference, method, factor):
+    if image is None or reference is None:
+        return None
+        
+    # Convert to torch tensors and normalize
+    image = torch.from_numpy(image).float() / 255.0
+    reference = torch.from_numpy(reference).float() / 255.0
+    
+    # Add batch dimension and rearrange to BCHW
+    image = image.unsqueeze(0).permute(0, 3, 1, 2)
+    reference = reference.unsqueeze(0).permute(0, 3, 1, 2)
+    
+    if method == "pytorch":
+        # Apply PyTorch-based histogram matching
+        matcher = HistogramMatcher(differentiable=True)
+        matched = matcher(image, reference)
+    else:  # skimage
+        # Convert back to numpy for skimage
+        matched = match_histograms(
+            image.permute(0, 2, 3, 1).numpy(),
+            reference.permute(0, 2, 3, 1).numpy(),
+            channel_axis=3
+        )
+        matched = torch.from_numpy(matched).permute(0, 3, 1, 2)
+    
+    # Apply factor blending
+    result = factor * matched + (1 - factor) * image
+    
+    # Convert back to HWC format and to uint8
+    output = result.squeeze(0).permute(1, 2, 0)
+    output = (output.clamp(0, 1).numpy() * 255).astype(np.uint8)
+    
+    return output
+
+def create_histogram_tab():
+    with gr.Tab("Histogram Matching"):
+        gr.Markdown("Match histograms between images using PyTorch or scikit-image methods")
+        
+        with gr.Row():
+            with gr.Column():
+                input_image = gr.Image(label="Input Image")
+                reference_image = gr.Image(label="Reference Image")
+                
+                method = gr.Radio(
+                    choices=["pytorch", "skimage"],
+                    value="pytorch",
+                    label="Matching Method"
+                )
+                
+                factor = gr.Slider(
+                    minimum=0.0,
+                    maximum=1.0,
+                    value=1.0,
+                    step=0.05,
+                    label="Blend Factor"
+                )
+                
+                match_btn = gr.Button("Apply Histogram Matching")
+            
+            with gr.Column():
+                output_image = gr.Image(label="Matched Image")
+        
+        match_btn.click(
+            fn=apply_histogram_matching,
+            inputs=[input_image, reference_image, method, factor],
+            outputs=output_image
+        )

requirements.txt

@@ -3,4 +3,5 @@ torch
 torchvision
 numpy==1.26.4
 colour-science
-kornia
+kornia
+scikit-image