/*---------------------------------------------------------------------------*/
/* Program:  ahe.c                                                           */
/*                                                                           */
/* Purpose:  This program implements Prof. Pizer's adaptive histogram        */
/*           equalization algorithm for constrast enhancement.  The only     */
/*           option presently supported is the partitioning of the input     */
/*           image into N by N subimages and interpolating histograms        */
/*           to obtain the enhanced output image.                            */
/*                                                                           */
/* Author:   John Gauch                                                      */
/*                                                                           */
/* Date:     April 21, 1994    - Original program.                           */
/*           October 5, 1994   - Improved interpolation.                     */
/*           August 1, 1997    - Support for color and vector images.        */
/*                                                                           */
/* Note:     Copyright (C) The University of Kansas, 1994                    */
/*---------------------------------------------------------------------------*/
#include <IM.h>
#define MAXBIN 256

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

/*---------------------------------------------------------------------------*/
/* Purpose:  This routine performs AHE.                                      */
/*---------------------------------------------------------------------------*/
void AHE(BYTE_TYPE ** Data1, FLOAT_TYPE ** Data2)
{
   IM_TYPE *Image3;
   FLOAT_TYPE ***Map;
   int i, x, y, r, c;
   int Xstart, Xend, Ystart, Yend;
   float Xsize, Ysize, Scale;
   float Histo[MAXBIN];
   int w1, w2, w3, w4, w5;

   /* Create temporary mapping function image */
   Image3 = im_create("/dev/null", FLOAT, Number, Number, MAXBIN);
   Map = (FLOAT_TYPE ***) im_alloc3D(Image3, FLOAT);

   /* Initialize output image */
   for (y = 0; y < Ydim; y++)
      for (x = 0; x < Xdim; x++)
	 Data2[y][x] = 0.0;

   /* Preprocess each partition of the image */
   Xsize = Xdim / (float) Number;
   Ysize = Ydim / (float) Number;
   for (r = 0; r < Number; r++)
      for (c = 0; c < Number; c++)
      {
	 /* Calculate bounding rectangle for patch */
	 Xstart = (int) (c * Xsize);
	 Xend = (int) ((c + 1) * Xsize);
	 if (Xend >= Xdim)
	    Xend = Xdim;
	 Ystart = (int) (r * Ysize);
	 Yend = (int) ((r + 1) * Ysize);
	 if (Yend >= Ydim)
	    Yend = Ydim;
	 if (Debug == TRUE)
	    printf("patch %d %d range %d %d to %d %d\n",
		   c, r, Xstart, Ystart, Xend, Yend);

	 /* Calculate cumulative intensity histogram */
	 for (i = 0; i < MAXBIN; i++)
	    Histo[i] = 0;
	 for (y = Ystart; y < Yend; y++)
	    for (x = Xstart; x < Xend; x++)
	       Histo[Data1[y][x]]++;
	 for (i = 1; i < MAXBIN; i++)
	    Histo[i] = Histo[i] + Histo[i - 1];

	 /* Calculate mapping function */
	 Scale = (float) 255.0 / (1 + Histo[MAXBIN - 1]);
	 for (i = 0; i < MAXBIN; i++)
	    Map[i][r][c] = Histo[i] * Scale;
      }

   /* Process internal partitions of the image */
   for (r = 0; r < (Number - 1); r++)
      for (c = 0; c < (Number - 1); c++)
      {
	 /* Calculate bounding rectangle for patch */
	 Xstart = (int) (c * Xsize + Xsize / 2);
	 Xend = (int) ((c + 1) * Xsize + Xsize / 2);
	 Ystart = (int) (r * Ysize + Ysize / 2);
	 Yend = (int) ((r + 1) * Ysize + Ysize / 2);
	 if (Debug == TRUE)
	    printf("patch %d %d range %d %d to %d %d\n",
		   c, r, Xstart, Ystart, Xend, Yend);

	 /* Peform interpolated histogram equalization */
	 for (y = Ystart; y < Yend; y++)
	    for (x = Xstart; x < Xend; x++)
	    {
	       w1 = (Xend - 1 - x) * (Yend - 1 - y);
	       w2 = (x - Xstart) * (Yend - 1 - y);
	       w3 = (Xend - 1 - x) * (y - Ystart);
	       w4 = (x - Xstart) * (y - Ystart);
	       w5 = w1 + w2 + w3 + w4;
	       Data2[y][x] = (w1 * Map[Data1[y][x]][r][c]
			      + w2 * Map[Data1[y][x]][r][c + 1]
			      + w3 * Map[Data1[y][x]][r + 1][c]
			      + w4 * Map[Data1[y][x]][r + 1][c + 1]) / w5;
	    }
      }

   /* Process four corner partitions of the image */
   Xstart = (int) (Xsize / 2);
   Ystart = (int) (Ysize / 2);
   Xend = (int) ((Number - 1) * Xsize + Xsize / 2);
   Yend = (int) ((Number - 1) * Ysize + Ysize / 2);
   for (y = 0; y < Ystart; y++)
      for (x = 0; x < Xstart; x++)
	 Data2[y][x] = Map[Data1[y][x]][0][0];
   for (y = Yend; y < Ydim; y++)
      for (x = 0; x < Xstart; x++)
	 Data2[y][x] = Map[Data1[y][x]][Number - 1][0];
   for (y = 0; y < Ystart; y++)
      for (x = Xend; x < Xdim; x++)
	 Data2[y][x] = Map[Data1[y][x]][0][Number - 1];
   for (y = Yend; y < Ydim; y++)
      for (x = Xend; x < Xdim; x++)
	 Data2[y][x] = Map[Data1[y][x]][Number - 1][Number - 1];

   /* Process top and bottom partitions of the image */
   for (c = 0; c < (Number - 1); c++)
   {
      Xstart = (int) (c * Xsize + Xsize / 2);
      Xend = (int) ((c + 1) * Xsize + Xsize / 2);
      if (Debug == TRUE)
	 printf("column %d range %d %d\n", c, Xstart, Xend);
      for (x = Xstart; x < Xend; x++)
      {
	 w1 = Xend - 1 - x;
	 w2 = x - Xstart;
	 w3 = w1 + w2;

	 /* Handle top */
	 Ystart = 0;
	 Yend = (int) (Ysize / 2);
	 for (y = Ystart; y < Yend; y++)
	    Data2[y][x] = (w1 * Map[Data1[y][x]][0][c]
			   + w2 * Map[Data1[y][x]][0][c + 1]) / w3;

	 /* Handle bottom */
	 Ystart = (int) ((Number - 1) * Ysize + Ysize / 2);
	 Yend = Ydim;
	 for (y = Ystart; y < Yend; y++)
	    Data2[y][x] = (w1 * Map[Data1[y][x]][Number - 1][c]
			   + w2 * Map[Data1[y][x]][Number - 1][c + 1]) / w3;
      }
   }

   /* Process left and right partitions of the image */
   for (r = 0; r < (Number - 1); r++)
   {
      Ystart = (int) (r * Ysize + Ysize / 2);
      Yend = (int) ((r + 1) * Ysize + Ysize / 2);
      if (Debug == TRUE)
	 printf("row %d range %d %d\n", r, Ystart, Yend);
      for (y = Ystart; y < Yend; y++)
      {
	 w1 = Yend - 1 - y;
	 w2 = y - Ystart;
	 w3 = w1 + w2;

	 /* Handle left */
	 Xstart = 0;
	 Xend = (int) (Xsize / 2);
	 for (x = Xstart; x < Xend; x++)
	    Data2[y][x] = (w1 * Map[Data1[y][x]][r][0]
			   + w2 * Map[Data1[y][x]][r + 1][0]) / w3;

	 /* Handle right */
	 Xstart = (int) ((Number - 1) * Xsize + Xsize / 2);
	 Xend = Xdim;
	 for (x = Xstart; x < Xend; x++)
	    Data2[y][x] = (w1 * Map[Data1[y][x]][r][Number - 1]
			   + w2 * Map[Data1[y][x]][r + 1][Number - 1]) / w3;
      }
   }
}

