tvinpaint.c

Go to the documentation of this file.

#include <math.h>
#include <stdio.h>
#include <string.h>
#include <ctype.h>
#include "num.h"
#include "tvreg.h"
#include "imageio.h"

#define DISPLAY_SCALING             255

#define JPEGQUALITY                 95

#ifdef NUM_SINGLE
#define IMAGEIO_NUM           (IMAGEIO_SINGLE)
#else
#define IMAGEIO_NUM           (IMAGEIO_DOUBLE)
#endif

typedef struct
{
    num *Data;
    int Width;
    int Height;
    int NumChannels;
} image;


void PrintHelpMessage()
{
    puts(
    "Total variation regularized inpainting IPOL demo, P. Getreuer, 2012\n\n"
    "Syntax: tvinpaint <D> <lambda> <input> <inpainted>\n");
        puts("where <D>, <input>, and <inpainted> are " 
    READIMAGE_FORMATS_SUPPORTED " images.\n");        
        puts("Example:\n"
    "  tvinpaint mountains-D.bmp 1e3 mountains-f.bmp inpainted.bmp\n");
}

int Inpaint(image u, image f, image D, num Lambda);
num ComputeRmse(image f, image u);
void ThresholdD(image D, num Lambda);
int IsGrayscale(image f);


int main(int argc, char **argv)
{
    const char *InputFile, *DomainFile, *OutputFile;
    image f = {NULL, 0, 0, 0}, u = {NULL, 0, 0, 0}, D = {NULL, 0, 0, 0};
    num Lambda;
    int Status = 1;
    
    if(argc != 5)
    {
        PrintHelpMessage();
        return 0;
    }
    
    /* Read command line arguments */    
    DomainFile = argv[1];
    Lambda = (num)atof(argv[2]);
    InputFile = argv[3];    
    OutputFile = argv[4];    
    
    /* Read the input images */
    if(!(f.Data = (num *)ReadImage(&f.Width, &f.Height, InputFile,
        IMAGEIO_RGB | IMAGEIO_PLANAR | IMAGEIO_NUM)) || 
        !(D.Data = (num *)ReadImage(&D.Width, &D.Height, DomainFile,
        IMAGEIO_RGB | IMAGEIO_PLANAR | IMAGEIO_NUM)))
        goto Catch;
    else if(f.Width != D.Width || f.Height != D.Height)
    {
        fprintf(stderr, "Size mismatch: D is %dx%d but f is %dx%d\n",
            D.Width, D.Height, f.Width, f.Height);
        goto Catch;
    }
    
    f.NumChannels = IsGrayscale(f) ? 1 : 3;
    u = f;
    
    /* Allocate space for the inpainted image */
    if(!(u.Data = (num *)Malloc(sizeof(num) * ((size_t)f.Width)
        * ((size_t)f.Height) * f.NumChannels)))
    {
        fprintf(stderr, "Memory allocation failed.\n");
        goto Catch;
    }
    
    if(!Inpaint(u, f, D, Lambda))
    {
        fprintf(stderr, "Failure!\n");
        goto Catch;
    }
    
    /* Write inpainted image */
    if(!WriteImage(u.Data, u.Width, u.Height, OutputFile, 
        ((u.NumChannels == 1) ? IMAGEIO_GRAYSCALE : IMAGEIO_RGB)
        | IMAGEIO_PLANAR | IMAGEIO_NUM, JPEGQUALITY))    
        fprintf(stderr, "Error writing to \"%s\".\n", OutputFile);
    
    Status = 0;
Catch:
    if(u.Data)
        Free(u.Data);
    if(D.Data)
        Free(D.Data);
    if(f.Data)
        Free(f.Data);
    return Status;
}


int Inpaint(image u, image f, image D, num Lambda)
{
    tvregopt *Opt = NULL;
    const long NumPixels = ((long)f.Width) * ((long)f.Height);
    num *Red = D.Data;
    num *Green = D.Data + NumPixels;
    num *Blue = D.Data + 2*NumPixels;
    long n, k;
    int Success = 0;
    
    if(!(Opt = TvRegNewOpt()))
    {
        fprintf(stderr, "Memory allocation failed\n");
        return 0;
    }
    
    memcpy(u.Data, f.Data, sizeof(num)*f.Width*f.Height*f.NumChannels);
    
    /* Convert the mask into spatially-varing lambda */       
    for(n = 0; n < NumPixels; n++)
        if(0.299*Red[n] + 0.587*Green[n] + 0.114*Blue[n] > 0.5)
        {
            D.Data[n] = 0;          /* Inside of the inpainting domain */
            
            /* Set u = 0.5 within D */
            for(k = 0; k < u.NumChannels; k++)
                u.Data[n + k*NumPixels] = 0.5;
        }
        else
            D.Data[n] = Lambda;    /* Outside of the inpainting domain */
    
    TvRegSetVaryingLambda(Opt, D.Data, D.Width, D.Height);
    TvRegSetMaxIter(Opt, 250);
    TvRegSetTol(Opt, (num)1e-5);
    
    /* TvRestore performs the split Bregman inpainting */
    if(!TvRestore(u.Data, f.Data, f.Width, f.Height, f.NumChannels, Opt))
    {
        fprintf(stderr, "Error in computation.\n");
        goto Catch;
    }    
    
    Success = 1;
Catch:
    TvRegFreeOpt(Opt);
    return Success;
}


int IsGrayscale(image f)
{
    const long NumPixels = ((long)f.Width) * ((long)f.Height);
    const num *Red = f.Data;
    const num *Green = f.Data + NumPixels;
    const num *Blue = f.Data + 2*NumPixels;
    long n;
    
    for(n = 0; n < NumPixels; n++)
        if(Red[n] != Green[n] || Red[n] != Blue[n])
            return 0;
    
    return 1;
}