Core matte input for keying node
[blender.git] / source / blender / compositor / nodes / COM_KeyingNode.cpp
1 /*
2  * Copyright 2012, Blender Foundation.
3  *
4  * This program is free software; you can redistribute it and/or
5  * modify it under the terms of the GNU General Public License
6  * as published by the Free Software Foundation; either version 2
7  * of the License, or (at your option) any later version.
8  *
9  * This program is distributed in the hope that it will be useful,
10  * but WITHOUT ANY WARRANTY; without even the implied warranty of
11  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
12  * GNU General Public License for more details.
13  *
14  * You should have received a copy of the GNU General Public License
15  * along with this program; if not, write to the Free Software Foundation,
16  * Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
17  *
18  * Contributor:
19  *              Jeroen Bakker
20  *              Monique Dewanchand
21  *              Sergey Sharybin
22  */
23
24 #include "COM_KeyingNode.h"
25
26 #include "COM_ExecutionSystem.h"
27
28 #include "COM_KeyingOperation.h"
29 #include "COM_KeyingBlurOperation.h"
30 #include "COM_KeyingDespillOperation.h"
31 #include "COM_KeyingClipOperation.h"
32
33 #include "COM_SeparateChannelOperation.h"
34 #include "COM_CombineChannelsOperation.h"
35 #include "COM_ConvertRGBToYCCOperation.h"
36 #include "COM_ConvertYCCToRGBOperation.h"
37 #include "COM_GaussianBokehBlurOperation.h"
38 #include "COM_SetValueOperation.h"
39
40 #include "COM_DilateErodeOperation.h"
41
42 #include "COM_SetAlphaOperation.h"
43
44 KeyingNode::KeyingNode(bNode *editorNode): Node(editorNode)
45 {
46 }
47
48 OutputSocket *KeyingNode::setupPreBlur(ExecutionSystem *graph, InputSocket *inputImage, int size, OutputSocket **originalImage)
49 {
50         ConvertRGBToYCCOperation *convertRGBToYCCOperation = new ConvertRGBToYCCOperation();
51         convertRGBToYCCOperation->setMode(0);  /* ITU 601 */
52
53         inputImage->relinkConnections(convertRGBToYCCOperation->getInputSocket(0), 0, graph);
54         graph->addOperation(convertRGBToYCCOperation);
55
56         CombineChannelsOperation *combineOperation = new CombineChannelsOperation();
57         graph->addOperation(combineOperation);
58
59         for (int channel = 0; channel < 4; channel++) {
60                 SeparateChannelOperation *separateOperation = new SeparateChannelOperation();
61                 separateOperation->setChannel(channel);
62                 addLink(graph, convertRGBToYCCOperation->getOutputSocket(0), separateOperation->getInputSocket(0));
63                 graph->addOperation(separateOperation);
64
65                 if (channel == 0 || channel == 3) {
66                         addLink(graph, separateOperation->getOutputSocket(0), combineOperation->getInputSocket(channel));
67                 }
68                 else {
69                         KeyingBlurOperation *blurOperation = new KeyingBlurOperation();
70
71                         blurOperation->setSize(size);
72
73                         addLink(graph, separateOperation->getOutputSocket(0), blurOperation->getInputSocket(0));
74                         addLink(graph, blurOperation->getOutputSocket(0), combineOperation->getInputSocket(channel));
75                         graph->addOperation(blurOperation);
76                 }
77         }
78
79         ConvertYCCToRGBOperation *convertYCCToRGBOperation = new ConvertYCCToRGBOperation();
80         convertYCCToRGBOperation->setMode(0);  /* ITU 601 */
81         addLink(graph, combineOperation->getOutputSocket(0), convertYCCToRGBOperation->getInputSocket(0));
82         graph->addOperation(convertYCCToRGBOperation);
83
84         *originalImage = convertRGBToYCCOperation->getInputSocket(0)->getConnection()->getFromSocket();
85
86         return convertYCCToRGBOperation->getOutputSocket(0);
87 }
88
89 OutputSocket *KeyingNode::setupPostBlur(ExecutionSystem *graph, OutputSocket *postBLurInput, int size)
90 {
91         KeyingBlurOperation *blurOperation = new KeyingBlurOperation();
92
93         blurOperation->setSize(size);
94
95         addLink(graph, postBLurInput, blurOperation->getInputSocket(0));
96
97         graph->addOperation(blurOperation);
98
99         return blurOperation->getOutputSocket();
100 }
101
102 OutputSocket *KeyingNode::setupDilateErode(ExecutionSystem *graph, OutputSocket *dilateErodeInput, int distance)
103 {
104         DilateStepOperation *dilateErodeOperation;
105
106         if (distance > 0) {
107                 dilateErodeOperation = new DilateStepOperation();
108                 dilateErodeOperation->setIterations(distance);
109         }
110         else {
111                 dilateErodeOperation = new ErodeStepOperation();
112                 dilateErodeOperation->setIterations(-distance);
113         }
114
115         addLink(graph, dilateErodeInput, dilateErodeOperation->getInputSocket(0));
116
117         graph->addOperation(dilateErodeOperation);
118
119         return dilateErodeOperation->getOutputSocket(0);
120 }
121
122 OutputSocket *KeyingNode::setupDespill(ExecutionSystem *graph, OutputSocket *despillInput, OutputSocket *inputScreen, float factor)
123 {
124         KeyingDespillOperation *despillOperation = new KeyingDespillOperation();
125
126         despillOperation->setDespillFactor(factor);
127
128         addLink(graph, despillInput, despillOperation->getInputSocket(0));
129         addLink(graph, inputScreen, despillOperation->getInputSocket(1));
130
131         graph->addOperation(despillOperation);
132
133         return despillOperation->getOutputSocket(0);
134 }
135
136 OutputSocket *KeyingNode::setupClip(ExecutionSystem *graph, OutputSocket *clipInput, int kernelRadius, float kernelTolerance,
137                                     float clipBlack, float clipWhite, bool edgeMatte)
138 {
139         KeyingClipOperation *clipOperation = new KeyingClipOperation();
140
141         clipOperation->setKernelRadius(kernelRadius);
142         clipOperation->setKernelTolerance(kernelTolerance);
143
144         clipOperation->setClipBlack(clipBlack);
145         clipOperation->setClipWhite(clipWhite);
146         clipOperation->setIsEdgeMatte(edgeMatte);
147
148         addLink(graph, clipInput, clipOperation->getInputSocket(0));
149
150         graph->addOperation(clipOperation);
151
152         return clipOperation->getOutputSocket(0);
153 }
154
155 void KeyingNode::convertToOperations(ExecutionSystem *graph, CompositorContext *context)
156 {
157         InputSocket *inputImage = this->getInputSocket(0);
158         InputSocket *inputScreen = this->getInputSocket(1);
159         InputSocket *inputGarbageMatte = this->getInputSocket(2);
160         InputSocket *inputCoreMatte = this->getInputSocket(3);
161         OutputSocket *outputImage = this->getOutputSocket(0);
162         OutputSocket *outputMatte = this->getOutputSocket(1);
163         OutputSocket *outputEdges = this->getOutputSocket(2);
164         OutputSocket *postprocessedMatte, *postprocessedImage, *originalImage, *edgesMatte;
165
166         bNode *editorNode = this->getbNode();
167         NodeKeyingData *keying_data = (NodeKeyingData *) editorNode->storage;
168
169         /* keying operation */
170         KeyingOperation *keyingOperation = new KeyingOperation();
171
172         keyingOperation->setScreenBalance(keying_data->screen_balance);
173
174         inputScreen->relinkConnections(keyingOperation->getInputSocket(1), 1, graph);
175         inputGarbageMatte->relinkConnections(keyingOperation->getInputSocket(2), 2, graph);
176         inputCoreMatte->relinkConnections(keyingOperation->getInputSocket(3), 3, graph);
177
178         if (keying_data->blur_pre) {
179                 /* chroma preblur operation for input of keying operation  */
180                 OutputSocket *preBluredImage = setupPreBlur(graph, inputImage, keying_data->blur_pre, &originalImage);
181                 addLink(graph, preBluredImage, keyingOperation->getInputSocket(0));
182         }
183         else {
184                 inputImage->relinkConnections(keyingOperation->getInputSocket(0), 0, graph);
185                 originalImage = keyingOperation->getInputSocket(0)->getConnection()->getFromSocket();
186         }
187
188         graph->addOperation(keyingOperation);
189
190         postprocessedMatte = keyingOperation->getOutputSocket();
191
192         if (keying_data->clip_black > 0.0f || keying_data->clip_white< 1.0f) {
193                 postprocessedMatte = setupClip(graph, postprocessedMatte,
194                                                keying_data->edge_kernel_radius, keying_data->edge_kernel_tolerance,
195                                                keying_data->clip_black, keying_data->clip_white, false);
196         }
197
198         edgesMatte = setupClip(graph, postprocessedMatte,
199                                keying_data->edge_kernel_radius, keying_data->edge_kernel_tolerance,
200                                keying_data->clip_black, keying_data->clip_white, true);
201
202         /* apply blur on matte if needed */
203         if (keying_data->blur_post)
204                 postprocessedMatte = setupPostBlur(graph, postprocessedMatte, keying_data->blur_post);
205
206         /* matte dilate/erode */
207         if (keying_data->dilate_distance != 0) {
208                 postprocessedMatte = setupDilateErode(graph, postprocessedMatte, keying_data->dilate_distance);
209         }
210
211         /* set alpha channel to output image */
212         SetAlphaOperation *alphaOperation = new SetAlphaOperation();
213         addLink(graph, originalImage, alphaOperation->getInputSocket(0));
214         addLink(graph, postprocessedMatte, alphaOperation->getInputSocket(1));
215
216         postprocessedImage = alphaOperation->getOutputSocket();
217
218         /* despill output image */
219         if (keying_data->despill_factor > 0.0f) {
220                 postprocessedImage = setupDespill(graph, postprocessedImage,
221                                                   keyingOperation->getInputSocket(1)->getConnection()->getFromSocket(),
222                                                   keying_data->despill_factor);
223         }
224
225         /* connect result to output sockets */
226         outputImage->relinkConnections(postprocessedImage);
227         outputMatte->relinkConnections(postprocessedMatte);
228         outputEdges->relinkConnections(edgesMatte);
229
230         graph->addOperation(alphaOperation);
231 }