newdenoise.cc

Motion JPEG编解码器源代码

// This file (C) 2004 Steven Boswell.  All rights reserved.
// Released to the public under the GNU General Public License.
// See the file COPYING for more information.

#ifdef HAVE_CONFIG_H
#include "config.h"
#endif

#include <assert.h>
#include <unistd.h>
#include <string.h>
#include <math.h>
#include <pthread.h>
#include "mjpeg_types.h"
#include "mjpeg_logging.h"
#include "yuv4mpeg.h"
#include <stdio.h>
#include "newdenoise.hh"
#include "MotionSearcher.hh"

// The denoisers.  (We have to make these classes, in order to keep gdb
// from crashing.  I didn't even know one could crash gdb. ;-)
#if 0
typedef MotionSearcher<uint8_t, 1, int32_t, int16_t, int32_t, 4, 2,
		uint16_t, PixelY, ReferencePixelY, ReferenceFrameY>
	MotionSearcherY;
typedef MotionSearcher<uint8_t, 2, int32_t, int16_t, int32_t, 2, 2,
		uint16_t, PixelCbCr, ReferencePixelCbCr, ReferenceFrameCbCr>
	MotionSearcherCbCr;
#else
class MotionSearcherY
	: public MotionSearcher<uint8_t, 1, int32_t, int16_t, int32_t, 4, 2,
		uint16_t, PixelY, ReferencePixelY, ReferenceFrameY> {};
class MotionSearcherCbCr
	: public MotionSearcher<uint8_t, 2, int32_t, int16_t, int32_t, 2, 2,
		uint16_t, PixelCbCr, ReferencePixelCbCr, ReferenceFrameCbCr> {};
#endif
MotionSearcherY g_oMotionSearcherY;
MotionSearcherCbCr g_oMotionSearcherCbCr;

// Whether the denoisers should be used.
bool g_bMotionSearcherY;
bool g_bMotionSearcherCbCr;

// Pixel buffers, used to translate provided input into the form the
// denoiser needs.
int g_nPixelsY, g_nWidthY, g_nHeightY;
MotionSearcherY::Pixel_t *g_pPixelsY;
int g_nPixelsCbCr, g_nWidthCbCr, g_nHeightCbCr;
MotionSearcherCbCr::Pixel_t *g_pPixelsCbCr;

// Internal methods to output a frame.
static void output_frame
	(const MotionSearcherY::ReferenceFrame_t *a_pFrameY,
	const MotionSearcherCbCr::ReferenceFrame_t *a_pFrameCbCr,
	uint8_t *a_pOutputY, uint8_t *a_pOutputCb, uint8_t *a_pOutputCr);
static void output_field
	(int a_nMask, const MotionSearcherY::ReferenceFrame_t *a_pFrameY,
	const MotionSearcherCbCr::ReferenceFrame_t *a_pFrameCbCr,
	uint8_t *a_pOutputY, uint8_t *a_pOutputCb, uint8_t *a_pOutputCr);



// A class wrapper for mutexes.
class ThreadMutex
{
public:
	ThreadMutex();
		// Default constructor.
	
	~ThreadMutex();
		// Destructor.
	
	void Lock (void);
		// Lock the mutex.
	
	void Unlock (void);
		// Unlock the mutex.
	
	operator pthread_mutex_t * (void) { return &m_oMutex; }
		// Accessor.

private:
	pthread_mutex_t m_oMutex;
		// The wrapped mutex.

#ifndef NDEBUG
public:
	bool m_bLocked;
		// true if the mutex is locked.
#endif // NDEBUG
};



// A class to communicate conditions between threads.
class ThreadCondition
{
public:
	ThreadCondition (ThreadMutex &a_rMutex);
		// Default constructor.
	
	~ThreadCondition();
		// Destructor.
	
	void Signal (void);
		// Signal the condition.
		// The mutex must be locked to call this.
	
	void Wait (void);
		// Wait for the condition to be signaled.
		// The mutex must be locked to call this.

private:
	ThreadMutex &m_rMutex;
		// Provides mutual exclusion for the condition.
	pthread_cond_t m_oCondition;
		// The way the condition gets signaled if another thread is
		// waiting for it.  Protected by the mutex.
	bool m_bFlag;
		// The way the condition gets signaled if another thread is
		// not waiting for it.  Protected by the mutex.
};

