greenscreening with python and opencv (color differencing method)

pythonGreenscreen.py

import cv
import sys
import math
import os
import Image

USE_PYGAME = True
SHOW_RGB = False

mainWin = "RESULT"
inputWin = "INPUT"
redWin = "RED"
greenWin = "GREEN"
blueWin = "BLUE"
redSubWin = "RED SUBTRACTED"
greenSubWin = "GREEN SUBTRACTED"
blueSubWin = "BLUE SUBTRACTED"
maskRGWin = "MASK Red - Green"
maskDetailWin = "MASK Detail"
maskWin = "MASK Combined"
guiWin = "GUI WINDOW"
edgeWin = "EDGE DETECTION"

rgbWin = "RGB original"
rgbWinSubtracted = "RGB background subtracted"

##START CAPTURE
capture = cv.CreateCameraCapture(0)
# set the wanted image size from the camera
#cv.SetCaptureProperty (capture, cv._CAP_PROP_FRAME_WIDTH, 640)
#cv.SetCaptureProperty (capture, cv._CAP_PROP_FRAME_HEIGHT, 480)
frame = cv.QueryFrame (capture)

##OBJECTS
imgSize = cv.GetSize(frame)
bg = [50, 200, 50]

##IMAGES
depth = cv.IPL_DEPTH_8U
checkers = cv.LoadImageM(os.path.dirname(os.path.realpath(__file__))+"/checkers.png", cv.CV_LOAD_IMAGE_UNCHANGED)
image = cv.CreateImage (imgSize, depth, 3)
red = cv.CreateImage (imgSize, depth, 1)
green = cv.CreateImage (imgSize, depth, 1)
blue = cv.CreateImage (imgSize, depth, 1)
greenNorm = cv.CreateImage (imgSize, cv.IPL_DEPTH_32F, 1)
greenInverted = cv.CreateImage (imgSize, depth, 1)
redSub = cv.CreateImage (imgSize, depth, 1)
greenSub = cv.CreateImage (imgSize, depth, 1)
blueSub = cv.CreateImage (imgSize, depth, 1)
maskRG = cv.CreateImage (imgSize, depth, 1)
maskDetail = cv.CreateImage (imgSize, depth, 1)
maskDetailNorm = cv.CreateImage (imgSize, cv.IPL_DEPTH_32F, 1)
mask = cv.CreateImage (imgSize, depth, 1)
composition = cv.CreateImage (imgSize, depth, 4)
sub = cv.CreateImage (imgSize, depth, 1)
compositionPygame = cv.CreateImage (imgSize, depth, 4)
edgeF = cv.CreateImage (imgSize, cv.IPL_DEPTH_32F, 1)
edge = cv.CreateImage (imgSize, depth, 1)

##START THE WINDOW
offX = 0
offY = 0

cv.NamedWindow (inputWin, cv.CV_WINDOW_AUTOSIZE)
cv.MoveWindow(inputWin, imgSize[0]/2, offY)



if not USE_PYGAME:
    cv.NamedWindow (mainWin, cv.CV_WINDOW_AUTOSIZE)
    cv.MoveWindow(mainWin, offX, offY)

#MASK WIN
offX = imgSize[0]+100
cv.NamedWindow (maskRGWin, cv.CV_WINDOW_AUTOSIZE)
cv.NamedWindow (maskDetailWin, cv.CV_WINDOW_AUTOSIZE)
cv.NamedWindow (maskWin, cv.CV_WINDOW_AUTOSIZE)

cv.MoveWindow(maskRGWin, offX, offY)
cv.MoveWindow(maskDetailWin, offX+int(imgSize[0]*.5), offY)
cv.MoveWindow(maskWin, offX+imgSize[0], offY)

