mirror of https://github.com/kritiksoman/GIMP-ML
You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
131 lines
4.2 KiB
Python
131 lines
4.2 KiB
Python
import os
|
|
|
|
baseLoc = os.path.dirname(os.path.realpath(__file__)) + '/'
|
|
savePath = '/'.join(baseLoc.split('/')[:-2]) + '/output/interpolateframes'
|
|
from gimpfu import *
|
|
import sys
|
|
|
|
sys.path.extend([baseLoc + 'gimpenv/lib/python2.7', baseLoc + 'gimpenv/lib/python2.7/site-packages',
|
|
baseLoc + 'gimpenv/lib/python2.7/site-packages/setuptools', baseLoc + 'RIFE'])
|
|
|
|
import cv2
|
|
import torch
|
|
from torch.nn import functional as F
|
|
from model import RIFE
|
|
import numpy as np
|
|
|
|
|
|
|
|
def channelData(layer): # convert gimp image to numpy
|
|
region = layer.get_pixel_rgn(0, 0, layer.width, layer.height)
|
|
pixChars = region[:, :] # Take whole layer
|
|
bpp = region.bpp
|
|
# return np.frombuffer(pixChars,dtype=np.uint8).reshape(len(pixChars)/bpp,bpp)
|
|
return np.frombuffer(pixChars, dtype=np.uint8).reshape(layer.height, layer.width, bpp)
|
|
|
|
|
|
def createResultLayer(image, name, result):
|
|
rlBytes = np.uint8(result).tobytes();
|
|
rl = gimp.Layer(image, name, image.width, image.height, 0, 100, NORMAL_MODE)
|
|
region = rl.get_pixel_rgn(0, 0, rl.width, rl.height, True)
|
|
region[:, :] = rlBytes
|
|
image.add_layer(rl, 0)
|
|
gimp.displays_flush()
|
|
|
|
|
|
def getinter(img_s, img_e, c_flag, string_path):
|
|
exp = 4
|
|
out_path = string_path
|
|
|
|
model = RIFE.Model(c_flag)
|
|
model.load_model(baseLoc + 'weights' + '/interpolateframes')
|
|
model.eval()
|
|
model.device(c_flag)
|
|
|
|
img0 = img_s
|
|
img1 = img_e
|
|
|
|
img0 = (torch.tensor(img0.transpose(2, 0, 1)) / 255.).unsqueeze(0)
|
|
img1 = (torch.tensor(img1.transpose(2, 0, 1))/ 255.).unsqueeze(0)
|
|
if torch.cuda.is_available() and not c_flag:
|
|
device = torch.device("cuda")
|
|
else:
|
|
device = torch.device("cpu")
|
|
|
|
img0 = img0.to(device)
|
|
img1 = img1.to(device)
|
|
|
|
n, c, h, w = img0.shape
|
|
ph = ((h - 1) // 32 + 1) * 32
|
|
pw = ((w - 1) // 32 + 1) * 32
|
|
padding = (0, pw - w, 0, ph - h)
|
|
img0 = F.pad(img0, padding)
|
|
img1 = F.pad(img1, padding)
|
|
|
|
img_list = [img0, img1]
|
|
idx=0
|
|
t=exp * (len(img_list) - 1)
|
|
for i in range(exp):
|
|
tmp = []
|
|
for j in range(len(img_list) - 1):
|
|
mid = model.inference(img_list[j], img_list[j + 1])
|
|
tmp.append(img_list[j])
|
|
tmp.append(mid)
|
|
idx=idx+1
|
|
gimp.progress_update(float(idx)/float(t))
|
|
gimp.displays_flush()
|
|
tmp.append(img1)
|
|
img_list = tmp
|
|
|
|
if not os.path.exists(out_path):
|
|
os.makedirs(out_path)
|
|
for i in range(len(img_list)):
|
|
cv2.imwrite(out_path + '/img{}.png'.format(i),
|
|
(img_list[i][0] * 255).byte().cpu().numpy().transpose(1, 2, 0)[:h, :w, ::-1])
|
|
|
|
|
|
|
|
|
|
def interpolateframes(imggimp, curlayer, string_path, layer_s, layer_e, c_flag):
|
|
layer_1 = channelData(layer_s)
|
|
layer_2 = channelData(layer_e)
|
|
|
|
if layer_1.shape[0] != imggimp.height or layer_1.shape[1] != imggimp.width or layer_2.shape[0] != imggimp.height or layer_2.shape[1] != imggimp.width:
|
|
pdb.gimp_message(" Do (Layer -> Layer to Image Size) for both layers and try again.")
|
|
else:
|
|
if torch.cuda.is_available() and not c_flag:
|
|
gimp.progress_init("(Using GPU) Running slomo and saving frames in "+string_path)
|
|
# device = torch.device("cuda")
|
|
else:
|
|
gimp.progress_init("(Using CPU) Running slomo and saving frames in "+string_path)
|
|
# device = torch.device("cpu")
|
|
|
|
if layer_1.shape[2] == 4: # get rid of alpha channel
|
|
layer_1 = layer_1[:, :, 0:3]
|
|
if layer_2.shape[2] == 4: # get rid of alpha channel
|
|
layer_2 = layer_2[:, :, 0:3]
|
|
getinter(layer_1, layer_2, c_flag, string_path)
|
|
# pdb.gimp_message("Saved")
|
|
|
|
|
|
register(
|
|
"interpolate-frames",
|
|
"interpolate-frames",
|
|
"Running slomo...",
|
|
"Kritik Soman",
|
|
"Your",
|
|
"2020",
|
|
"interpolate-frames...",
|
|
"*", # Alternately use RGB, RGB*, GRAY*, INDEXED etc.
|
|
[(PF_IMAGE, "image", "Input image", None),
|
|
(PF_DRAWABLE, "drawable", "Input drawable", None),
|
|
(PF_STRING, "string", "Output Location", savePath),
|
|
(PF_LAYER, "drawinglayer", "Start Frame:", None),
|
|
(PF_LAYER, "drawinglayer", "End Frame:", None),
|
|
(PF_BOOL, "fcpu", "Force CPU", False)
|
|
],
|
|
[],
|
|
interpolateframes, menu="<Image>/Layer/GIML-ML")
|
|
|
|
main()
|