54 lines
1.2 KiB
Plaintext
54 lines
1.2 KiB
Plaintext
|
#!/usr/bin/env python3
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import itertools, functools
|
||
|
|
import array
|
||
|
|
import os
|
||
|
|
import struct
|
||
|
|
import time
|
||
|
|
import torch
|
||
|
|
from colorizers import eccv16
|
||
|
|
import numpy as np
|
||
|
|
from PIL import Image
|
||
|
|
from skimage import color
|
||
|
|
|
||
|
|
def run(args):
|
||
|
|
model = eccv16(pretrained=True).eval()
|
||
|
|
|
||
|
|
img_path = args.src
|
||
|
|
print(f'Processing {img_path}...')
|
||
|
|
img = np.asarray(Image.open(img_path))
|
||
|
|
if img.ndim == 2:
|
||
|
|
img = np.tile(img[:,:,None], 3)
|
||
|
|
img = img[:,:,:3]
|
||
|
|
img = np.asarray(Image.fromarray(img).resize((256, 256), resample=3))
|
||
|
|
img = color.rgb2lab(img)[:,:,0]
|
||
|
|
|
||
|
|
N = 100
|
||
|
|
iters = 100
|
||
|
|
|
||
|
|
img = np.tile(img, (N, 1, 1, 1))
|
||
|
|
print(img.shape)
|
||
|
|
tin = torch.Tensor(img)
|
||
|
|
|
||
|
|
model = model.cuda()
|
||
|
|
tin = tin.cuda()
|
||
|
|
|
||
|
|
tsum = 0
|
||
|
|
for i in range(iters):
|
||
|
|
st = time.time()
|
||
|
|
ref = model(tin)
|
||
|
|
et = time.time()
|
||
|
|
print(f'iter {i} time {et - st}')
|
||
|
|
tsum += et - st
|
||
|
|
print(f'avg {tsum / iters} throughput {(iters * N) / tsum}')
|
||
|
|
|
||
|
|
def main():
|
||
|
|
parser = argparse.ArgumentParser()
|
||
|
|
parser.add_argument('src', help='Input image.')
|
||
|
|
args = parser.parse_args()
|
||
|
|
run(args)
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
main()
|