// A basic input/output-oriented thread.
class BasicThread
{
public:
	BasicThread();
		// Default constructor.
	
	virtual ~BasicThread();
		// Destructor.

protected:
	virtual int Work (void) = 0;
		// Denoise the current frame.
		// Must be implemented by the subclass.

	void Initialize (void);
		// Start the thread.
		// Called by the subclass' Initialize() method.
	
	void Lock (void);
		// Lock the mutex that guards the conditions.
	
	void Unlock (void);
		// Unlock the mutex that guards the conditions.

	void SignalInput (void);
		// Signal that input has been provided to the thread.
	
	void WaitForInput (void);
		// Wait for input to be provided.
	
	void SignalOutput (void);
		// Signal that output has been provided to the thread.
	
	void WaitForOutput (void);
		// Wait for output to be provided.
	
	void Shutdown (void);
		// Stop the thread.

	pthread_t m_oThreadInfo;
		// OS-level thread state.
	
private:
	ThreadMutex m_oMutex;
		// A mutex to guard the following conditions.

	ThreadCondition m_oInputCondition, m_oOutputCondition;
		// Conditions to signal input-ready and output-ready.
	
	static void *WorkLoop (void *a_pThread);
		// The thread function stub.  Calls the virtual WorkLoop().
	
protected:
	bool m_bWaitingForInput, m_bWaitingForOutput;
		// true if we're waiting for input/output (i.e. if they need
		// to be signaled).

	bool m_bWorkLoop;
		// true if WorkLoop() should continue running, false if it
		// should exit.
	
	virtual void WorkLoop (void);
		// The thread function.  Loop, calling Work(), until told to
		// stop.

	int m_nWorkRetval;
		// The value returned by Work().
};

// A class to run denoisers in a separate thread.
class DenoiserThread : public BasicThread
{
private:
	typedef BasicThread BaseClass;
		// Keep track of who our base class is.

public:
	DenoiserThread();
		// Default constructor.
	
	virtual ~DenoiserThread();
		// Destructor.

	// Prototypes for methods that subclasses must implement.

	//void Initialize (subclass parameters);
		// Initialize.  Set up all private thread data and start the
		// worker thread.
	
	//void AddFrame (subclass parameters);
		// Add a frame to the denoiser.
	
	//void WaitForAddFrame (subclass parameters);
		// Wait for the frame to be added & output possibly generated.
	
	//void Shutdown (subclass parameters);
		// Shut down the worker thread.
	
	void Shutdown (void);
		// Stop the thread.

protected:

	virtual void WorkLoop (void);
		// The thread function.  Loop, calling Work(), until told to
		// stop.
};

// A class to run the intensity denoiser in a separate thread.
class DenoiserThreadY : public DenoiserThread
{
private:
	typedef DenoiserThread BaseClass;
		// Keep track of who our base class is.

public:
	DenoiserThreadY();
		// Default constructor.
	
	virtual ~DenoiserThreadY();
		// Destructor.

	void Initialize (void);
		// Initialize.  Set up all private thread data and start the
		// worker thread.
	
	void AddFrame (const uint8_t *a_pInputY, uint8_t *a_pOutputY);
		// Add a frame to the denoiser.
	
	int WaitForAddFrame (void);
		// Get the next denoised frame, if any.
		// Returns the result of Work().
	
	void Shutdown (void);
		// Stop the thread.

protected:
	virtual int Work (void);
		// Denoise the current frame.

private:
	const uint8_t *m_pInputY;
	uint8_t *m_pOutputY;
		// Input/output buffers.
};

// A class to run the color denoiser in a separate thread.
class DenoiserThreadCbCr : public DenoiserThread
{
private:
	typedef DenoiserThread BaseClass;
		// Keep track of who our base class is.

public:
	DenoiserThreadCbCr();
		// Default constructor.
	
	virtual ~DenoiserThreadCbCr();
		// Destructor.

	void Initialize (void);
		// Initialize.  Set up all private thread data and start the
		// worker thread.
	
	void AddFrame (const uint8_t *a_pInputCb,
			const uint8_t *a_pInputCr, uint8_t *a_pOutputCb,
			uint8_t *a_pOutputCr);
		// Add a frame to the denoiser.
	
