class Classify(object):
"""
Class for portrait segmentation
"""
def __init__(self, model_path, model_width, model_height):
self._model_path = model_path
self._model_width = model_width
self._model_height = model_height
self._img_width = 0
self._img_height = 0
self._model = None
self._dvpp = None
def init(self):
"""
Initialize
"""
self._dvpp = Dvpp()
# Load model
self._model = Model(self._model_path)
return const.SUCCESS
@utils.display_time
def pre_process(self, image):
"""
preprocess
"""
image_dvpp = image.copy_to_dvpp()
yuv_image = self._dvpp.jpegd(image_dvpp)
resized_image = self._dvpp.resize(yuv_image, self._model_width,
self._model_height)
return resized_image
@utils.display_time
def inference(self, input_data):
"""
model inference
"""
return self._model.execute(input_data)
@utils.display_time
def post_process(self, infer_output, image_file):
"""
Post-processing, analysis of inference results
"""
output_path = os.path.join(OUTPUT_DIR, os.path.basename(image_file))
infer_result = infer_output[0]
vals = infer_result.flatten()
pre_index = vals.argsort()[-1]
origin_img = Image.open(image_file)
draw = ImageDraw.Draw(origin_img)
font = ImageFont.load_default()
draw.text((10, 50), CLS[pre_index], font=font, fill=255)
origin_img.save(output_path)
class Cartoonization(object):
"""
class for Cartoonization
"""
def __init__(self, model_path, model_width, model_height):
self._model_path = model_path
self._model_width = model_width
self._model_height = model_height
self.device_id = 0
self._dvpp = None
self._model = None
def init(self):
"""
Initialize
"""
# Initialize dvpp
self._dvpp = Dvpp()
# Load model
self._model = Model(self._model_path)
return const.SUCCESS
@utils.display_time
def pre_process(self, image):
"""
image preprocess
"""
image_dvpp = image.copy_to_dvpp()
yuv_image = self._dvpp.jpegd(image_dvpp)
crop_and_paste_image = self._dvpp.crop_and_paste_get_roi(yuv_image, image.width, image.height, \
self._model_width, self._model_height)
return crop_and_paste_image
@utils.display_time
def inference(self, resized_image):
"""
model inference
"""
return self._model.execute(resized_image)
@utils.display_time
def post_process(self, infer_output, image_file, origin_image):
"""
post process
"""
data = ((np.squeeze(infer_output[0]) + 1) * 127.5)
img = cv2.cvtColor(data, cv2.COLOR_RGB2BGR)
img = cv2.resize(img, (origin_image.width, origin_image.height))
output_path = os.path.join("../outputs", os.path.basename(image_file))
cv2.imwrite(output_path, img)
class Gesture(object):
"""
define gesture class
"""
def __init__(self, model_path, model_width, model_height):
self.device_id = 0
self.context = None
self.stream = None
self._model_path = model_path
self._model_width = model_width
self._model_height = model_height
self._dvpp = None
self._model = None
def init(self):
"""
Initialize
"""
# Initialize dvpp
self._dvpp = Dvpp()
# Load model
self._model = Model(self._model_path)
return const.SUCCESS
def pre_process(self, image):
"""
pre_precess
"""
image_dvpp = image.copy_to_dvpp()
yuv_image = self._dvpp.jpegd(image_dvpp)
print("decode jpeg end")
resized_image = self._dvpp.resize(yuv_image, self._model_width,
self._model_height)
print("resize yuv end")
return resized_image
def inference(self, resized_image):
"""
inference
"""
return self._model.execute(resized_image)
def post_process(self, infer_output, image_file):
"""
post_process
"""
print("post process")
data = infer_output[0]
vals = data.flatten()
top_k = vals.argsort()[-1:-2:-1]
print("images:{}".format(image_file))
print("======== top5 inference results: =============")
for n in top_k:
object_class = get_gesture_categories(n)
print("label:%d confidence: %f, class: %s" %
(n, vals[n], object_class))
if len(top_k):
object_class = get_gesture_categories(top_k[0])
output_path = os.path.join(os.path.join(SRC_PATH, "../outputs"),
os.path.basename(image_file))
origin_img = Image.open(image_file)
draw = ImageDraw.Draw(origin_img)
font = ImageFont.load_default()
draw.text((10, 50), object_class, font=font, fill=255)
origin_img.save(output_path)
class CrowdCount(object):
"""
crowdCount
"""
def __init__(self, model_path, model_width, model_height):
self.device_id = 0
self.context = None
self.stream = None
self._model = None
self.run_mode = None
self._model_path = model_path
self._model_width = model_width
self._model_height = model_height
self._dvpp = None
self._model = None
def init(self):
"""
init
"""
self._dvpp = Dvpp()
self._model = Model(self._model_path)
return constants.SUCCESS
def pre_process(self, image):
"""
image preprocess
"""
image_dvpp = image.copy_to_dvpp()
yuv_image = self._dvpp.jpegd(image_dvpp)
crop_and_paste_image = \
self._dvpp.crop_and_paste(yuv_image, image.width, image.height, self._model_width, self._model_height)
return crop_and_paste_image
def inference(self, input_data):
"""
model inference
"""
return self._model.execute(input_data)
def post_process(self, infer_output, image_file):
"""
Post-processing, analysis of inference results
"""
orig = cv2.imread(image_file, 1)
orig = cv2.resize(orig, (1200, 800))
data = infer_output[0]
vals = data.flatten()
res = np.sum(vals, axis=0)
result = round(res / 1000.0)
data_2 = data.reshape(800, 1408)
heatMap = data_2[:800, :1200]
heatMap = heatMap.astype(np.uint8)
heatMap = cv2.GaussianBlur(heatMap, (5, 5), cv2.BORDER_DEFAULT)
cv2.normalize(heatMap, heatMap, 0, 255, cv2.NORM_MINMAX, cv2.CV_8UC1)
heatMap = cv2.applyColorMap(heatMap, cv2.COLORMAP_JET)
add_img = cv2.addWeighted(orig, 1, heatMap, 0.5, 0.0)
cv2.putText(add_img, str(result), (30, 60), cv2.FONT_HERSHEY_PLAIN, 5, (0, 0, 255), 8)
output_path = os.path.join("./outputs", os.path.basename(image_file))
cv2.imwrite(output_path, add_img)
