Commits

marcus  committed 46a7864

Added weight support to PixelArray.extract() and PixelArray.replace().

  • Participants
  • Parent commits 123fae0

Comments (0)

Files changed (6)

File docs/ref/pixelarray.html

 <a name="PixelArray.replace">
 <big><b>PixelArray.replace</big></b><br><ul>
   <i>Replaces the passed color in the PixelArray with another one.</i><br>
-  <tt>PixelArray.replace (color, repcolor, distance=0): Return None</tt><br>
+  <tt>PixelArray.replace (color, repcolor, distance=0, weights=(0.299, 0.587, 0.114)): Return None</tt><br>
 <p>Replaces the pixels with the passed color in the PixelArray by changing them them to the passed replacement color. </p>
 <p>It uses a simple weighted euclidian distance formula to calculate the distance between the colors. The distance space ranges from <tt>0.0</tt> to <tt>1.0</tt> and is used as threshold for the color detection. This causes the replacement to take pixels with a similar but not exactly identical color, into account as well. </p>
 <p>This is an in place operation that directly affects the pixels of the PixelArray. </p>
 <a name="PixelArray.extract">
 <big><b>PixelArray.extract</big></b><br><ul>
   <i>Extracts the passed color from the PixelArray.</i><br>
-  <tt>PixelArray.extract (color, distance=0): Return PixelArray</tt><br>
+  <tt>PixelArray.extract (color, distance=0, weights=(0.299, 0.587, 0.114)): Return PixelArray</tt><br>
 <p>Extracts the passed color by changing all matching pixels to white, while non-matching pixels are changed to black. This returns a new PixelArray with the black/white color mask. </p>
 <p>It uses a simple weighted euclidian distance formula to calculate the distance between the colors. The distance space ranges from <tt>0.0</tt> to <tt>1.0</tt> and is used as threshold for the color detection. This causes the extraction to take pixels with a similar but not exactly identical color, into account as well. </p>
 <p>New in <tt>pygame.1.8.1</tt>. </p>

File src/pixelarray.c

                 /* Construct the columns */
                 pixel = (Uint32) *((Uint8 *) pixels + x + y * array->padding);
                 PyString_ConcatAndDel (&string, PyString_FromFormat
-                    ("%ld, ", pixel));
+                    ("%ld, ", (long)pixel));
                 x += array->xstep;
                 posx += absxstep;
             }
             pixel = (Uint32) *((Uint8 *) pixels + x + y * array->padding);
             PyString_ConcatAndDel (&string,
-                PyString_FromFormat ("%ld]", pixel));
+                PyString_FromFormat ("%ld]", (long)pixel));
             y += array->ystep;
             posy += absystep;
         }
                 pixel = (Uint32)
                     *((Uint16 *) (pixels + y * array->padding) + x);
                 PyString_ConcatAndDel (&string, PyString_FromFormat
-                    ("%ld, ", pixel));
+                    ("%ld, ", (long)pixel));
                 x += array->xstep;
                 posx += absxstep;
             }
             pixel = (Uint32) *((Uint16 *) (pixels + y * array->padding) + x);
             PyString_ConcatAndDel (&string,
-                PyString_FromFormat ("%ld]", pixel));
+                PyString_FromFormat ("%ld]", (long)pixel));
             y += array->ystep;
             posy += absystep;
         }
                 pixel = (px24[2]) + (px24[1] << 8) + (px24[0] << 16);
 #endif
                 PyString_ConcatAndDel (&string, PyString_FromFormat
-                    ("%ld, ", pixel));
+                    ("%ld, ", (long)pixel));
                 x += array->xstep;
                 posx += absxstep;
             }
             pixel = (px24[2]) + (px24[1] << 8) + (px24[0] << 16);
 #endif
             PyString_ConcatAndDel (&string,
-                PyString_FromFormat ("%ld]", pixel));
+                PyString_FromFormat ("%ld]", (long)pixel));
             y += array->ystep;
             posy += absystep;
         }
                 /* Construct the columns */
                 pixel = *((Uint32 *) (pixels + y * array->padding) + x);
                 PyString_ConcatAndDel (&string, PyString_FromFormat
-                    ("%ld, ", pixel));
+                    ("%ld, ", (long)pixel));
                 x += array->xstep;
                 posx += absxstep;
             }
             pixel = *((Uint32 *) (pixels + y * array->padding) + x);
             PyString_ConcatAndDel (&string,
-                PyString_FromFormat ("%ld]", pixel));
+                PyString_FromFormat ("%ld]", (long)pixel));
             y += array->ystep;
             posy += absystep;
         }

File src/pixelarray.doc

 
 replace
 Replaces the passed color in the PixelArray with another one.
