croma_key.py

import cv2 
import numpy as np 
import pandas as pd
from os import system,name
from tqdm import tqdm
import time

import segment

def click_and_crop(event, x, y, flags, param):

    global refPt, cropping,cropped

    if cropped==True:
        return    
    if event == cv2.EVENT_LBUTTONDOWN:
        refPt = [(x, y)]
        cropping = True
    
    elif event == cv2.EVENT_LBUTTONUP:
        refPt.append((x, y))
        cropping = False
        cv2.rectangle(frame, refPt[0], refPt[1], (0, 255, 0), 2)    
        cv2.imshow("video", frame)

def change_background(mask,image,background):

    image=cv2.cvtColor(image, cv2.COLOR_HSV2RGB)
    background = cv2.resize(background,(1920,1080))    
       
    image[mask==255]=0    
    image=resize_with_aspect_ratio(image)    
    foreground=np.zeros((1080,1920,3),np.uint8)
    foreground[1080-image.shape[0]:,100:image.shape[1]+100]=image
    # cv2.namedWindow('foreground', cv2.WINDOW_NORMAL)
    # cv2.resizeWindow('foreground', 1920//2, 1080//2)
    # cv2.imshow('foreground',foreground)

    mask=resize_with_aspect_ratio(mask)
    HD_mask=np.ones((1080,1920),np.uint8)*255
    HD_mask[1080-image.shape[0]:,100:image.shape[1]+100]=mask
    # cv2.namedWindow('HD_mask', cv2.WINDOW_NORMAL)
    # cv2.resizeWindow('HD_mask', 1920//2, 1080//2)
    # cv2.imshow('HD_mask',HD_mask)

    background[HD_mask==0]=0    
    combined=foreground+background

    return combined

def find_green_range(greens):
    hand_flattened = []
    l=len(greens.flatten())//3   
    greens=greens.reshape((l,3))

    for index in range(3):
        bg=greens[:,index].flatten()
        bg=bg[bg != 0]
        top_and_bottom_series = pd.Series(bg)

        hand_flattened.append(top_and_bottom_series)

    low = []
    high = []

    def restrict(color_component):
        return np.clip(color_component, 0, 255)

    z_value = 5.5

    for i in range(3):
        mu = hand_flattened[i].values.mean()
        sigma = hand_flattened[i].values.std()
        deviation = z_value*sigma
        low.append(restrict(mu-deviation))
        high.append(restrict(mu+deviation))

    mask_lower = np.array([low[0], low[1], low[2]])
    mask_higher = np.array([high[0], high[1], high[2]])

    return mask_lower,mask_higher

def get_green_values(mask,frame_copy):
    image=np.copy(frame_copy)
    image[mask==255]=0
    return image

def find_segment_mask(frame_copy):

    seg_img=segment.run_visualization(frame_copy)
    seg_img = cv2.cvtColor(seg_img,cv2.COLOR_BGR2GRAY)
    ret, seg_img = cv2.threshold(seg_img, 10, 255, cv2.THRESH_BINARY)

    seg_img=cv2.resize(seg_img,(frame_copy.shape[1],frame_copy.shape[0]))
    kernel = np.ones((5,5),np.uint8)
    seg_img = cv2.erode(seg_img,kernel,iterations = 1)

    return seg_img

def tune_mask(mask):
    kernel = np.ones((3,3),np.uint8)
    mask=cv2.dilate(mask,kernel,iterations = 2)
    mask = cv2.medianBlur(mask,3)

    return mask

def resize_with_aspect_ratio(image):
    (h, w) = image.shape[:2]
    r = 1080 / float(h)
    dim = (int(w * r), 1080)
    image=cv2.resize(image, dim, interpolation = cv2.INTER_AREA)
    return image


refPt = []
cropping = False
cropped=False
got_color=False

input_video_path="input/agni_puranam_1.mp4"
cap  = cv2.VideoCapture(input_video_path)
background = cv2.imread("input/bg.jpeg")

cv2.namedWindow("video",cv2.WINDOW_NORMAL)
cv2.resizeWindow('video', 1920//2, 1080//2)

cv2.namedWindow("output",cv2.WINDOW_NORMAL)
cv2.resizeWindow('output', 1920//2, 1080//2)


cv2.setMouseCallback("video", click_and_crop)

video_length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
# fourcc = cv2.VideoWriter_fourcc(*'MP4V')
frame_width = 1920
frame_height = 1080 
size = (frame_width, frame_height) 
fps = cap.get(cv2.CAP_PROP_FPS)
save_path='output/out_'+input_video_path.split('/')[1]
# out = cv2.VideoWriter(save_path, fourcc, fps, size)
out = cv2.VideoWriter(save_path,cv2.VideoWriter_fourcc('M','J','P','G'), fps, size)
# out = cv2.VideoWriter('outpy2.mp4',cv2.VideoWriter_fourcc('M','J','P','G'), 10, (frame_width,frame_height))
# print(length,fps)


# while True:
for i in tqdm(range(500)):

    s1=time.time()
    ret, frame = cap.read() 
    e1=time.time()
    print('io time'+str(s1-e1))

    s2=time.time()

    frame=cv2.cvtColor(frame, cv2.COLOR_RGB2HSV)
    frame_copy=np.copy(frame)

    if cropped==False:
        cv2.imshow('video',frame)
        key=cv2.waitKey(0)
        if key == ord("r"):
            cropped=False
            continue
        elif key == ord("k"):
            cropped=True

    frame_copy = frame_copy[refPt[0][1]:refPt[1][1], refPt[0][0]:refPt[1][0]]    
    
    if got_color==False:        
        seg_mask=find_segment_mask(frame_copy)                
        greens=get_green_values(seg_mask,frame_copy)          
        mask_lower,mask_higher=find_green_range(greens)              
        got_color=True  
          
    
            
    mask= cv2.inRange(frame_copy, mask_lower, mask_higher)
    mask=tune_mask(mask)
    output=change_background(mask,frame_copy,background) 

    # cv2.imshow("output",output) 
    out.write(output)
    # cv2.imwrite('frame_'+str(i)+'.jpg',output)

    # if cv2.waitKey(25) == 1:
    #     break 
    e2=time.time()
    print('pro time'+str(s2-e2))

cap.release()
out.release()
cv2.destroyAllWindows()