if SHOW_RGB:
    #RGB WIN
    offX = 0
    offY = imgSize[1]
    cv.NamedWindow (redWin, cv.CV_WINDOW_AUTOSIZE)
    cv.NamedWindow (greenWin, cv.CV_WINDOW_AUTOSIZE)
    cv.NamedWindow (blueWin, cv.CV_WINDOW_AUTOSIZE)
    
    cv.MoveWindow(redWin, offX, offY)
    cv.MoveWindow(greenWin, offX+imgSize[0], offY)
    cv.MoveWindow(blueWin, offX+imgSize[0]*2, offY)
    
    #RGB SUB
    offY = imgSize[1]+int(imgSize[1]*.5)
    cv.NamedWindow (redSubWin, cv.CV_WINDOW_AUTOSIZE)
    cv.NamedWindow (greenSubWin, cv.CV_WINDOW_AUTOSIZE)
    cv.NamedWindow (blueSubWin, cv.CV_WINDOW_AUTOSIZE)
    
    cv.MoveWindow(redSubWin, offX, offY)
    cv.MoveWindow(greenSubWin, offX+imgSize[0], offY)
    cv.MoveWindow(blueSubWin, offX+imgSize[0]*2, offY)

##GUI WINDOW
offX = 0
cv.NamedWindow(guiWin, cv.CV_WINDOW_AUTOSIZE)
cv.MoveWindow(guiWin, offX, imgSize[1])

cv.NamedWindow(edgeWin, cv.CV_WINDOW_AUTOSIZE)
cv.MoveWindow(edgeWin, imgSize[0], imgSize[1])

##MODFIABLE VARIABLES
RGSub = 200
greenRemoval = 1
detailMaskMin = 10
detailMaskMax = 200

##callbacks to modify vars
def RGSubChange(val):
    global RGSub
    RGSub = val
    
def GreenRemovalChange(val):
    global greenRemoval
    greenRemoval = val*1/100.

def detailMaskMinChange(val):
    global detailMaskMin
    detailMaskMin = val
    
def detailMaskMaxChange(val):
    global detailMaskMax
    detailMaskMax = val

##CREATE SLIDERS
cv.CreateTrackbar("RED GREEN SUBTRACT", guiWin, RGSub, 255, RGSubChange)
cv.CreateTrackbar("GREEN REMOVAL", guiWin, int(greenRemoval*100), 100, GreenRemovalChange)
cv.CreateTrackbar("DETAIL MASK MIN", guiWin, detailMaskMin, 255, detailMaskMinChange)
cv.CreateTrackbar("DETAIL MASK MAX", guiWin, detailMaskMax, 255, detailMaskMaxChange)

##SETUP PYGAME
if USE_PYGAME:
    os.environ["SDL_VIDEO_WINDOW_POS"] = "0,0"
    import pygame #have to use pygame to display the image with alpha channel
    pygame.init()
    pygame.display.set_caption("OUTPUT")
    display = pygame.display.set_mode((1280,960), pygame.HWSURFACE)
    checkers = pygame.image.load(os.path.dirname(os.path.realpath(__file__))+"/checkers.png").convert()


def ofMap(value, inputMin, inputMax, outputMin, outputMax, clamp = True):
    outVal = ((value - inputMin) / (inputMax - inputMin) * (outputMax - outputMin) + outputMin)
    if clamp:
        if outputMax < outputMin:
            if outVal < outputMax: outVal = outputMax
            elif outVal > outputMin: outVal = outputMin
        else:
            if outVal > outputMax: outVal = outputMax
            elif outVal < outputMin: outVal = outputMin
    return outVal;

def remapImage(src, target, min, max):
    lut = cv.CreateMat(256,1,cv.CV_8U)
    for i in range(256):
        lut[i,0] = ofMap(i, min, max, 0., 255.)
    
    cv.LUT(src, target, lut)

def updateBg(x, y):
    global greenInverted, red, blue
    bg[0] = cv.Get2D(red, y, x)[0]
    bg[1] = cv.Get2D(greenInverted, y, x)[0]
    bg[2] = cv.Get2D(blue, y, x)[0]

def onMouse (event, x, y, flags, param):
    if event is cv.CV_EVENT_LBUTTONUP:
        updateBg(x, y)


cv.SetMouseCallback (inputWin, onMouse, None)
if not USE_PYGAME:
    cv.SetMouseCallback (mainWin, onMouse, None)