/*---------------------------------------------------------------------------*/
/* Purpose:  Calculate AHE for a scalar image.                               */
/*---------------------------------------------------------------------------*/
void AheScalar()
{
   BYTE_TYPE **Data1;
   FLOAT_TYPE **Data2;

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

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

   /* Perform AHE */
   AHE(Data1, Data2);

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

/*---------------------------------------------------------------------------*/
/* Purpose:  Calculate AHE for a vector image.                               */
/*---------------------------------------------------------------------------*/
void AheVector()
{
   BYTE_TYPE ***Data1;
   FLOAT_TYPE ***Data2;
   int z;

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

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

   /* Perform AHE */
   for (z = 0; z < Zdim; z++)
      AHE(Data1[z], Data2[z]);

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

/*---------------------------------------------------------------------------*/
/* Purpose:  Calculate AHE for a color image (separate colors).              */
/*---------------------------------------------------------------------------*/
void AheColor()
{
   COLOR_TYPE **Data1;
   BYTE_TYPE **Data2;
   FLOAT_TYPE **Data3;
   int x, y;

   /* 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 = (BYTE_TYPE **) im_alloc2D(Image2, BYTE);
   Data3 = (FLOAT_TYPE **) im_alloc2D(Image2, FLOAT);

   /* Copy R pixels */
   for (y = 0; y < Ydim; y++)
      for (x = 0; x < Xdim; x++)
	 Data2[y][x] = Data1[y][x].r;

   /* Perform AHE on R pixels */
   AHE(Data2, Data3);

   /* Copy R, G pixels */
   for (y = 0; y < Ydim; y++)
      for (x = 0; x < Xdim; x++)
      {
	 Data1[y][x].r = (BYTE_TYPE) Data3[y][x];
	 Data2[y][x] = Data1[y][x].g;
      }

   /* Perform AHE on G pixels */
   AHE(Data2, Data3);

   /* Copy G, B pixels */
   for (y = 0; y < Ydim; y++)
      for (x = 0; x < Xdim; x++)
      {
	 Data1[y][x].g = (BYTE_TYPE) Data3[y][x];
	 Data2[y][x] = Data1[y][x].b;
      }

   /* Perform AHE on B pixels */
   AHE(Data2, Data3);

   /* Copy B pixels */
   for (y = 0; y < Ydim; y++)
      for (x = 0; x < Xdim; x++)
	 Data1[y][x].b = (BYTE_TYPE) Data3[y][x];

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

/*---------------------------------------------------------------------------*/
/* Purpose:  Calculate AHE for a color image (intensity only).               */
/*---------------------------------------------------------------------------*/
void AheColorInten()
{
   COLOR_TYPE **Data1;
   BYTE_TYPE **Data2;
   FLOAT_TYPE **Data3;
   int x, y, 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 = (BYTE_TYPE **) im_alloc2D(Image2, BYTE);
   Data3 = (FLOAT_TYPE **) im_alloc2D(Image2, FLOAT);

   /* Calculate pixel intensity */
   for (y = 0; y < Ydim; y++)
      for (x = 0; x < Xdim; x++)
	 Data2[y][x] = Data1[y][x].r / 3
	    + Data1[y][x].g / 3
	    + Data1[y][x].b / 3;

   /* Perform AHE */
   AHE(Data2, Data3);

   /* Increase color contrast */
   for (y = 0; y < Ydim; y++)
      for (x = 0; x < Xdim; x++)
	 if (Data2[y][x] > 0)
	 {
	    temp = (int) (Data1[y][x].r * Data3[y][x]) / Data2[y][x];
	    if (temp < 0)
	       temp = 0;
	    else if (temp > 255)
	       temp = 255;
	    Data1[y][x].r = temp;
	    temp = (int) (Data1[y][x].g * Data3[y][x]) / Data2[y][x];
	    if (temp < 0)
	       temp = 0;
	    else if (temp > 255)
	       temp = 255;
	    Data1[y][x].g = temp;
	    temp = (int) (Data1[y][x].b * Data3[y][x]) / Data2[y][x];
	    if (temp < 0)
	       temp = 0;
	    else if (temp > 255)
	       temp = 255;
	    Data1[y][x].b = temp;
	 }

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

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

   /* Interpret program options */
   printf("AHE 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");
	 break;
      case 'd':
	 Debug = TRUE;
	 break;
      default:
	 Error("Invalid option encountered");
	 break;
      }

   /* Check number of file names */
   if (argc - i != 2)
   {
      fprintf(stderr, "Usage: ahe [options] infile outfile\n");
      fprintf(stderr, "       [-d]    Print debugging information\n");
      fprintf(stderr, "       [-n #]  Number of subregions in x and y (%d)\n",
	      Number);
      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");

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

   /* Check range on Number parameter */
   if (Number < 1)
      Number = 1;
   if (Number > Xdim / 2)
      Number = Xdim / 2;
   if (Number > Ydim / 2)
      Number = Ydim / 2;

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

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

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

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