/*---------------------------------------------------------------------------*/
/* Program:  cve.c                                                           */
/*                                                                           */
/* Purpose:  This program performs Constant Variance image Enhancement.      */
/*           The pixels in a neighborhood about each point in the            */
/*           image are used to estimate the local mean and variance.         */
/*           The size of this neighborhood is set by the -n switch.          */
/*           The value of the center pixel in the neighborhood is then       */
/*           normalized by subtracting the mean and dividing by the          */
/*           standard deviation (the sqrt of the variance).  This value      */
/*           is scaled so the resulting image has the same global mean       */
/*           and variance as the input image.                                */
/*                                                                           */
/*           For debugging purposes, it is also possible to output the       */
/*           local mean (using -m) or local standard deviation (using -s).   */
/*           The amount of contrast enhancement performed is controlled      */
/*           by a gain factor (specified by the -g switch).  A gain of 0.0   */
/*           returns the original image while a gains greater than 1.0       */
/*           yields very enhanced images.                                    */
/*                                                                           */
/* Author:   John Gauch, Edu Metz                                            */
/*                                                                           */
/* Date:     September 25, 1995 - Original program.                          */
/*           July 24, 1996      - Color support added.                       */
/*           July 30, 1997      - Simplified color, vector, scalar code.     */
/*                                                                           */
/* Note:     Copyright (C) The University of Kansas, 1995-1996               */
/*---------------------------------------------------------------------------*/
#include <IM.h>

/* Output options */
#define  CVE    0
#define  MEAN   1
#define  STDDEV 2
#define  SMALL  0.0001

/* Program variables */
int Debug = FALSE;
int Output = CVE;
int Number = 9;
float Gain = 0.5;

/* Image variables */
char Name1[50];
char Name2[50];
IM_TYPE *Image1;
IM_TYPE *Image2;
int PixType, Xdim, Ydim, Zdim, DimCnt;

/*---------------------------------------------------------------------------*/
/* Purpose:  This routine performs CVE on a scalar image.                    */
/*---------------------------------------------------------------------------*/
void CveScalar()
{
   /* Program variables */
   FLOAT_TYPE **Data1;
   FLOAT_TYPE **Data2;
   float GlobalMean, GlobalVariance;
   float GlobalMax, GlobalMin;
   float LocalMean, LocalVariance;
   int StartX, EndX, StartY, EndY;
   int Count, x, y, nx, ny;

   /* Read input image */
   Data1 = (FLOAT_TYPE **) im_alloc2D(Image1, FLOAT);
   im_read(Image1, FLOAT, (char *) &(Data1[0][0]));

   /* Create output image */
   Image2 = im_create(Name2, PixType, Xdim, Ydim, Zdim);
   Data2 = (FLOAT_TYPE **) im_alloc2D(Image2, FLOAT);

   /* Initialize variables */
   GlobalMean = 0.0;
   GlobalVariance = 0.0;
   GlobalMax = Data1[0][0];
   GlobalMin = Data1[0][0];

   /* Calculate global mean and variance */
   /* ---------------------------------- */
   for (y = 0; y < Ydim; y++)
      for (x = 0; x < Xdim; x++)
      {
	 GlobalMean += Data1[y][x];
	 GlobalVariance += Data1[y][x] * Data1[y][x];
	 if (Data1[y][x] > GlobalMax)
	    GlobalMax = Data1[y][x];
	 if (Data1[y][x] < GlobalMin)
	    GlobalMin = Data1[y][x];
      }

   GlobalMean = GlobalMean / (Xdim * Ydim);
   GlobalVariance = GlobalVariance / (Xdim * Ydim)
      - GlobalMean * GlobalMean;

   /* Calculate local statistics and CVE output */
   /* ----------------------------------------- */
   for (y = 0; y < Ydim; y++)
      for (x = 0; x < Xdim; x++)
      {
	 /* Specify local neighborhood */
	 StartX = x - Number / 2;
	 if (StartX < 0)
	    StartX = 0;

	 EndX = StartX + Number;
	 if (EndX >= Xdim)
	    EndX = Xdim - 1;

	 StartY = y - Number / 2;
	 if (StartY < 0)
	    StartY = 0;

	 EndY = StartY + Number;
	 if (EndY >= Ydim)
	    EndY = Ydim - 1;

	 /* Init statistics */
	 Count = (EndX - StartX) * (EndY - StartY);
	 LocalMean = 0.0;
	 LocalVariance = 0.0;

	 /* Gather local statistics */
	 for (ny = StartY; ny < EndY; ny++)
	    for (nx = StartX; nx < EndX; nx++)
	    {
	       LocalMean += Data1[ny][nx];
	       LocalVariance += Data1[ny][nx] * Data1[ny][nx];
	    }

	 /* Calculate local statistics */
	 if (Count == 0)
	    Count = 1;
	 LocalMean = LocalMean / Count;
	 LocalVariance = (LocalVariance / Count) - (LocalMean * LocalMean);

	 /* Generate output value */
	 if (Output == MEAN)
	    Data2[y][x] = LocalMean;
	 else if (Output == STDDEV)
	    Data2[y][x] = (float) sqrt((double) LocalVariance);
	 else if (LocalVariance < SMALL)
	    Data2[y][x] = Data1[y][x];
	 else
	 {
	    /* Perform CVE */
	    Data2[y][x] = (float) (Data1[y][x]
				   + Gain * (Data1[y][x] - LocalMean)
				   * sqrt(GlobalVariance / LocalVariance));
	    if (Data2[y][x] < GlobalMin)
	       Data2[y][x] = GlobalMin;
	    if (Data2[y][x] > GlobalMax)
	       Data2[y][x] = GlobalMax;
	 }
      }

   /* Write information to output image */
   im_write(Image2, FLOAT, (char *) &(Data2[0][0]));
   im_free2D((char **) Data1);
   im_free2D((char **) Data2);
}

