/* determine average color of insulator image

              _                                                   
             (_)_ __   __ _ _ __ ___   ___  __ _ _ __         ___ 
             | | '_ \ / _` | '_ ` _ \ / _ \/ _` | '_ \       / __|
             | | |_) | (_| | | | | | |  __/ (_| | | | |  _  | (__ 
            _/ | .__/ \__, |_| |_| |_|\___|\__,_|_| |_| (_)  \___|
           |__/|_|    |___/                                       

   ian macky feb-19-2004
*/

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <setjmp.h>
#include <string.h>
#include "jpgmean.h"

#define MIN_PIXELS		10	/* min number of pixel samples */
#define WHITE_THRESH		40	/* if white% over this, white mode */
#define BLACK_THRESH		40	/* if black% over this, black mode */

/* return distance between RGB coordinates (r1,g1,b1) and (r2,g2,b2) */

double rgb_distance(int r1, int g1, int b1, int r2, int g2, int b2)
{
    int d_r = (r1 - r2), d_g = (g1 - g2), d_b = (b1 - b2);

    return sqrt((d_r * d_r) + (d_g * d_g) + (d_b * d_b));
}

/* compute mean RGB of jpg and optionally write split-image map.

    ARGUMENT    DIR  DESCRIPTION
    ---------- 	---  ------------------------------------------------
    image	 I   path of image to be processed
    ofile	 I   if generating split image, the output jpg's path
    white_dist	 I   distance from (255,255,255) considered "white"
    black_dist	 I   distance from (0,0,0) considered "black"
    std_units	 I   pixels must be within this many standard units
    mean_red	 O   returned mean red value for selected pixels
    mean_green	 O   returned mean green value for selected pixels
    mean_blue	 O   returned mean blue value for selected pixels
    width	 O   width of image (NULL if don't want it)
    height	 O   height of image (NULL if don't want it)
    debug	 I   debug output level, 0=none, ++ for more detail

    computes the mean R G B values of selected pixels, returning them
    through <mean_red>, <mean_green>, and <mean_blue>.  the image's width
    and height are also returned through <width> and <height> if not NULL.

    the average is computed by three passes:

    PASS 1: WHITE, BLACK, RAW AVERAGE

    the first pass does an initial filtering for white, black, and colored
    pixels.  typically about 1/4 of the pixels in an image will be white
    (not counting any border, which is first stripped).  so for auto-white
    mode, if the white% is > 1/3, it's assumed the insulator is light colored,
    and is rescanned with white_dist 0 and black_dist 200.  a hard white
    pixel (255,255,255) is presumed to be background, but all other light
    pixels are kept.  most dark pixels are throw out; typically they are
    from optical effects or are opaque liners or other parts.

    similarly, for auto black level, if the black% is > 1/3, the black_dist
    is set to 5 and white_dist to 200, and the first pass repeated.  there
    is no black border to be stripped.

    the thinking here is: a blackglass insulator, so you want to ignore
    the background and specular highlights, pretty much all whitish pixels,
    and pay attention to almost all dark ones.

    finally, a raw average is computed for red, green, and blue.  with that,
    we can do the next step.

    PASS 2: STANDARD DEVIATION

    the image is rescanned with the same white_dist and black_dist as
    the previous round, and the sum of the differences from the raw mean
    is computed, leading to the standard deviation for each channel.  this
    is used for the final filtering pass.

    PASS 3: FILTER DEVIANTS

    one more pass is made to filter out deviant pixels.  white_dist and
    black_dist are as before, but now filter must pass a third gauntlet:
    they must be within std_units standard units or be excluded.  a sum
    of all pixels passing all three tests is made, and the final average
    computed.

    if std_units is specified as -1, it means use the default, 2.0.
    standard units for each channel are simply (red - mean_red) / red_s,
    e.g. difference from the mean red value divided by red's standard
    deviation.
*/

#define START_DECOMPRESS \
    jd->err = jpeg_std_error(&jerr.pub); \
    jerr.pub.error_exit = jpgmean_error_exit; \
    if (setjmp(jerr.env)) \
	goto jpeglib_error; \
    jpeg_create_decompress(jd); \
    jpeg_stdio_src(jd, in_f); \
    jpeg_read_header(jd, TRUE); \
    jpeg_start_decompress(jd);
 
#define RESTART_DECOMPRESS \
    rewind(in_f); \
    START_DECOMPRESS

struct jpgmean_error_mgr {
    struct jpeg_error_mgr pub;
    jmp_buf               env;
} jerr;

METHODDEF(void) jpgmean_error_exit(j_common_ptr cinfo) { longjmp(jerr.env, 1); }