	int WaitForAddFrame (void);
		// Get the next denoised frame, if any.
		// Returns the result of Work().
	
	void Shutdown (void);
		// Stop the thread.

protected:
	virtual int Work (void);
		// Denoise the current frame.

private:
	const uint8_t *m_pInputCb;
	const uint8_t *m_pInputCr;
	uint8_t *m_pOutputCb;
	uint8_t *m_pOutputCr;
		// Input/output buffers.
};

// A class to read/write raw-video in a separate thread.
class ReadWriteThread : public BasicThread
{
private:
	typedef BasicThread BaseClass;
		// Keep track of who our base class is.

public:
	ReadWriteThread();
		// Default constructor.
	
	virtual ~ReadWriteThread();
		// Destructor.

	void Initialize (int a_nFD, const y4m_stream_info_t *a_pStreamInfo,
			y4m_frame_info_t *a_pFrameInfo, int a_nWidthY,
			int a_nHeightY, int a_nWidthCbCr, int a_nHeightCbCr);
		// Start the thread.

protected:
	int m_nFD;
		// The file descriptor from which to read/write raw video.
	const y4m_stream_info_t *m_pStreamInfo;
		// Information on the raw video stream.
	y4m_frame_info_t *m_pFrameInfo;
		// Information on the current frame in the raw video stream.
	struct Frame { uint8_t *planes[3]; Frame *next; };
		// The type of a frame.
	Frame *m_apFrames;
		// Space for frames being read from input.
	Frame *m_pValidFramesHead, *m_pValidFramesTail, *m_pCurrentFrame,
			*m_pFreeFramesHead;
		// A list of frames containing data, the current frame, and
		// a list of free frames.
	
	Frame *GetFirstValidFrame (void);
		// Remove the first valid frame from the list and return it.

	void AddFrameToValidList (Frame *a_pFrame);
		// Add the given frame to the end of the valid-frame list.

	Frame *GetFreeFrame (void);
		// Remove a frame from the free-frames list and return it.

	void AddFrameToFreeList (Frame *a_pFrame);
		// Add the given frame to the free-frames list.

	void MoveValidFrameToCurrent (void);
		// Remove the first valid frame, and make it the current frame.
	
	void MoveCurrentFrameToValidList (void);
		// Move the current frame to the end of the valid-frame list.
	
	void MoveFreeFrameToCurrent (void);
		// Remove a frame from the free list, and make it the current
		// frame.
	
	void MoveCurrentFrameToFreeList (void);
		// Move the current frame to the free-frame list.
};

// A class to read raw-video in a separate thread.
class DenoiserThreadRead : public ReadWriteThread
{
private:
	typedef ReadWriteThread BaseClass;
		// Keep track of who our base class is.

public:
	DenoiserThreadRead();
		// Default constructor.
	
	virtual ~DenoiserThreadRead();
		// Destructor.

	int ReadFrame (uint8_t **a_apPlanes);
		// Read a frame from input.  m_apPlanes[] gets backpatched
		// with pointers to valid frame data, and they are valid until
		// the next call to ReadFrame().
		// Returns Y4M_OK if it succeeds, Y4M_ERR_EOF at the end of
		// the stream.  (Returns other errors too.)
	
	void Shutdown (void);
		// Stop the thread.

protected:
	virtual int Work (void);
		// Read frames from the raw-video stream.
};

// A class to write raw-video in a separate thread.
class DenoiserThreadWrite : public ReadWriteThread
{
private:
	typedef ReadWriteThread BaseClass;
		// Keep track of who our base class is.

public:
	DenoiserThreadWrite();
		// Default constructor.
	
	virtual ~DenoiserThreadWrite();
		// Destructor.

	int GetSpaceToWriteFrame (uint8_t **a_apPlanes);
		// Get space for a frame to write to output.  a_apPlanes[] gets
		// backpatched with pointers to valid frame data.
		// Returns Y4M_OK if it succeeds, something else if it fails.
	
	void WriteFrame (void);
		// Write a frame to output.  The a_apPlanes[] previously set up
		// by GetSpaceToWriteFrame() must be filled with video data by
		// the client.
	