/*---------------------------------------------------------------------------*/
/* Purpose:  This routine performs CVE on a vector image.                    */
/*---------------------------------------------------------------------------*/
void CveVector()
{
   FLOAT_TYPE ***Data1;
   FLOAT_TYPE ***Data2;
   float GlobalMean, GlobalVariance;
   float GlobalMax, GlobalMin;
   float LocalMean, LocalVariance;
   int StartX, EndX, StartY, EndY;
   int Count, x, y, z, nx, ny;

   /* Read input image */
   Data1 = (FLOAT_TYPE ***) im_alloc3D(Image1, FLOAT);
   im_read(Image1, FLOAT, (char *) &(Data1[0][0][0]));

   /* Create output image */
   Image2 = im_create(Name2, PixType, Xdim, Ydim, Zdim);
   Data2 = (FLOAT_TYPE ***) im_alloc3D(Image2, FLOAT);

   /* Process all images in Z dimension */
   for (z = 0; z < Zdim; z++)
   {
      /* Initialize variables */
      GlobalMean = 0.0;
      GlobalVariance = 0.0;
      GlobalMax = Data1[z][0][0];
      GlobalMin = Data1[z][0][0];

      /* Calculate global mean and variance */
      /* ---------------------------------- */
      for (y = 0; y < Ydim; y++)
	 for (x = 0; x < Xdim; x++)
	 {
	    GlobalMean += Data1[z][y][x];
	    GlobalVariance += Data1[z][y][x] * Data1[z][y][x];
	    if (Data1[z][y][x] > GlobalMax)
	       GlobalMax = Data1[z][y][x];
	    if (Data1[z][y][x] < GlobalMin)
	       GlobalMin = Data1[z][y][x];
	 }

      GlobalMean = GlobalMean / (Xdim * Ydim);
      GlobalVariance = GlobalVariance / (Xdim * Ydim)
	 - GlobalMean * GlobalMean;

      /* Calculate local statistics and CVE output */
      /* ----------------------------------------- */
      for (y = 0; y < Ydim; y++)
	 for (x = 0; x < Xdim; x++)
	 {
	    /* Specify local neighborhood */
	    StartX = x - Number / 2;
	    if (StartX < 0)
	       StartX = 0;

	    EndX = StartX + Number;
	    if (EndX >= Xdim)
	       EndX = Xdim - 1;

	    StartY = y - Number / 2;
	    if (StartY < 0)
	       StartY = 0;

	    EndY = StartY + Number;
	    if (EndY >= Ydim)
	       EndY = Ydim - 1;

	    /* Init statistics */
	    Count = (EndX - StartX) * (EndY - StartY);
	    LocalMean = 0.0;
	    LocalVariance = 0.0;

	    /* Gather local statistics */
	    for (ny = StartY; ny < EndY; ny++)
	       for (nx = StartX; nx < EndX; nx++)
	       {
		  LocalMean += Data1[z][ny][nx];
		  LocalVariance += Data1[z][ny][nx] * Data1[z][ny][nx];
	       }

	    /* Calculate local statistics */
	    if (Count == 0)
	       Count = 1;
	    LocalMean = LocalMean / Count;
	    LocalVariance = (LocalVariance / Count) - (LocalMean * LocalMean);

	    /* Generate output value */
	    if (Output == MEAN)
	       Data2[z][y][x] = LocalMean;
	    else if (Output == STDDEV)
	       Data2[z][y][x] = (float) sqrt((double) LocalVariance);
	    else if (LocalVariance < SMALL)
	       Data2[z][y][x] = Data1[z][y][x];
	    else
	    {
	       /* Perform CVE */
	       Data2[z][y][x] = (float) (Data1[z][y][x]
				       + Gain * (Data1[z][y][x] - LocalMean)
				    * sqrt(GlobalVariance / LocalVariance));
	       if (Data2[z][y][x] < GlobalMin)
		  Data2[z][y][x] = GlobalMin;
	       if (Data2[z][y][x] > GlobalMax)
		  Data2[z][y][x] = GlobalMax;
	    }
	 }
   }

   /* Write information to output image */
   im_write(Image2, FLOAT, (char *) &(Data2[0][0][0]));
   im_free3D((char ***) Data1);
   im_free3D((char ***) Data2);
}

