Automatic Rotoscoping…For Free 1
Support ProVideo Coalition
Shop with Filmtools Logo

What Do You Think? Let Us Know.

Subscribe
Notify of
guest
1 Comment
Oldest
Newest Most Voted
Inline Feedbacks
View all comments
Randy Curtiss

I found a way to get this running with CUDA:

Here’s an altered FindMattes.py. (If you want to get the colors as they originally were, replace the decode_segmap function with the original…)

from PIL import Image
import torch
import torchvision.transforms as T
from torchvision import models
import numpy as np
device = torch.device(“cuda” if torch.cuda.is_available() else “cpu”)
torch.set_default_tensor_type(‘torch.cuda.FloatTensor’)
fcn = None

def getRotoModel():
global fcn
fcn = models.segmentation.fcn_resnet101(pretrained=True).eval()

# Define the helper function
def decode_segmap(image, nc=21):

label_colors = np.array([(0, 0, 0), # 0=background
# 1=aeroplane, 2=bicycle, 3=bird, 4=boat, 5=bottle
(0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0),
# 6=bus, 7=car, 8=cat, 9=chair, 10=cow
(0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0),
# 11=dining table, 12=dog, 13=horse, 14=motorbike, 15=person
(0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (255, 255, 255),
# 16=potted plant, 17=sheep, 18=sofa, 19=train, 20=tv/monitor
(0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0), (0, 0, 0)])

r = np.zeros_like(image).astype(np.uint8)
g = np.zeros_like(image).astype(np.uint8)
b = np.zeros_like(image).astype(np.uint8)

for l in range(0, nc):
idx = image == l
r[idx] = label_colors[l, 0]
g[idx] = label_colors[l, 1]
b[idx] = label_colors[l, 2]

rgb = np.stack([r, g, b], axis=2)
return rgb

def createMatte(filename, matteName, size):
img = Image.open(filename).convert(‘RGB’)
trf = T.Compose([T.Resize(size),
T.ToTensor(),
T.Normalize(mean = [0.485, 0.456, 0.406],
std = [0.229, 0.224, 0.225])])
inp = trf(img).unsqueeze(0)
if (fcn == None): getRotoModel()
out = fcn(inp.cuda())[‘out’]
om = torch.argmax(out.squeeze(), dim=0).detach().cpu().numpy()
rgb = decode_segmap(om)
im = Image.fromarray(rgb)
im.save(matteName)

Last edited 8 days ago by Randy Curtiss

You Might Also Like