	void Shutdown (void);
		// Stop the thread.

protected:
	virtual int Work (void);
		// Write frames to the raw-video stream.

	virtual void WorkLoop (void);
		// The thread function.  Loop, calling Work(), until told to
		// stop.
};

// Threads for denoising intensity & color.
DenoiserThreadY g_oDenoiserThreadY;
DenoiserThreadCbCr g_oDenoiserThreadCbCr;

// Threads for reading and writing raw video.
DenoiserThreadRead g_oDenoiserThreadRead;
DenoiserThreadWrite g_oDenoiserThreadWrite;



// Initialize the denoising system.
int newdenoise_init (int a_nFrames, int a_nWidthY, int a_nHeightY,
	int a_nWidthCbCr, int a_nHeightCbCr, int a_nInputFD,
	int a_nOutputFD, const y4m_stream_info_t *a_pStreamInfo,
	y4m_frame_info_t *a_pFrameInfo)
{
	Status_t eStatus;
		// An error that may occur.
	int nInterlace;
		// A factor to apply to frames/frame-height because of
		// interlacing.
	
	// No errors yet.
	eStatus = g_kNoError;

	// Save the width and height.
	g_nWidthY = a_nWidthY;
	g_nHeightY = a_nHeightY;
	g_nWidthCbCr = a_nWidthCbCr;
	g_nHeightCbCr = a_nHeightCbCr;

	// If the video is interlaced, that means the denoiser will see
	// twice as many frames, half their original height.
	nInterlace = (denoiser.interlaced) ? 2 : 1;
	
	// If input/output should be handled in separate threads, set that
	// up.
	if (denoiser.threads >= 1)
	{
		g_oDenoiserThreadRead.Initialize (a_nInputFD,
			a_pStreamInfo, a_pFrameInfo, a_nWidthY, a_nHeightY,
			a_nWidthCbCr, a_nHeightCbCr);
		g_oDenoiserThreadWrite.Initialize (a_nOutputFD,
			a_pStreamInfo, a_pFrameInfo, a_nWidthY, a_nHeightY,
			a_nWidthCbCr, a_nHeightCbCr);
	}

	// If intensity should be denoised, set it up.
	if (a_nWidthY != 0 && a_nHeightY != 0)
	{
		g_bMotionSearcherY = true;
		g_nPixelsY = a_nWidthY * a_nHeightY / nInterlace;
		g_pPixelsY = new MotionSearcherY::Pixel_t [g_nPixelsY];
		if (g_pPixelsY == NULL)
			return -1;
		g_oMotionSearcherY.Init (eStatus, nInterlace * a_nFrames,
			a_nWidthY, a_nHeightY / nInterlace,
			denoiser.radiusY, denoiser.radiusY,
			denoiser.zThresholdY, denoiser.thresholdY,
			denoiser.matchCountThrottle, denoiser.matchSizeThrottle);
		if (eStatus != g_kNoError)
		{
			delete[] g_pPixelsY;
			return -1;
		}
	}
	else
		g_bMotionSearcherY = false;

	// If color should be denoised, set it up.
	if (a_nWidthCbCr != 0 && a_nHeightCbCr != 0)
	{
		g_bMotionSearcherCbCr = true;
		g_nPixelsCbCr = a_nWidthCbCr * a_nHeightCbCr / nInterlace;
		g_pPixelsCbCr = new MotionSearcherCbCr::Pixel_t [g_nPixelsCbCr];
		if (g_pPixelsCbCr == NULL)
			return -1;
		g_oMotionSearcherCbCr.Init (eStatus, nInterlace * a_nFrames,
			a_nWidthCbCr, a_nHeightCbCr / nInterlace,
			denoiser.radiusCbCr / denoiser.frame.ss_h,
			denoiser.radiusCbCr / denoiser.frame.ss_v,
			denoiser.zThresholdCbCr, denoiser.thresholdCbCr,
			denoiser.matchCountThrottle, denoiser.matchSizeThrottle);
		if (eStatus != g_kNoError)
			return -1;

		// If color should be denoised in a separate thread, set that
		// up.
		if (denoiser.threads == 2)
			g_oDenoiserThreadCbCr.Initialize();
	}
	else
		g_bMotionSearcherCbCr = false;