116 lines
3.2 KiB
Python
116 lines
3.2 KiB
Python
import os
|
|
from os.path import isfile, join
|
|
|
|
import imutils as imutils
|
|
import numpy as np
|
|
import cv2
|
|
import constant
|
|
import mimetypes
|
|
|
|
try:
|
|
from cv2 import cv2
|
|
finally:
|
|
pass
|
|
|
|
proto_file = "./model/colorization_deploy_v2.prototxt"
|
|
model = "./model/colorization_release_v2.caffemodel"
|
|
points_file = "./model/pts_in_hull.npy"
|
|
network = cv2.dnn.readNetFromCaffe(proto_file, model)
|
|
points = np.load(points_file)
|
|
|
|
points = points.transpose().reshape(2, 313, 1, 1)
|
|
|
|
class_layer = network.getLayerId("class8_ab")
|
|
convolution_layer = network.getLayerId("conv8_313_rh")
|
|
|
|
network.getLayer(class_layer).blobs = [points.astype(constant.FLOAT32)]
|
|
network.getLayer(convolution_layer).blobs = [np.full([1, 313], 2.606, dtype=constant.FLOAT32)]
|
|
|
|
|
|
# Convert from grayscale to gray to remove all little noises, and then again to rgb
|
|
def clear_image(image):
|
|
cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
|
cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
|
|
|
|
return image
|
|
|
|
|
|
def convert_frames_to_video(file_name):
|
|
path_in = 'videos/colorized_frames/'
|
|
path_out = 'videos/output/' + file_name.replace('/', '_')
|
|
frame_array = []
|
|
files = [f for f in os.listdir(path_in) if isfile(join(path_in, f))]
|
|
files.sort(key=lambda x: int(x[5:-4]))
|
|
|
|
for i in range(len(files)):
|
|
filename = path_in + files[i]
|
|
img = cv2.imread(filename)
|
|
height, width, layers = img.shape
|
|
size = (width, height)
|
|
frame_array.append(img)
|
|
|
|
out = cv2.VideoWriter(path_out, cv2.VideoWriter_fourcc(*'MJPG'), 30.0, size)
|
|
|
|
for i in range(len(frame_array)):
|
|
out.write(frame_array[i])
|
|
|
|
out.release()
|
|
|
|
|
|
def colorize_image(image):
|
|
image_scaled = image.astype(constant.FLOAT32) / 255.0
|
|
lab = cv2.cvtColor(image_scaled, cv2.COLOR_RGB2LAB)
|
|
resized = cv2.resize(lab, (224, 224))
|
|
image_l_channel = cv2.split(resized)[0] - 50
|
|
|
|
network.setInput(cv2.dnn.blobFromImage(image_l_channel))
|
|
ab = network.forward()[0, :, :, :].transpose((1, 2, 0))
|
|
ab = cv2.resize(ab, (image.shape[1], image.shape[0]))
|
|
|
|
image_l_channel = cv2.split(lab)[0]
|
|
colorized = np.concatenate((image_l_channel[:, :, np.newaxis], ab), axis=2)
|
|
|
|
colorized = cv2.cvtColor(colorized, cv2.COLOR_LAB2RGB)
|
|
colorized = np.clip(colorized, 0, 1)
|
|
|
|
return (255 * colorized).astype("uint8")
|
|
|
|
|
|
def colorize(file_path: str) -> None:
|
|
file_type = mimetypes.guess_type(file_path)[0].split('/')
|
|
|
|
if file_type[0] == 'image':
|
|
image = cv2.imread(file_path)
|
|
colorized = colorize_image(image)
|
|
|
|
cv2.imwrite("images/output/" + file_path.replace('/', '_'), cv2.cvtColor(colorized, cv2.COLOR_RGB2BGR))
|
|
|
|
print('Success')
|
|
elif file_type[0] == 'video':
|
|
vs = cv2.VideoCapture(file_path)
|
|
count = 0
|
|
success = True
|
|
|
|
while success:
|
|
success, frame = vs.read()
|
|
|
|
if frame is None:
|
|
break
|
|
|
|
frame = imutils.resize(frame, 500)
|
|
|
|
colorized = colorize_image(frame)
|
|
|
|
cv2.imwrite("./videos/colorized_frames/frame%d.jpg" % count, colorized)
|
|
count += 1
|
|
key = cv2.waitKey(1) & 0xFF
|
|
if key == ord("q"):
|
|
break
|
|
|
|
vs.release()
|
|
convert_frames_to_video(file_path)
|
|
|
|
print('Success')
|
|
else:
|
|
print('Wrong media format')
|