|
|
import cv2 |
|
|
import torch |
|
|
import imageio |
|
|
from methods.img2pixl import pixL |
|
|
|
|
|
|
|
|
device = "cuda" if torch.cuda.is_available() else "cpu" |
|
|
face2paint = torch.hub.load("bryandlee/animegan2-pytorch:main", "face2paint", device=device, size=512) |
|
|
model = torch.hub.load("bryandlee/animegan2-pytorch", "generator", device=device).eval() |
|
|
|
|
|
class Media: |
|
|
|
|
|
def __init__(self,fname = None,pixel_size = None): |
|
|
self.fname = fname |
|
|
self.pixel_size = pixel_size |
|
|
|
|
|
def split(self,fname,pixel_size, mediaType): |
|
|
media = cv2.VideoCapture(fname) |
|
|
frames = [] |
|
|
while True: |
|
|
ret, cv2Image = media.read() |
|
|
if not ret: |
|
|
break |
|
|
frames.append(cv2Image) |
|
|
frames = pixL().toThePixL(frames, pixel_size) |
|
|
if mediaType == 'gif': |
|
|
imageio.mimsave('cache.gif', frames) |
|
|
return [None, 'cache.gif', 'cache.gif'] |
|
|
else: |
|
|
output_file = "cache.mp4" |
|
|
out = cv2.VideoWriter(output_file,cv2.VideoWriter_fourcc(*'h264'), 15, (frames[0].shape[1],frames[0].shape[0])) |
|
|
for i in range(len(frames)): |
|
|
out.write(frames[i]) |
|
|
out.release() |
|
|
return [output_file, None, output_file] |