/*---------------------------------------------------------------------------*/
/* Purpose:  This routine performs CVE on a color image.                     */
/*---------------------------------------------------------------------------*/
void CveColor()
{
   /* Program variables */
   COLOR_TYPE **Data1;
   COLOR_TYPE **Data2;
   float GlobalMeanR, GlobalMeanG, GlobalMeanB;
   float GlobalVarianceR, GlobalVarianceG, GlobalVarianceB;
   float LocalMeanR, LocalMeanG, LocalMeanB;
   float LocalVarianceR, LocalVarianceG, LocalVarianceB;
   int StartX, EndX, StartY, EndY;
   int Count, x, y, nx, ny;
   int Temp;

   /* Read input image */
   Data1 = (COLOR_TYPE **) im_alloc2D(Image1, COLOR);
   im_read(Image1, COLOR, (char *) &(Data1[0][0]));

   /* Create output image */
   Image2 = im_create(Name2, PixType, Xdim, Ydim, Zdim);
   Data2 = (COLOR_TYPE **) im_alloc2D(Image2, COLOR);

   /* Initialize variables */
   GlobalMeanR = GlobalMeanG = GlobalMeanB = 0.0;
   GlobalVarianceR = GlobalVarianceG = GlobalVarianceB = 0.0;

   /* Calculate global mean and variance */
   /* ---------------------------------- */
   for (y = 0; y < Ydim; y++)
      for (x = 0; x < Xdim; x++)
      {
	 GlobalMeanR += Data1[y][x].r;
	 GlobalMeanG += Data1[y][x].g;
	 GlobalMeanB += Data1[y][x].b;
	 GlobalVarianceR += (float) Data1[y][x].r * (float) Data1[y][x].r;
	 GlobalVarianceG += (float) Data1[y][x].g * (float) Data1[y][x].g;
	 GlobalVarianceB += (float) Data1[y][x].b * (float) Data1[y][x].b;
      }

   GlobalMeanR = GlobalMeanR / (Xdim * Ydim);
   GlobalMeanG = GlobalMeanG / (Xdim * Ydim);
   GlobalMeanB = GlobalMeanB / (Xdim * Ydim);
   GlobalVarianceR = GlobalVarianceR / (Xdim * Ydim)
      - GlobalMeanR * GlobalMeanR;
   GlobalVarianceG = GlobalVarianceG / (Xdim * Ydim)
      - GlobalMeanG * GlobalMeanG;
   GlobalVarianceB = GlobalVarianceB / (Xdim * Ydim)
      - GlobalMeanB * GlobalMeanB;

   /* Calculate local statistics and CVE output */
   /* ----------------------------------------- */
   for (y = 0; y < Ydim; y++)
      for (x = 0; x < Xdim; x++)
      {
	 /* Specify local neighborhood */
	 StartX = x - Number / 2;
	 if (StartX < 0)
	    StartX = 0;

	 EndX = StartX + Number;
	 if (EndX >= Xdim)
	    EndX = Xdim - 1;

	 StartY = y - Number / 2;
	 if (StartY < 0)
	    StartY = 0;

	 EndY = StartY + Number;
	 if (EndY >= Ydim)
	    EndY = Ydim - 1;

	 /* Init statistics */
	 Count = (EndX - StartX) * (EndY - StartY);
	 LocalMeanR = LocalMeanG = LocalMeanB = 0.0;
	 LocalVarianceR = LocalVarianceG = LocalVarianceB = 0.0;

	 /* Gather local statistics */
	 for (ny = StartY; ny < EndY; ny++)
	    for (nx = StartX; nx < EndX; nx++)
	    {
	       LocalMeanR += Data1[ny][nx].r;
	       LocalMeanG += Data1[ny][nx].g;
	       LocalMeanB += Data1[ny][nx].b;
	       LocalVarianceR += Data1[ny][nx].r * Data1[ny][nx].r;
	       LocalVarianceG += Data1[ny][nx].g * Data1[ny][nx].g;
	       LocalVarianceB += Data1[ny][nx].b * Data1[ny][nx].b;
	    }

	 /* Calculate local statistics */
	 if (Count == 0)
	    Count = 1;
	 LocalMeanR = LocalMeanR / Count;
	 LocalMeanG = LocalMeanG / Count;
	 LocalMeanB = LocalMeanB / Count;
	 LocalVarianceR = (LocalVarianceR / Count) - (LocalMeanR * LocalMeanR);
	 LocalVarianceG = (LocalVarianceG / Count) - (LocalMeanG * LocalMeanG);
	 LocalVarianceB = (LocalVarianceB / Count) - (LocalMeanB * LocalMeanB);

	 /* Generate output value */
	 if (Output == MEAN)
	 {
	    Data2[y][x].r = (BYTE_TYPE) LocalMeanR;
	    Data2[y][x].g = (BYTE_TYPE) LocalMeanG;
	    Data2[y][x].b = (BYTE_TYPE) LocalMeanB;
	 }
	 else if (Output == STDDEV)
	 {
	    Data2[y][x].r = (BYTE_TYPE) sqrt((double) LocalVarianceR);
	    Data2[y][x].g = (BYTE_TYPE) sqrt((double) LocalVarianceG);
	    Data2[y][x].b = (BYTE_TYPE) sqrt((double) LocalVarianceB);
	 }
	 else if ((LocalVarianceR < SMALL) ||
		  (LocalVarianceG < SMALL) ||
		  (LocalVarianceB < SMALL))
	 {
	    Data2[y][x].r = Data1[y][x].r;
	    Data2[y][x].g = Data1[y][x].g;
	    Data2[y][x].b = Data1[y][x].b;
	 }
	 else
	 {
	    /* Perform CVE on R channel */
	    Temp = (int) (Data1[y][x].r
			  + Gain * (Data1[y][x].r - LocalMeanR)
			  * sqrt(GlobalVarianceR / LocalVarianceR));
	    if (Temp < 0)
	       Temp = 0;
	    else if (Temp > 255)
	       Temp = 255;
	    Data2[y][x].r = (BYTE_TYPE) Temp;

	    /* Perform CVE on G channel */
	    Temp = (int) (Data1[y][x].g
			  + Gain * (Data1[y][x].g - LocalMeanG)
			  * sqrt(GlobalVarianceG / LocalVarianceG));
	    if (Temp < 0)
	       Temp = 0;
	    else if (Temp > 255)
	       Temp = 255;
	    Data2[y][x].g = (BYTE_TYPE) Temp;

	    /* Perform CVE on B channel */
	    Temp = (int) (Data1[y][x].b
			  + Gain * (Data1[y][x].b - LocalMeanB)
			  * sqrt(GlobalVarianceB / LocalVarianceB));
	    if (Temp < 0)
	       Temp = 0;
	    else if (Temp > 255)
	       Temp = 255;
	    Data2[y][x].b = (BYTE_TYPE) Temp;
	 }
      }

   /* Write information to output image */
   im_write(Image2, COLOR, (char *) &(Data2[0][0]));
   im_free2D((char **) Data1);
   im_free2D((char **) Data2);
}