-PixelArray.replace (color, repcolor, distance=0): Return None
+PixelArray.replace (color, repcolor, distance=0, weights=(0.299, 0.587, 0.114)): Return None
 
 Replaces the pixels with the passed color in the PixelArray by changing
 them them to the passed replacement color.
 
 extract
 Extracts the passed color from the PixelArray.
-PixelArray.extract (color, distance=0): Return PixelArray
+PixelArray.extract (color, distance=0, weights=(0.299, 0.587, 0.114)): Return PixelArray
 
 Extracts the passed color by changing all matching pixels to white,
 while non-matching pixels are changed to black. This returns a new

File src/pixelarray_methods.c

  * It receives RGB values in the range 0-255 and returns a distance
  * value between 0.0 and 1.0.
  */
-#define COLOR_DIFF_RGB(r1,g1,b1,r2,g2,b2) \
-    (sqrt (.299 * (r1 - r2) * (r1 - r2) + \
-           .587 * (g1 - g2) * (g1 - g2) + \
-           .114 * (b1 - b2) * (b1 - b2)) / 255.0)
+#define COLOR_DIFF_RGB(wr,wg,wb,r1,g1,b1,r2,g2,b2) \
+    (sqrt (wr * (r1 - r2) * (r1 - r2) + \
+           wg * (g1 - g2) * (g1 - g2) + \
+           wb * (b1 - b2) * (b1 - b2)) / 255.0)
 
+#define WR_NTSC 0.299
+#define WG_NTSC 0.587
+#define WB_NTSC 0.114
 
 /**
  * Tries to retrieve a valid color for a Surface.
     return newsf;
 }
 
+static int
+_get_weights (PyObject *weights, float *wr, float *wg, float *wb)
+{
+    int success = 1;
+    float rgb[3] = { 0 };
+    
+    if (!weights)
+    {
+        *wr = WR_NTSC;
+        *wg = WG_NTSC;
+        *wb = WB_NTSC;
+        return 1;
+    }
+    
+    if (!PySequence_Check (weights))
+    {
+        PyErr_SetString (PyExc_TypeError, "weights must be a sequence");
+        success = 0;
+    }
+    else if (PySequence_Size (weights) < 3)
+    {
+        PyErr_SetString (PyExc_TypeError,
+            "weights must contain at least 3 values");
+        success = 0;
+    }
+    else
+    {
+        PyObject *item;
+        int i;
+        
+        for (i = 0; i < 3; i++)
+        {
+            item = PySequence_GetItem (weights, i);
+            if (PyNumber_Check (item))
+            {
+                PyObject *num = NULL;
+                if ((num = PyNumber_Float (item)) != NULL)
+                {
+                    rgb[i] = (float) PyFloat_AsDouble (num);
+                    Py_DECREF (num);
+                }
+                else if ((num = PyNumber_Int (item)) != NULL)
+                {
+                    rgb[i] = (float) PyInt_AsLong (num);
+                    if (rgb[i] == -1 && PyErr_Occurred ())
+                        success = 0;
+                    Py_DECREF (num);
+                }
+                else if ((num = PyNumber_Long (item)) != NULL)
+                {
+                    rgb[i] = (float) PyLong_AsLong (num);
+                    if (PyErr_Occurred () &&
+                        PyErr_ExceptionMatches (PyExc_OverflowError))
+                        success = 0;
+                    Py_DECREF (num);
+                }
+            }
+            else
+            {
+                PyErr_SetString (PyExc_TypeError, "invalid weights");
+                success = 0;
+            }
+            Py_XDECREF (item);
+            if (!success)
+                break;
+        }
+    }
+    
+    if (success)
+    {
+        float sum = 0;
+        
+        *wr = rgb[0];
+        *wg = rgb[1];
+        *wb = rgb[2];
+        if ((*wr < 0 || *wg < 0 || *wb < 0) ||
+            (*wr == 0 && *wg == 0 && *wb == 0))
+        {
+            PyErr_SetString (PyExc_ValueError,
+                "weights must be positive and greater than 0");
+            return 0;
+        }
+        /* Build the average weight values. */
+        sum = *wr + *wg + *wb;
+        *wr = *wr / sum;
+        *wg = *wg / sum;
+        *wb = *wb / sum;
+        
+        return success;
+    }
+    return 0;
+}
+
 static PyObject*
 _replace_color (PyPixelArray *array, PyObject *args, PyObject *kwds)
 {
+    PyObject *weights = NULL;
     PyObject *delcolor = NULL;
     PyObject *replcolor = NULL;
     Uint32 dcolor;
     Uint8 r1, g1, b1, r2, g2, b2, a2;
     SDL_Surface *surface;
     float distance = 0;
-    char *type;
+    float wr, wg, wb;
 
     Uint32 x = 0;
     Uint32 y = 0;
     Sint32 absystep;
     Uint8 *pixels;
 
-    static char *keys[] = { "color", "repcolor", "distance", NULL };
+    static char *keys[] = { "color", "repcolor", "distance", "weights", NULL };
     
-    if (!PyArg_ParseTupleAndKeywords (args, kwds, "OO|f", keys, &delcolor,
-            &replcolor, &distance))
+    if (!PyArg_ParseTupleAndKeywords (args, kwds, "OO|fO", keys, &delcolor,
+            &replcolor, &distance, &weights))
         return NULL;
 
     if (distance < 0 || distance > 1)
         !_get_color_from_object (replcolor, surface->format, &rcolor))
         return NULL;
 
