class Worker(threading.Thread):
def __init__(self,name,upstream_task_queue,downstream_task_queue):
super(Worker, self).__init__()
self._name = name
self._upstream_task_queue = upstream_task_queue
self._downstream_task_queue = downstream_task_queue
self._stop_flag = False
caffe.set_mode_cpu()
logging.info("Running net on CPU.")
# initialize style transfer
self._style_transfer = StyleTransfer('googlenet')
def _create_art(self,style_image,content_image,params):
# perform style transfer
start = timeit.default_timer()
n_iters = self._style_transfer.transfer_style(style_image, content_image, length=params["length"], init="mixed", ratio=np.float(params["ratio"]), n_iter=params["num_iters"], verbose=params["verbose"])
end = timeit.default_timer()
logging.info("Ran {0} iterations in {1:.0f}s.".format(n_iters, end-start))
img_out = self._style_transfer.get_generated()
img_out = img_as_ubyte(img_out)
return img_out;
def stop(self):
self._stop_flag = True
def run(self):
'''
'''
logging.info("Runing worker ("+self._name+")...")
while self._stop_flag == False:
try:
task = self._upstream_task_queue.get(False)
# validate task
if self._validate_task(task) == False:
self._upstream_task_queue.task_done()
continue
# create art
task.output_image = self._create_art(task.style_image, task.content_image, task.params)
#send event to downstream
self._downstream_task_queue.put(task)
self._upstream_task_queue.task_done()
except Empty:
logging.info('not task request on local queue, retry in 5 seconds later.')
time.sleep(5)
logging.info("Completed worker ("+self._name+").")
def _validate_task(self,task):
return True
# -*- coding: utf-8 -*-
"""
Created on Thu Nov 24 16:52:41 2016
@author: zhouc
"""
# library imports
import caffe
import cv2
import numpy as np
# local imports
from style import StyleTransfer
caffe.set_mode_gpu()
style_img_path = 'images/style/the_scream.jpg'
content_img_path = 'images/content/tubingen.jpg'
img_style = caffe.io.load_image(style_img_path)
img_content = caffe.io.load_image(content_img_path)
args = {"length": 600, "ratio": 2e5, "n_iter": 32, "init": "content"}
st = StyleTransfer()
st.transfer_style(img_style, img_content, **args)
img_out = st.get_generated()
# show the image
cv2.imshow("Style", cv2.cvtColor(img_out, cv2.COLOR_RGB2BGR))
cv2.waitKey()
cv2.destroyWindow("Style")
