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)