while True:
    #######################################
    ## CAPTURE VIDEO AND PREPARE CHANNELS
    ######################################
    frame = cv.QueryFrame (capture)
    if frame is None:
        pygame.quit()
        break

    #split into channels
    cv.Split(frame, red, green, blue, None)
    
    
    #######################################
    ## CALCULATE THE DETAIL MASK
    ######################################
    
    #invert green
    cv.Not(green, greenInverted)
    
    #subtract the bg
    cv.Set(sub, bg[0])
    cv.Sub(red, sub, redSub)
    
    cv.Set(sub, bg[1])
    cv.Sub(greenInverted, sub, greenSub)
    
    cv.Set(sub, bg[2])
    cv.Sub(blue, sub, blueSub)
    
    #add all subtracted backgrounds
    cv.Add(blueSub, greenSub, maskDetail)
    cv.Add(maskDetail, redSub, maskDetail)
    remapImage(maskDetail, maskDetail, float(detailMaskMin), float(detailMaskMax))
    
    #######################################
    ## CALCULATE THE GREEN - RED MASK
    ######################################
    cv.Sub(green, red, maskRG)
    cv.Not(maskRG, maskRG)
    cv.Set(sub, RGSub)
    cv.Sub(maskRG, sub, maskRG)
    cv.Smooth(maskRG, maskRG, cv.CV_GAUSSIAN, 5)
    cv.Erode(maskRG, maskRG)
    cv.Dilate(maskRG, maskRG)
    #remapImage(maskRG, maskRG, 20, 200)

    
    #######################################
    ## REMOVE GREEN FROM IMAGE
    ######################################
    cv.ConvertScale(maskDetail, maskDetailNorm, 1/255.)
    cv.ConvertScale(green, greenNorm, 1/255.)
    cv.Mul(greenNorm, maskDetailNorm, greenNorm, greenRemoval)
    cv.ConvertScale(greenNorm, green, 255)
    
    #######################################
    ## COMBINE THE MASKS
    ######################################
    cv.Add(maskRG, maskDetail, mask)
    remapImage(mask, mask, 60., 170.)
    
    #######################################
    ## SOFTEN EDGES
    ######################################
    cv.Canny(mask, edge, 50, 200)
    #cv.Sobel(mask, edgeF, 1, 1)
    #cv.ConvertScale(edgeF, edge, 255)
    
    cv.Smooth(edge, edge, cv.CV_BLUR, 5, 5)
    remapImage(edge, edge, 20., 60.)
    
    
    #######################################
    ## CREATE COMPOSITION
    ######################################
    
    cv.Merge(red, green, blue, mask, composition)
    
    

    
    #######################################
    ## SHOW IMAGES
    ######################################
    cv.ShowImage(inputWin, frame)
    
    if SHOW_RGB:
        cv.ShowImage(redWin, red)
        cv.ShowImage(greenWin, greenInverted)
        cv.ShowImage(blueWin, blue)
        
        cv.ShowImage(redSubWin, redSub)
        cv.ShowImage(greenSubWin, greenSub)
        cv.ShowImage(blueSubWin, blueSub)
    
    #show image with pygame
    if USE_PYGAME:
        display.blit(checkers, (0,0))
        #cv.CvtColor(composition, compositionPygame, cv.CV_BGR2RGB)
        #cv.Set(sub, 255)
        cv.Merge(blue, green, red, mask, compositionPygame)
        comp = pygame.image.frombuffer(compositionPygame.tostring(), (640, 480), 'RGBA')
        comp = pygame.transform.scale(comp, (1280, 960))
        display.blit(comp, (0,0))
        pygame.display.flip()
    else:
        cv.ShowImage(mainWin, composition)

    cv.ShowImage(maskRGWin, maskRG)
    cv.ShowImage(maskDetailWin, maskDetail)
    cv.ShowImage(maskWin, mask)
    
    cv.ShowImage(edgeWin, edge)
    
    if USE_PYGAME:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                break
            elif event.type == pygame.KEYDOWN and event.key == pygame.K_ESCAPE:
                pygame.quit()
                break
            elif event.type == pygame.KEYDOWN and pygame.K_SPACE:
                cv.SaveImage( os.path.dirname(os.path.realpath(__file__))+"/output.png", composition)
            elif event.type == pygame.MOUSEBUTTONUP:
                x, y = pygame.mouse.get_pos()
                #updateBg(x,y)
    
    #INPUT METHOD FOR PURE OPENCV
    k = cv.WaitKey(10)
    if k == 32:
        cv.SaveImage( os.path.dirname(os.path.realpath(__file__))+"/output.png", composition)
    if k == 27: #exit on escape
        pygame.quit()
        break