+    if (!_get_weights (weights, &wr, &wg, &wb))
+        return NULL;
+
     surface = PySurface_AsSurface (array->surface);
     pixels = surface->pixels;
     absxstep = ABS (array->xstep);
                 if (distance)
                 {
                     GET_PIXELVALS_1 (r2, g2, b2, a2, pixel, surface->format);
-                    if (COLOR_DIFF_RGB (r1, g1, b1, r2, g2, b2) <= distance)
+                    if (COLOR_DIFF_RGB (wr, wg, wb, r1, g1, b1, r2, g2, b2) <=
+                        distance)
                         *pixel = (Uint8) rcolor;
                 }
                 else if (*pixel == dcolor)
                 {
                     GET_PIXELVALS (r2, g2, b2, a2, (Uint32) *pixel,
                         surface->format);
-                    if (COLOR_DIFF_RGB (r1, g1, b1, r2, g2, b2) <= distance)
+                    if (COLOR_DIFF_RGB (wr, wg, wb, r1, g1, b1, r2, g2, b2) <=
+                        distance)
                         *pixel = (Uint16) rcolor;
                 }
                 else if (*pixel == dcolor)
                 if (distance)
                 {
                     GET_PIXELVALS (r2, g2, b2, a2, pxcolor, surface->format);
-                    if (COLOR_DIFF_RGB (r1, g1, b1, r2, g2, b2) <= distance)
+                    if (COLOR_DIFF_RGB (wr, wg, wb, r1, g1, b1, r2, g2, b2) <=
+                        distance)
                     {
                         *(px + (format->Rshift >> 3)) = (Uint8) (rcolor >> 16);
                         *(px + (format->Gshift >> 3)) = (Uint8) (rcolor >> 8);
                 if (distance)
                 {
                     GET_PIXELVALS (r2, g2, b2, a2, pxcolor, surface->format);
-                    if (COLOR_DIFF_RGB (r1, g1, b1, r2, g2, b2) <= distance)
+                    if (COLOR_DIFF_RGB (wr, wg, wb, r1, g1, b1, r2, g2, b2) <=
+                        distance)
                     {
                         *(px + 2 - (format->Rshift >> 3)) =
                             (Uint8) (rcolor >> 16);
                 if (distance)
                 {
                     GET_PIXELVALS (r2, g2, b2, a2, *pixel, surface->format);
-                    if (COLOR_DIFF_RGB (r1, g1, b1, r2, g2, b2) <= distance)
+                    if (COLOR_DIFF_RGB (wr, wg, wb, r1, g1, b1, r2, g2, b2) <=
+                        distance)
                         *pixel = rcolor;
                 }
                 else if (*pixel == dcolor)
 static PyObject*
 _extract_color (PyPixelArray *array, PyObject *args, PyObject *kwds)
 {
+    PyObject *weights = NULL;
     PyObject *sf = NULL;
     PyObject *excolor = NULL;
     PyPixelArray *newarray = NULL;
     Uint32 black;
     SDL_Surface *surface;
     float distance = 0;
+    float wr, wg, wb;
     Uint8 r1, g1, b1, r2, g2, b2, a2;
 
     Uint32 x = 0;
     Sint32 absystep;
     Uint8 *pixels;
 
-    static char *keys[] = { "color", "distance", NULL };
+    static char *keys[] = { "color", "distance", "weights", NULL };
 
-    if (!PyArg_ParseTupleAndKeywords (args, kwds, "O|f", keys, &excolor,
-            &distance))
+    if (!PyArg_ParseTupleAndKeywords (args, kwds, "O|fO", keys, &excolor,
+            &distance, &weights))
         return NULL;
 
     if (distance < 0 || distance > 1)
     if (!_get_color_from_object (excolor, surface->format, &color))
         return NULL;
 
+    if (!_get_weights (weights, &wr, &wg, &wb))
+        return NULL;
+
     /* Create the b/w mask surface. */
     sf = _make_surface (array);
     if (!sf)
                 if (distance)
                 {
                     GET_PIXELVALS_1 (r2, g2, b2, a2, pixel, surface->format);
-                    if (COLOR_DIFF_RGB (r1, g1, b1, r2, g2, b2) <= distance)
+                    if (COLOR_DIFF_RGB (wr, wg, wb, r1, g1, b1, r2, g2, b2) <=
+                        distance)
                         *pixel = (Uint8) white;
                     else
                         *pixel = (Uint8) black;
                 {
                     GET_PIXELVALS (r2, g2, b2, a2, (Uint32) *pixel,
                         surface->format);
-                    if (COLOR_DIFF_RGB (r1, g1, b1, r2, g2, b2) <= distance)
+                    if (COLOR_DIFF_RGB (wr, wg, wb, r1, g1, b1, r2, g2, b2) <=
+                        distance)
                         *pixel = (Uint16) white;
                     else
                         *pixel = (Uint16) black;
                 if (distance)
                 {
                     GET_PIXELVALS (r2, g2, b2, a2, pxcolor, surface->format);
-                    if (COLOR_DIFF_RGB (r1, g1, b1, r2, g2, b2) <=  distance)
+                    if (COLOR_DIFF_RGB (wr, wg, wb, r1, g1, b1, r2, g2, b2) <=
+                        distance)
                     {
                         *(px + (format->Rshift >> 3)) = (Uint8) (white >> 16);
                         *(px + (format->Gshift >> 3)) = (Uint8) (white >> 8);
                 if (distance)
                 {
                     GET_PIXELVALS (r2, g2, b2, a2, pxcolor, surface->format);
-                    if (COLOR_DIFF_RGB (r1, g1, b1, r2, g2, b2) <=
+                    if (COLOR_DIFF_RGB (wr, wg, wb, r1, g1, b1, r2, g2, b2) <=
                         distance)
                     {
                         *(px + 2 - (format->Rshift >> 3)) =
                 if (distance)
                 {
                     GET_PIXELVALS (r2, g2, b2, a2, *pixel, surface->format);
-                    if (COLOR_DIFF_RGB (r1, g1, b1, r2, g2, b2) <= distance)
+                    if (COLOR_DIFF_RGB (wr, wg, wb, r1, g1, b1, r2, g2, b2) <=
+                        distance)
                         *pixel = white;
                     else
                         *pixel = black;

File src/pygamedocs.h

 
 #define DOC_PIXELARRAYMAKESURFACE "PixelArray.make_surface (): Return Surface\nCreates a new Surface from the current PixelArray."
 
-#define DOC_PIXELARRAYREPLACE "PixelArray.replace (color, repcolor, distance=0): Return None\nReplaces the passed color in the PixelArray with another one."
+#define DOC_PIXELARRAYREPLACE "PixelArray.replace (color, repcolor, distance=0, weights=(0.299, 0.587, 0.114)): Return None\nReplaces the passed color in the PixelArray with another one."
 
-#define DOC_PIXELARRAYEXTRACT "PixelArray.extract (color, distance=0): Return PixelArray\nExtracts the passed color from the PixelArray."
+#define DOC_PIXELARRAYEXTRACT "PixelArray.extract (color, distance=0, weights=(0.299, 0.587, 0.114)): Return PixelArray\nExtracts the passed color from the PixelArray."
 
 #define DOC_PYGAMERECT "pygame.Rect(left, top, width, height): return Rect\npygame.Rect((left, top), (width, height)): return Rect\npygame.Rect(object): return Rect\npygame object for storing rectangular coordinates"
 

File test/pixelarray_test.py

             self.assertEqual (ar[3][6], oval)
             self.assertEqual (ar[8][9], rval)
             self.assertEqual (ar[9][9], oval)
+            
+            ar[::2].replace ((0, 0, 255), (255, 0, 0), weights=(10, 20, 50))
+            self.assertEqual (ar[0][0], oval)
+            self.assertEqual (ar[2][3], oval)
+            self.assertEqual (ar[3][6], oval)
+            self.assertEqual (ar[8][9], oval)
+            self.assertEqual (ar[9][9], oval)
 
     def test_extract (self):
         for bpp in (8, 16, 24, 32):
             self.assertEqual (newar[8][9], black)
             self.assertEqual (newar[9][9], black)
 
+            newar = ar.extract ((255, 0, 0), weights=(10, 0.1, 50))
+            self.assertEqual (newar[0][0], black)
+            self.assertEqual (newar[1][0], black)
+            self.assertEqual (newar[2][3], white)
+            self.assertEqual (newar[3][6], white)
+            self.assertEqual (newar[8][9], black)
+            self.assertEqual (newar[9][9], black)
+
 if __name__ == '__main__':
     unittest.main()