Add new function MSRCR

Add new function MSRCR. Add comments to the procedures.
HomeOfVapourSynthEvolution · Oct 28, 2014 · d4aa7bc · d4aa7bc
1 parent a48a710
commit d4aa7bc
Show file tree

Hide file tree

Showing 10 changed files with 669 additions and 91 deletions.
diff --git a/README.md b/README.md
@@ -6,7 +6,7 @@ VapourSynth plugin
 
 namespace: retinex
 
-functions: MSRCP
+functions: MSRCP, MSRCR
 
 ## About Retinex
 
@@ -94,3 +94,68 @@ i = core.lsmas.LWLibavSource(r'Image.png')
 i = core.fmtc.bitdepth(i, bits=16)
 i = core.retinex.MSRCP(i)
 ```
+
+## MSRCR
+
+### Description
+
+MSRCR(Multi Scale Retinex with Color Restoration) is based on MSR. It applies MSR to each spectral channel (e.g. R, G and B), and modify the MSR output by multiplying it by a color restoration function of the chromaticity.
+
+When MSR is applied to each spectral channel, it assumes the image obey gray world assumption. Otherwise, if the image violates gray world assumption, the MSR will produce grayish image by decreasing the color saturation, thus the color restoration step is proposed to solve this problem. However, for images with nice color balance, MSRCR still produces a desaturated look. Hence it is recommended to use MSRCP in most cases, and only apply MSRCR to the images with color cast. Also since MSRCR applies MSR to each spectral channel instead of intensity channel, it is slower than MSRCP.
+
+This function only accept 8-16bit integer RGB input.
+
+### Usage
+
+```python
+retinex.MSRCR(clip input, float[] sigmaS=[25,80,250], float lower_thr=0.001, float upper_thr=0.001, bool fulls=True, bool fulld=fulls, float restore=125)
+```
+
+- input:<br />
+    clip to process
+
+- sigma: (Default: [25,80,250])<br />
+    The same as MSRCP.
+
+- lower_thr: (Default: 0.001)<br />
+    The same as MSRCP.
+
+- upper_thr: (Default: 0.001)<br />
+    The same as MSRCP.
+
+- fulls: (Default: True)<br />
+    The same as MSRCP.
+
+- fulld: (Default: fulls)<br />
+    The same as MSRCP.
+
+- restore: (Default: 125)<br />
+    The strength of the nonlinearity for color restoration function, larger value result in stronger restoration, available range is [0, +inf).<br />
+    It is a multiplier in a log function, so try to adjust it in a large scale (e.g. multiply it by a power of 10) if you want to see any difference.
+
+### Example
+
+TV range YUV420P8 input, filtered in PC range RGB48, output PC range RGB48
+
+```python
+v = core.fmtc.resample(v, csp=vs.YUV444P16)
+v = core.fmtc.matrix(v, mat="709", csp=vs.RGB48)
+v = core.retinex.MSRCR(v)
+```
+
+JPEG image(PC range YUV420P8 with MPEG-1 chroma placement) input, filtered in PC range RGB48 without color restoration (pure MSR), output PC range RGB48
+
+```python
+i = core.lsmas.LWLibavSource(r'Image.jpg')
+i = core.fmtc.resample(i, csp=vs.YUV444P16, fulls=True, cplace="MPEG1")
+i = core.fmtc.matrix(i, mat="601", fulls=True, csp=vs.RGB48)
+i = core.retinex.MSRCR(i, restore=0)
+```
+
+PNG image(PC range RGB24) input, filtered in PC range RGB48, output PC range RGB48
+
+```python
+i = core.lsmas.LWLibavSource(r'Image.png')
+i = core.fmtc.bitdepth(i, bits=16)
+i = core.retinex.MSRCR(i)
+```
diff --git a/include/Helper.h b/include/Helper.h
@@ -192,22 +192,23 @@ class VSProcess
     const VSFormat *fi = nullptr;
     VSFrameRef *dst = nullptr;
 
+    int PlaneCount;
     int Bps;
     int bps;
 
-    int stride;
-    int width;
     int height;
+    int width;
+    int stride;
     int pcount;
 
-    int src_stride[VSMaxPlaneCount];
-    int src_width[VSMaxPlaneCount];
     int src_height[VSMaxPlaneCount];
+    int src_width[VSMaxPlaneCount];
+    int src_stride[VSMaxPlaneCount];
     int src_pcount[VSMaxPlaneCount];
 
-    int dst_stride[VSMaxPlaneCount];
-    int dst_width[VSMaxPlaneCount];
     int dst_height[VSMaxPlaneCount];
+    int dst_width[VSMaxPlaneCount];
+    int dst_stride[VSMaxPlaneCount];
     int dst_pcount[VSMaxPlaneCount];
 
 private:
@@ -225,32 +226,37 @@ class VSProcess
         src = vsapi->getFrameFilter(n, d.node, frameCtx);
         fi = vsapi->getFrameFormat(src);
 
+        PlaneCount = fi->numPlanes;
         Bps = fi->bytesPerSample;
         bps = fi->bitsPerSample;
 
-        stride = vsapi->getStride(src, 0) / Bps;
-        width = vsapi->getFrameWidth(src, 0);
         height = vsapi->getFrameHeight(src, 0);
+        width = vsapi->getFrameWidth(src, 0);
+        stride = vsapi->getStride(src, 0) / Bps;
         pcount = stride * height;
 
-        const int planes[VSMaxPlaneCount] = { 0, 1, 2 };
-        const VSFrameRef * cp_planes[VSMaxPlaneCount] = { d.process[0] ? nullptr : src, d.process[1] ? nullptr : src, d.process[2] ? nullptr : src };
-        dst = vsapi->newVideoFrame2(fi, width, height, cp_planes, planes, src, core);
+        int planes[VSMaxPlaneCount];
+        const VSFrameRef *cp_planes[VSMaxPlaneCount];
 
         for (int i = 0; i < VSMaxPlaneCount; i++)
         {
-            if (d.process[i])
-            {
-                src_stride[i] = vsapi->getStride(src, i) / Bps;
-                src_width[i] = vsapi->getFrameWidth(src, i);
-                src_height[i] = vsapi->getFrameHeight(src, i);
-                src_pcount[i] = src_stride[i] * src_height[i];
-
-                dst_stride[i] = vsapi->getStride(dst, i) / Bps;
-                dst_width[i] = vsapi->getFrameWidth(dst, i);
-                dst_height[i] = vsapi->getFrameHeight(dst, i);
-                dst_pcount[i] = dst_stride[i] * dst_height[i];
-            }
+            planes[i] = i;
+            cp_planes[i] = d.process[i] ? nullptr : src;
+        }
+
+        dst = vsapi->newVideoFrame2(fi, width, height, cp_planes, planes, src, core);
+
+        for (int i = 0; i < PlaneCount; i++)
+        {
+            src_height[i] = vsapi->getFrameHeight(src, i);
+            src_width[i] = vsapi->getFrameWidth(src, i);
+            src_stride[i] = vsapi->getStride(src, i) / Bps;
+            src_pcount[i] = src_stride[i] * src_height[i];
+
+            dst_height[i] = vsapi->getFrameHeight(dst, i);
+            dst_width[i] = vsapi->getFrameWidth(dst, i);
+            dst_stride[i] = vsapi->getStride(dst, i) / Bps;
+            dst_pcount[i] = dst_stride[i] * dst_height[i];
         }
     }
 
@@ -263,11 +269,11 @@ class VSProcess
     {
         int i;
 
-        for (i = 0; i < VSMaxPlaneCount; i++)
+        for (i = 0; i < PlaneCount; i++)
         {
             if (d.process[i]) break;
         }
-        if (i >= VSMaxPlaneCount) return dst;
+        if (i >= PlaneCount) return dst;
 
         else if (Bps == 1)
         {

diff --git a/include/MSR.h b/include/MSR.h
@@ -188,12 +188,123 @@ class MSRProcess
 
     virtual ~MSRProcess() {}
 
-    int SimplestColorBalance(FLType *odata, const FLType *idata) const;
+    // Multi Scale Retinex process kernel for floating point data
     int MSRKernel(FLType *odata, const FLType *idata) const;
+
+    // Simplest color balance with pixel clipping on either side of the dynamic range
+    int SimplestColorBalance(FLType *odata, const FLType *idata) const; // odata as input and output, idata as source
+    template < typename T >
+    int SimplestColorBalance(T *dst, FLType *odata, const T *src, T dFloor, T dCeil) const; // odata as input, dst as output, src as source
 };
 
 
 ////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
 
 
+template < typename T >
+int MSRProcess::SimplestColorBalance(T *dst, FLType *odata, const T *src, T dFloor, T dCeil) const
+{
+    int i, j, upper;
+
+    FLType offset, gain;
+    FLType min = FLType_MAX;
+    FLType max = -FLType_MAX;
+
+    FLType dFloorFL = static_cast<FLType>(dFloor);
+    FLType dCeilFL = static_cast<FLType>(dCeil);
+    FLType dRangeFL = dCeilFL - dFloorFL;
+
+    for (j = 0; j < height; j++)
+    {
+        i = stride * j;
+        for (upper = i + width; i < upper; i++)
+        {
+            min = Min(min, odata[i]);
+            max = Max(max, odata[i]);
+        }
+    }
+
+    if (max <= min)
+    {
+        memcpy(dst, src, sizeof(T)*pcount);
+        return 1;
+    }
+
+    if (d.lower_thr > 0 || d.upper_thr > 0)
+    {
+        int h, HistBins = d.HistBins;
+        int Count, MaxCount;
+
+        int *Histogram = vs_aligned_malloc<int>(sizeof(int)*HistBins, Alignment);
+        memset(Histogram, 0, sizeof(int)*HistBins);
+
+        gain = (HistBins - 1) / (max - min);
+        offset = -min * gain;
+
+        for (j = 0; j < height; j++)
+        {
+            i = stride * j;
+            for (upper = i + width; i < upper; i++)
+            {
+                Histogram[static_cast<int>(odata[i] * gain + offset)]++;
+            }
+        }
+
+        gain = (max - min) / (HistBins - 1);
+        offset = min;
+
+        Count = 0;
+        MaxCount = static_cast<int>(width*height*d.lower_thr + 0.5);
+
+        for (h = 0; h < HistBins; h++)
+        {
+            Count += Histogram[h];
+            if (Count > MaxCount) break;
+        }
+
+        min = h * gain + offset;
+
+        Count = 0;
+        MaxCount = static_cast<int>(width*height*d.upper_thr + 0.5);
+
+        for (h = HistBins - 1; h >= 0; h--)
+        {
+            Count += Histogram[h];
+            if (Count > MaxCount) break;
+        }
+
+        max = h * gain + offset;
+
+        vs_aligned_free(Histogram);
+    }
+
+    gain = dRangeFL / (max - min);
+    offset = -min * gain + dFloorFL + FLType(0.5);
+
+    if (d.lower_thr > 0 || d.upper_thr > 0)
+    {
+        for (j = 0; j < height; j++)
+        {
+            i = stride * j;
+            for (upper = i + width; i < upper; i++)
+                dst[i] = static_cast<T>(Clip(odata[i] * gain + offset, dFloorFL, dCeilFL));
+        }
+    }
+    else
+    {
+        for (j = 0; j < height; j++)
+        {
+            i = stride * j;
+            for (upper = i + width; i < upper; i++)
+                dst[i] = static_cast<T>(odata[i] * gain + offset);
+        }
+    }
+
+    return 0;
+}
+
+
+////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
+
+
 #endif