boolean jpg_mean(char *image, char *ofile, char **error_string,
		 double white_dist, double black_dist, double std_units,
		 JSAMPLE *mean_red, JSAMPLE *mean_green, JSAMPLE *mean_blue,
		 unsigned *width, unsigned *height, boolean debug)
{
    /* jpeg input controls */
    JSAMPLE jinrow[MAX_IMAGE_WIDTH * 3];   /* 3=R G B for each pixel */
    JSAMPROW jinarray[1] = { jinrow };
    struct jpeg_decompress_struct ijd, *jd = &ijd;
    JSAMPLE joutrow[MAX_IMAGE_WIDTH * 3];
    JSAMPROW joutarray[1] = { joutrow };
    struct jpeg_compress_struct jc;
    JSAMPLE *jp;

    FILE         *in_f, *out_f;
    JSAMPLE       red, green, blue, *p;
    int           r, c, middle_column, white_pct, black_pct;
    int		  min_r, max_r, min_c, max_c, box_width, box_height;
    double        red_s, green_s, blue_s, delta;
    double        red_z, green_z, blue_z;
    double        sum_d_red, sum_d_green, sum_d_blue;
    double        raw_mean_red, raw_mean_green, raw_mean_blue;
    unsigned long colored_pixels, red_total, green_total, blue_total;
    unsigned long white_pixels, black_pixels, total_pixels, border_pixels;
    boolean       auto_white, auto_black, auto_su, in_auto;

    if (debug)
	printf("JPG Mean of image '%s'\n", image);

    if (!(in_f = fopen(image, "rb")))
    {
	printf("failed to open image '%s'\n", image);	
	return FALSE;
    }

    if ((auto_white = (white_dist < 0)))
	white_dist = WHITE_DISTANCE;

    if ((auto_black = (black_dist < 0)))
	black_dist = BLACK_DISTANCE;

    if ((auto_su = (std_units < 0)))
	std_units = STANDARD_UNITS;

    in_auto = FALSE;	/* oscillation preventer */

    /* --- PASS 1 ---------------------------------------------------------- */

    START_DECOMPRESS

    if (debug)
	printf("Image is %d rows by %d columns.\n",
		jd->image_height, jd->image_width);

    if (jd->image_height > MAX_IMAGE_WIDTH)
    {
	printf("Image too wide (%d), maximum allowed is %d.\n", 
		jd->image_width, MAX_IMAGE_WIDTH);
	return FALSE;
    }

restart:
    if (width)				/* not required to return these */
	*width = jd->image_width;
    if (height)
	*height = jd->image_height;

    red = green = blue = 0;
    red_total = green_total = blue_total = 0;
    colored_pixels = white_pixels = black_pixels = 0;

    min_r = jd->image_height;
    min_c = jd->image_width;
    max_r = max_c = 0;

    /* scan input image one line at a time */
    for (r = 0; r < jd->image_height; r++)	/* for each scanline */
    {
	/* grab next scan line */
	jpeg_read_scanlines(&ijd, jinarray, 1);
	for (p = jinrow, c = 0; c < jd->image_width; c++)
	{
	    red = *p++; green = *p++; blue = *p++;
	    
	    if (rgb_distance(red, green, blue, 255,255,255) <= white_dist)
		white_pixels++;
	    else
	    {
		/* track image box (sans white border) */
		if (r < min_r) min_r = r;
		if (r > max_r) max_r = r;
		if (c < min_c) min_c = c;
		if (c > max_c) max_c = c;

		if (rgb_distance(red, green, blue, 0,0,0) <= black_dist)
		    black_pixels++;
		else
		{
		    red_total   += red;
		    green_total += green;
		    blue_total  += blue;
		    colored_pixels++;
		}
	    }
	}
    }

    /* actual area occupied by image */
    box_width = max_c - min_c + 1;
    box_height = max_r - min_r + 1;
    if (debug)
	printf("Colored area is rows %d to %d, columns %d to %d\n",
		min_r, max_r, min_c, max_c);

    /* total number of pixels in bounding box of colored pixels */
    total_pixels = box_width * box_height;

    /* number of pixels in the white border */
    border_pixels = (jd->image_height * jd->image_width) - total_pixels;

    /* we want the number of white pixels in the bounding box, so adjust */
    white_pixels -= border_pixels;

    /* seems to be light insulator with lots of white pixels, reset
       params for white mode and rescan
    */
    white_pct = white_pixels * 100 / total_pixels;
    if (auto_white && (white_pct > WHITE_THRESH) && !in_auto)
    {
	if (debug)
	    printf("White %d%% exceeds threshhold %d%%, restarting\n",
			white_pct, WHITE_THRESH);
	white_dist = 0;
	black_dist = 200;
	in_auto = TRUE;
	RESTART_DECOMPRESS
	goto restart;
    }

    /* seems to be black insulator with lots of black pixels, reset
       params for black mode and rescan
    */
    black_pct = black_pixels * 100 / total_pixels;
    if (auto_black && (black_pct > BLACK_THRESH) && !in_auto)
    {						/* seems to be blackish */
	if (debug)
	    printf("Black %d%% exceeds threshhold %d%%, restarting\n",
			black_pct, BLACK_THRESH);
	black_dist = 5;
	white_dist = 200;
	in_auto = TRUE;
	RESTART_DECOMPRESS 
	goto restart;
    }

    /* figure average R G B of colored pixels */
    raw_mean_red   = (double) red_total   / colored_pixels;
    raw_mean_green = (double) green_total / colored_pixels;
    raw_mean_blue  = (double) blue_total  / colored_pixels;

    if (debug)
    {
	printf("White: distance = %g, pixels = %ld = %ld%%\n",
	    white_dist, white_pixels, white_pixels * 100 / total_pixels);

	printf("Black: distance = %g, pixels = %ld = %ld%%\n",
	    black_dist, black_pixels, black_pixels * 100 / total_pixels);

	printf("Included %ld pixels of %ld = %ld%%\n",
		colored_pixels, total_pixels,
		colored_pixels * 100 / total_pixels);

	printf("    Red   mean = %3d\n", (int) raw_mean_red);
	printf("    Green mean = %3d\n", (int) raw_mean_green);
	printf("    Blue  mean = %3d\n", (int) raw_mean_blue);

	printf("RAW: %02X%02X%02X\n",
	    (unsigned) raw_mean_red, (unsigned) raw_mean_green,
	    (unsigned) raw_mean_blue);
    }

    /* --- PASS 2 ---------------------------------------------------------- */

    /* we have the mean, go again and compute standard deviation */

    RESTART_DECOMPRESS

    sum_d_red = sum_d_green = sum_d_blue = 0;
    for (r = 0; r < jd->image_height; r++)	/* for each scanline */
    {
	/* read next scan line */
	jpeg_read_scanlines(&ijd, jinarray, 1);
	if ((r < min_r) || (r > max_r))
	    continue;
	for (p = jinrow, c = 0; c < jd->image_width; c++)
	{
	    red = *p++; green = *p++; blue = *p++;
	    if ((c < min_c) || (c > max_c))
		continue;

	    /* if not a white or black pixel */
	    if ((rgb_distance(red, green, blue, 255,255,255) >= white_dist) &&
		(rgb_distance(red, green, blue, 0,0,0) >= black_dist))
	    {
		delta = (double) red - (double) raw_mean_red;
		sum_d_red += delta * delta;

		delta = (double) green - (double) raw_mean_green;
		sum_d_green += delta * delta;

		delta = (double) blue - (double) raw_mean_blue;
		sum_d_blue += delta * delta;
	    }
	}
    }

    /* note colored_pixels should be the same as from pass 1 */
    red_s   = sqrt(sum_d_red   / colored_pixels);
    green_s = sqrt(sum_d_green / colored_pixels);
    blue_s  = sqrt(sum_d_blue  / colored_pixels);

    if (debug)
    {
	printf("    Red   s = %g\n", red_s);
	printf("    Green s = %g\n", green_s);
	printf("    Blue  s = %g\n", blue_s);
    }

    /* --- PASS 3 ---------------------------------------------------------- */

    /* now that we have the standard deviation, once more to filter deviants */

    RESTART_DECOMPRESS

    /* scan input image */
    red_total = green_total = blue_total = 0;
    colored_pixels = 0;
    for (r = 0; r < jd->image_height; r++)	/* for each scanline */
    {
	jpeg_read_scanlines(&ijd, jinarray, 1);
	if ((r < min_r) || (r > max_r))
	    continue;
	for (p = jinrow, c = 0; c < jd->image_width; c++)
	{
	    red = *p++; green = *p++; blue = *p++;
	    if ((c < min_c) || (c > max_c))
		continue;

	    if ((rgb_distance(red, green, blue, 255,255,255) >= white_dist) &&
		(rgb_distance(red, green, blue, 0,0,0) >= black_dist))
	    {
		if ((red_z = ((double) red - raw_mean_red) / red_s) < 0)
		    red_z = -red_z;

		if ((green_z = ((double) green - raw_mean_green) / green_s) < 0)
		    green_z = -green_z;

		if ((blue_z  = ((double) blue  - raw_mean_blue) / blue_s) < 0)
		    blue_z = -blue_z;

		if ((red_z   < std_units) &&
		    (green_z < std_units) &&
		    (blue_z  < std_units))
		{
		    red_total   += red;
		    green_total += green;
		    blue_total  += blue;
		    colored_pixels++;
		}
	    }
	}
    }

    if (debug)
	printf("Included %ld pixels of %ld = %ld%%\n",
		colored_pixels, total_pixels,
		colored_pixels * 100 / total_pixels);

    if (colored_pixels < MIN_PIXELS)
    {
	*error_string = "Sorry, not enough pixels were selected "
			"to compute an average";
	if (in_f)
	    fclose(in_f);
	return FALSE;
    }

    /* return final mean RGB */
    *mean_red   = (double) red_total   / colored_pixels;
    *mean_green = (double) green_total / colored_pixels;
    *mean_blue  = (double) blue_total  / colored_pixels;

    if (debug)
    {
	printf("    Red   mean = %3d\n", (int) *mean_red);
	printf("    Green mean = %3d\n", (int) *mean_green);
	printf("    Blue  mean = %3d\n", (int) *mean_blue);

	printf("END: %02X%02X%02X\n", (unsigned) *mean_red,
		(unsigned) *mean_green, (unsigned) *mean_blue);
    }

    /* if all we're doing is computing average, then we're done */
    if (!ofile)
    {
	fclose(in_f);
	return TRUE;
    }

    if (debug)
	printf("Outputting split image to '%s'\n", ofile);

    /* they want a split image produced */
    if (!(out_f = fopen(ofile, "wb")))
    {
	printf("couldn't create output file '%s'", ofile);
	fclose(in_f);
	return FALSE;
    }

    RESTART_DECOMPRESS

    /* create compression structure */
    memset(&jc, 0, sizeof(jc));
    jc.err = jpeg_std_error(&jerr.pub);
    jpeg_create_compress(&jc);
    jpeg_stdio_dest(&jc, out_f);
    jc.image_width = box_width;
    jc.image_height = box_height;
    jc.in_color_space = JCS_RGB;
    jc.input_components = 3;     /* R G & B per pixel */
    jpeg_set_defaults(&jc);
    jpeg_start_compress(&jc, TRUE);

    /* left half gets original, right half shows pixels selection & average */
    middle_column = jd->image_width / 2;

    for (r = 0; r < jd->image_height; r++)	/* for each scanline */
    {
	/* read next scan line */
	jpeg_read_scanlines(&ijd, jinarray, 1);
	if ((r < min_r) || (r > max_r))
	    continue;

	jp = joutrow;				/* output scan line */

	for (p = jinrow, c = 0; c < jd->image_width; c++)
	{
	    red = *p++; green = *p++; blue = *p++;
	    if ((c < min_c) || (c > max_c))
		continue;

	    if (c <= middle_column)	/* left half? */
	    {
		*jp++ = red;		/* output gets original values */
		*jp++ = green;
		*jp++ = blue;
	    }
	    else			/* right half gets pixel map */
	    {
		if (rgb_distance(red, green, blue, 255,255,255) <= white_dist)
		{
		    *jp++ = 0xFF;	/* WHITE */
		    *jp++ = 0xFF;
		    *jp++ = 0xFF;
		}
		else if (rgb_distance(red, green, blue, 0,0,0) <= black_dist)
		{
		    *jp++ = 0;		/* BLACK */
		    *jp++ = 0;
		    *jp++ = 0;
		}
		else
		{
		    red_z = (double) (red - raw_mean_red) / red_s;
		    if (red_z < 0)
			red_z = -red_z;

		    green_z = (double) (green - raw_mean_green) / green_s;
		    if (green_z < 0)
			green_z = -green_z;

		    blue_z  = (double) (blue  - raw_mean_blue) / blue_s;
		    if (blue_z < 0)
			blue_z = -blue_z;

		    if ((red_z   < std_units) &&
			(green_z < std_units) &&
			(blue_z  < std_units))
		    {
			*jp++ = *mean_red;	/* pixel used for color avg */
			*jp++ = *mean_green;	/* it gets that avg color */
			*jp++ = *mean_blue;
		    }
		    else
		    {
			*jp++ = 0x7F;		/* GRAY */
			*jp++ = 0x7F;
			*jp++ = 0x7F;
		    }
		}
	    }
	}
	jpeg_write_scanlines(&jc, joutarray, 1); /* write one scanline */
    }
    jpeg_finish_compress(&jc);
    fclose(out_f);
    fclose(in_f);
    return TRUE;

jpeglib_error:
    if (in_f)
	fclose(in_f);
    *error_string = "JPEGLIB error";
    return FALSE;
}

/* $Id: jpgmean.c,v 1.6 2004/03/17 19:16:35 imacky Exp ian $ */