/*---------------------------------------------------------------------------*/
/* Purpose:  This is the main program.                                       */
/*---------------------------------------------------------------------------*/
int main(int argc, char *argv[])
{
   /* Program variables */
   int i = 0;

   /* Interpret program options */
   printf("CVE Program - KUIM Version 2.1\n\n");
   while ((++i < argc) && (argv[i][0] == '-'))
      switch (argv[i][1])
      {
      case 'n':
	 if (sscanf(argv[++i], "%d", &Number) == 0)
	    Error("Could not get integer argument");
	 if (Number < 1)
	    Number = 1;
	 break;
      case 'g':
	 if (sscanf(argv[++i], "%f", &Gain) == 0)
	    Error("Could not get float argument");
	 if (Gain < 0.0)
	    Gain = 0.0;
	 break;
      case 'm':
	 Output = MEAN;
	 break;
      case 's':
	 Output = STDDEV;
	 break;
      case 'd':
	 Debug = TRUE;
	 break;
      default:
	 Error("Invalid option encountered");
	 break;
      }

   /* Check number of file names */
   if (argc - i != 2)
   {
      fprintf(stderr, "Usage: cve [options] infile outfile\n");
      fprintf(stderr, "       [-d]    Print debugging information\n");
      fprintf(stderr, "       [-m]    Output local mean\n");
      fprintf(stderr, "       [-s]    Output local standard deviation\n");
      fprintf(stderr, "       [-n #]  Width of local neighborhood (9)\n");
      fprintf(stderr, "       [-g #]  Contrast gain factor (0.5)\n");
      exit(1);
   }

   /* Get image file names from argument list */
   if (sscanf(argv[i++], "%s", Name1) == 0)
      Error("Could not get input file name");
   if (sscanf(argv[i++], "%s", Name2) == 0)
      Error("Could not get output file name");

   /* Determine PixType */
   Image1 = im_open(Name1, &PixType, &Xdim, &Ydim, &Zdim, &DimCnt);

   /* Handle color images */
   if (PixType == COLOR || PixType == PSEUDO || PixType == JPEG_RGB)
      CveColor();

   /* Handle scalar images */
   else if (DimCnt == 2)
      CveScalar();

   /* Handle vector images */
   else if (DimCnt == 3)
      CveVector();

   /* Unsupported image type */
   else
      Error("Can not process image");
   return (0);
}