color_swap.py

import pygame as pg

       
def color_swap(source_image, swap_map):        
    """
    Creates a new Surface from the source_image with some or all colors
    swapped for new colors. Colors are swapped according to the
    color pairs in the swap_map dict. The keys and values in swap_map
    can be RGB tuples or pygame color-names. For each key in swap_map,
    all pixels of that color will be replaced by the color that key maps to. 
    For example, passing this dict:
    
    {(0,255,0): (255, 0, 255),
      "black": (255, 0, 0),
      "yellow": "green"}
      
    would result in green pixels recolored purple, black pixels recolored
    red and yellow pixels recolored green.
    NOTE: This will not work if Pygame's video mode has not been set
    (i.e., you need to call pygame.display.set_mode beforehand).
    """   
    img = source_image
    size = img.get_size()
    surf = pg.Surface(size).convert()
    color_surf = pg.Surface(size).convert()
    final = img.copy()
    for original_color, new_color in swap_map.items():
        if isinstance(original_color, str):
            original = pg.Color(original_color)
        else:
            original = original_color
        if isinstance(new_color, str):    
            recolor = pg.Color(new_color)
        else:
            recolor = new_color        
        color_surf.fill(original)
        surf.set_colorkey(original)
        pg.transform.threshold(surf, img, original, (0,0,0,0),
                               recolor, 1, color_surf, True)
        final.blit(surf, (0,0))
    return final
    
    
def make_test_image():
    """
    Returns a Surface with some different colored rectangles
    drawn on it for testing the color_swap function.
    """
    surf = pg.Surface((800, 800)).convert()
    surf.fill(pg.Color("black"))
    pg.draw.rect(surf, pg.Color("blue"), (50, 50, 50, 50))    
    pg.draw.rect(surf, pg.Color("red"), (300, 300, 50, 50))    
    pg.draw.rect(surf, pg.Color("green"), (50, 500, 50, 50))    
    pg.draw.rect(surf, pg.Color("purple"), (500, 400, 50, 50))    
    return surf
 
def swap_test():
    """Creates an image, recolors it and saves a copy of each."""
    pg.init()
    SCREEN = pg.display.set_mode((800, 800))
    black = pg.Color(0,0,0)
    img = make_test_image()
    swaps = {(0,0,0): "lightgray",
             "blue": (100,255,100),
             "red": "salmon",
             "purple": "orchid",
             "white": "chartreuse"}
    recolored = color_swap(img, swaps)
    
    pg.image.save(img, "original_image.png")
    pg.image.save(recolored, "recolored_image.png")
    
    
if __name__ == "__main__":
    swap_test()