1 /***********************************************************
2 * A new denoiser for the mjpegtools project *
3 * ------------------------------------------------------- *
4 * (C) 2004-2009 Steven Boswell. *
5 * Based on yuvdenoise/main.c, (C) 2001,2002 Stefan Fendt *
6 * *
7 * Licensed and protected by the GNU-General-Public- *
8 * License version 2. See the file COPYING for detailed *
9 * information. *
10 ***********************************************************/
11
12 #ifdef HAVE_CONFIG_H
13 #include "config.h"
14 #endif
15
16 #include <stdio.h>
17 #include <stdlib.h>
18 #include <string.h>
19 #include <unistd.h>
20 #include <assert.h>
21 #include "mjpeg_types.h"
22 #include "yuv4mpeg.h"
23 #include "mjpeg_logging.h"
24 #include "config.h"
25 #include "newdenoise.hh"
26
27 // HACK
28 #include <sys/stat.h>
29 #include <fcntl.h>
30
31 void process_commandline(int argc, char *argv[]);
32 void alloc_buffers (void);
33 void display_help (void);
34
35 DNSR_GLOBAL denoiser;
36 int frame = 0;
37 int verbose = 1;
38
39 /***********************************************************
40 * *
41 ***********************************************************/
42
main(int argc,char * argv[])43 int main(int argc, char *argv[])
44 {
45 int fd_in = 0;
46 int fd_out = 1;
47 int errno = 0;
48
49 y4m_frame_info_t frameinfo;
50 y4m_stream_info_t streaminfo;
51 int chroma_mode;
52 int bInputStreamEnded;
53
54 frame = 0;
55
56 /* initialize stream-information */
57 y4m_accept_extensions (1);
58 y4m_init_stream_info (&streaminfo);
59 y4m_init_frame_info (&frameinfo);
60
61 /* setup default values for denoiser's parameters */
62 denoiser.frames = 10;
63 denoiser.interlaced = -1;
64 denoiser.bwonly = 0;
65 denoiser.radiusY = 16;
66 denoiser.radiusCbCr = -1;
67 denoiser.zThresholdY = 0; /* assume medium noise material */
68 denoiser.zThresholdCbCr = -1;
69 denoiser.thresholdY = 3; /* assume medium noise material */
70 denoiser.thresholdCbCr = -1;
71 denoiser.matchCountThrottle = 16;
72 denoiser.matchSizeThrottle = 256;
73 denoiser.threads = 1;
74
75 /* process commandline */
76 process_commandline(argc, argv);
77
78 /* If color-specific settings were not specified, copy them
79 from the intensity settings. */
80 if (denoiser.radiusCbCr == -1)
81 denoiser.radiusCbCr = denoiser.radiusY;
82 if (denoiser.thresholdCbCr == -1)
83 denoiser.thresholdCbCr = denoiser.thresholdY;
84 if (denoiser.zThresholdCbCr == -1)
85 denoiser.zThresholdCbCr = denoiser.zThresholdY;
86
87 /* Make sure the parameters are consistent. */
88 {
89 int isConsistent = 1;
90
91 /* The zero-motion error threshold has to be lower than the
92 motion-compensated error threshold. The zero-motion
93 pass is just a way to speed up detection of unambiguous
94 zero-motion areas, so a higher error threshold would defeat
95 the purpose. */
96 if (denoiser.zThresholdY > denoiser.thresholdY)
97 {
98 mjpeg_error ("-z setting (%d) cannot be larger than "
99 "-t setting (%d)\n", denoiser.zThresholdY,
100 denoiser.thresholdY);
101 isConsistent = 0;
102 }
103 if (denoiser.zThresholdCbCr > denoiser.thresholdCbCr)
104 {
105 mjpeg_error ("-Z setting (%d) cannot be larger than "
106 "-T setting (%d)\n", denoiser.zThresholdCbCr,
107 denoiser.thresholdCbCr);
108 isConsistent = 0;
109 }
110
111 /* If the parameters weren't consistent, quit with an error. */
112 if (!isConsistent)
113 exit (1);
114 }
115
116 /* open input stream */
117 if ((errno = y4m_read_stream_header (fd_in, &streaminfo)) != Y4M_OK)
118 mjpeg_error_exit1 ("Couldn't read YUV4MPEG header: %s!",
119 y4m_strerr (errno));
120 if (y4m_si_get_plane_count (&streaminfo) != 3)
121 mjpeg_error_exit1 ("Only 3-plane formats supported.");
122 denoiser.frame.w = y4m_si_get_width (&streaminfo);
123 denoiser.frame.h = y4m_si_get_height (&streaminfo);
124 chroma_mode = y4m_si_get_chroma (&streaminfo);
125 denoiser.frame.ss_h = y4m_chroma_ss_x_ratio (chroma_mode).d;
126 denoiser.frame.ss_v = y4m_chroma_ss_y_ratio (chroma_mode).d;
127 denoiser.frame.Cw=denoiser.frame.w/denoiser.frame.ss_h;
128 denoiser.frame.Ch=denoiser.frame.h/denoiser.frame.ss_v;
129
130 /* write the outstream header */
131 y4m_write_stream_header (fd_out, &streaminfo);
132
133 /* allocate memory for frames */
134 if (!(denoiser.threads & 1))
135 alloc_buffers();
136
137 /* set interlacing, if it hasn't been yet. */
138 if (denoiser.interlaced == -1)
139 {
140 int n = y4m_si_get_interlace(&streaminfo);
141 switch (n)
142 {
143 case Y4M_ILACE_NONE:
144 denoiser.interlaced = 0;
145 break;
146 case Y4M_ILACE_TOP_FIRST:
147 denoiser.interlaced = 1;
148 break;
149 case Y4M_ILACE_BOTTOM_FIRST:
150 denoiser.interlaced = 2;
151 break;
152 default:
153 mjpeg_warn("Unknown interlacing '%d', assuming non-interlaced", n);
154 denoiser.interlaced = 0;
155 break;
156 }
157 }
158
159 /* if we're processing interlaced material, there is an additional
160 requirement that the number of reference frames be a multiple of
161 2 (since there are 2 fields per frame). */
162 if (denoiser.interlaced != 0 && (denoiser.frames & 1) != 0)
163 mjpeg_error_exit1 ("When denoising interlaced material, -f must be"
164 " a multiple of 2");
165
166 /* initialize the denoiser */
167 errno = newdenoise_init (denoiser.frames, denoiser.frame.w,
168 denoiser.frame.h, (denoiser.bwonly) ? 0 : denoiser.frame.Cw,
169 (denoiser.bwonly) ? 0 : denoiser.frame.Ch, fd_in, fd_out,
170 &streaminfo, &frameinfo);
171 if (errno == -1)
172 mjpeg_error_exit1( "Could not initialize denoiser");
173
174 /* get space for the first output frame, if necessary */
175 if (denoiser.threads & 1)
176 {
177 errno = newdenoise_get_write_frame (denoiser.frame.out);
178 if (errno)
179 mjpeg_error_exit1 ("Could not get space for frame %d",
180 frame);
181 }
182
183 /* read every single frame until the end of the input stream */
184 bInputStreamEnded = 0;
185 for (;;)
186 {
187 /* If there is another input frame, read it. We may find there
188 is no more input. */
189 if (!bInputStreamEnded)
190 {
191 /* Read the next frame. */
192 if (!(denoiser.threads & 1))
193 errno = y4m_read_frame (fd_in, &streaminfo, &frameinfo,
194 denoiser.frame.in);
195 else
196 errno = newdenoise_read_frame (denoiser.frame.in);
197
198 /* did stream end unexpectedly ? */
199 if(errno != Y4M_ERR_EOF && errno != Y4M_OK )
200 mjpeg_error_exit1( "%s", y4m_strerr( errno ) );
201
202 /* Note if the stream ended. */
203 if (errno == Y4M_ERR_EOF)
204 bInputStreamEnded = 1;
205 }
206
207 /* One more input frame. */
208 frame++;
209
210 //if (frame < 5395) continue; // MAJOR HACK MAJOR HACK MAJOR HACK
211 //if (frame == 5455) break; // MAJOR HACK MAJOR HACK MAJOR HACK
212 // fprintf (stderr, "Frame %d\r", frame); // HACK
213
214 /* denoise the current frame */
215 errno = ((denoiser.interlaced == 0) ? newdenoise_frame
216 : newdenoise_interlaced_frame)
217 (((bInputStreamEnded) ? NULL : denoiser.frame.in[0]),
218 ((bInputStreamEnded) ? NULL : denoiser.frame.in[1]),
219 ((bInputStreamEnded) ? NULL : denoiser.frame.in[2]),
220 denoiser.frame.out[0], denoiser.frame.out[1],
221 denoiser.frame.out[2]);
222 if (errno == -1)
223 mjpeg_error_exit1( "Could not denoise frame %d", frame);
224
225 /* if there was no output from the denoiser, leave. */
226 if (bInputStreamEnded && errno == 1)
227 break;
228
229 /* if there was some output from the denoiser, write it. */
230 if (errno == 0)
231 {
232 /* if b/w was selected, set the frame color to white */
233 if (denoiser.bwonly)
234 {
235 int i, iExtent = denoiser.frame.Cw * denoiser.frame.Ch;
236
237 for (i = 0; i < iExtent; ++i)
238 denoiser.frame.out[1][i] = 128;
239 for (i = 0; i < iExtent; ++i)
240 denoiser.frame.out[2][i] = 128;
241 }
242
243 /* Write the frame. */
244 if (!(denoiser.threads & 1))
245 errno = y4m_write_frame ( fd_out, &streaminfo,
246 &frameinfo, denoiser.frame.out);
247 else
248 errno = newdenoise_write_frame();
249 if (errno)
250 mjpeg_error_exit1 ("Could not write frame %d", frame);
251
252 /* Get space for the next frame.*/
253 if (denoiser.threads & 1)
254 {
255 errno = newdenoise_get_write_frame (denoiser.frame.out);
256 if (errno)
257 mjpeg_error_exit1 ("Could not get space for "
258 "frame %d", frame);
259 }
260 }
261 }
262
263 /* shut down the denoiser */
264 errno = newdenoise_shutdown();
265 if (errno != 0)
266 mjpeg_error_exit1 ("Could not shut down denoiser");
267
268 /* Exit gently */
269 return(0);
270 }
271
272
273
274 void
process_commandline(int argc,char * argv[])275 process_commandline(int argc, char *argv[])
276 {
277 char c;
278
279 while ((c = getopt (argc, argv, "h?z:Z:t:T:r:R:m:M:f:BI:p:v:")) != -1)
280 {
281 switch (c)
282 {
283 case 'h':
284 {
285 display_help();
286 exit (0);
287 break;
288 }
289 case '?':
290 {
291 display_help();
292 exit (0);
293 break;
294 }
295 case 'r':
296 {
297 denoiser.radiusY = atoi(optarg);
298 if(denoiser.radiusY<4)
299 {
300 denoiser.radiusY=4;
301 mjpeg_warn ("Minimum allowed search radius is 4 pixels.");
302 }
303 break;
304 }
305 case 'R':
306 {
307 denoiser.radiusCbCr = atoi(optarg);
308 if(denoiser.radiusCbCr<4)
309 {
310 denoiser.radiusCbCr=4;
311 mjpeg_warn ("Minimum allowed color search radius is 4 pixel.");
312 }
313 break;
314 }
315 case 'z':
316 {
317 denoiser.zThresholdY = atoi(optarg);
318 break;
319 }
320 case 'Z':
321 {
322 denoiser.zThresholdCbCr = atoi(optarg);
323 break;
324 }
325 case 't':
326 {
327 denoiser.thresholdY = atoi(optarg);
328 break;
329 }
330 case 'T':
331 {
332 denoiser.thresholdCbCr = atoi(optarg);
333 break;
334 }
335 case 'm':
336 {
337 denoiser.matchCountThrottle = atoi(optarg);
338 break;
339 }
340 case 'M':
341 {
342 denoiser.matchSizeThrottle = atoi(optarg);
343 break;
344 }
345 case 'f':
346 {
347 denoiser.frames = atoi(optarg);
348 break;
349 }
350 case 'B':
351 {
352 denoiser.bwonly = 1;
353 break;
354 }
355 case 'I':
356 {
357 int interlaced = atoi (optarg);
358 if (interlaced != 0 && interlaced != 1 && interlaced != 2)
359 {
360 mjpeg_error_exit1 ("-I must be either 0, 1, or 2");
361 }
362 denoiser.interlaced = interlaced;
363 break;
364 }
365 case 'p':
366 {
367 int threads = atoi (optarg);
368 if (threads < 0 || threads > 3)
369 {
370 mjpeg_error_exit1 ("-p must be either 0, 1, 2, or 3");
371 }
372 denoiser.threads = threads;
373 break;
374 }
375 case 'v':
376 verbose = atoi (optarg);
377 if (verbose < 0 || verbose > 2)
378 {
379 mjpeg_error_exit1 ("Verbose level must be [0..2]");
380 }
381 break;
382 default:
383 fprintf (stderr, "Unknown option '%c'\n", c);
384 display_help();
385 exit (1);
386 }
387 }
388 }
389
390
alloc_buffers(void)391 void alloc_buffers(void)
392 {
393 int luma_buffsize;
394 int chroma_buffsize;
395
396 luma_buffsize = denoiser.frame.w * denoiser.frame.h;
397 chroma_buffsize = denoiser.frame.Cw * denoiser.frame.Ch;
398
399 denoiser.frame.in[0] = (uint8_t *) malloc (luma_buffsize);
400 denoiser.frame.in[1] = (uint8_t *) malloc (chroma_buffsize);
401 denoiser.frame.in[2] = (uint8_t *) malloc (chroma_buffsize);
402
403 denoiser.frame.out[0] = (uint8_t *) malloc (luma_buffsize);
404 denoiser.frame.out[1] = (uint8_t *) malloc (chroma_buffsize);
405 denoiser.frame.out[2] = (uint8_t *) malloc (chroma_buffsize);
406
407 if ( denoiser.frame.in[0] == NULL
408 || denoiser.frame.in[1] == NULL
409 || denoiser.frame.in[2] == NULL
410 || denoiser.frame.out[0] == NULL
411 || denoiser.frame.out[1] == NULL
412 || denoiser.frame.out[2] == NULL)
413 mjpeg_error_exit1( "Out of memory "
414 "when allocating frame buffers");
415 }
416
417 void
display_help(void)418 display_help (void)
419 {
420 fprintf (stderr,
421 "y4mdenoise options\n"
422 "------------------\n"
423 "-p parallelism: 0=no threads, 1=r/w thread only, 2=do color in\n"
424 " separate thread (default: 1)\n"
425 "-r Radius for motion-search (default: 16)\n"
426 "-R Radius for color motion-search (default: -r setting)\n"
427 "-t Error tolerance (default: 3)\n"
428 "-T Color error tolerance (default: -t setting)\n"
429 "-z Error tolerance for zero-motion pass (default: 2)\n"
430 "-Z Error tolerance for color's zero-motion pass (default: -z setting)\n"
431 "-m Match-count throttle (keep this many of the best pixel-group\n"
432 " matches found in a radius search) (default: 15)\n"
433 "-M Match-size throttle (apply first match whose flood-fill is the\n"
434 " size of this many pixel-groups or greater) (default: 3)\n"
435 "-f Number of reference frames (default: 10)\n"
436 "-B Black-and-white mode; denoise intensity, set color to white\n"
437 "-I Interlacing type: 0=frame, 1=top-field-first, 2=bottom-field-first\n"
438 " (default: taken from stream header)\n"
439 "-v Verbosity (0=none, 1=normal 2=debug)\n"
440 "-h,-? Help\n"
441 "Intensity pixel-groups are 4x2 (i.e. 4 across and 2 down).\n"
442 "Color pixel-groups are 2x2.\n"
443 );
444 }
445