示例#1
0
def predict_mask(module, input_img, id, batch_size, extra=None, r_img=True):
    output_folder = "detection_result"
    global use_gpu
    method_name = module.desc.attr.map.data['default_signature'].s
    predict_method = getattr(module, method_name)
    try:
        data = {}
        if input_img is not None:
            input_img = {"image": input_img}
            data.update(input_img)
        if extra is not None:
            data.update(extra)
            r_img = False if "visual_result" in extra.keys() else True
        results = predict_method(data=data,
                                 visualization=r_img,
                                 use_gpu=use_gpu,
                                 batch_size=batch_size)
        results = utils.handle_mask_results(results)
    except Exception as err:
        curr = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(time.time()))
        print(curr, " - ", err)
        return {"result": "Please check data format!"}
    finally:
        base64_list = []
        results_pack = []
        if input_img is not None:
            if r_img is False:
                for index in range(len(results)):
                    results[index]["path"] = ""
                results_pack = results
            else:
                input_img = input_img.get("image", [])
                for index in range(len(input_img)):
                    item = input_img[index]
                    with open(os.path.join(output_folder, item), "rb") as fp:
                        b_head = "data:image/" + item.split(
                            ".")[-1] + ";base64"
                        b_body = base64.b64encode(fp.read())
                        b_body = str(b_body).replace("b'", "").replace("'", "")
                        b_img = b_head + "," + b_body
                        base64_list.append(b_img)
                        results[index]["path"] = results[index]["path"].replace(
                            id + "_", "") if results[index]["path"] != "" \
                            else ""

                        results[index].update({"base64": b_img})
                        results_pack.append(results[index])
                    os.remove(item)
                    os.remove(os.path.join(output_folder, item))
        else:
            results_pack = results

    return results_pack
示例#2
0
def predict_mask(module_info, input_img, id, extra=None, r_img=False):
    output_folder = "detection_result"
    method_name = module_info["method_name"]
    module = module_info["module"]
    predict_method = getattr(module, method_name)
    data_len = len(input_img) if input_img is not None else 0
    data = {}
    if input_img is not None:
        input_img = {"image": input_img}
        data.update(input_img)
    if extra is not None:
        data.update(extra)
        r_img = True if "visual_result" in extra.keys() else False

    predict_args = module_info["predict_args"].copy()
    predict_args.update({"data": data})
    if isinstance(predict_method, functools.partial):
        predict_method = predict_method.func
        predict_args.update({"sign_name": method_name})
    for item in predict_method.__code__.co_varnames:
        if item in module_info.keys():
            predict_args.update({item: module_info[item]})
    try:
        results = predict_method(**predict_args)
        results = utils.handle_mask_results(results, data_len)
    except Exception as err:
        curr = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(time.time()))
        print(curr, " - ", err)
        return gen_result("-1", "Please check data format!", "")
    finally:
        base64_list = []
        results_pack = []
        if input_img is not None:
            if r_img is False:
                for index in range(len(results)):
                    results[index]["path"] = ""
                results_pack = results
                str_id = id + "*"
                files_deleted = glob.glob(str_id)
                for path in files_deleted:
                    if os.path.exists(path):
                        os.remove(path)
            else:
                input_img = input_img.get("image", [])
                for index in range(len(input_img)):
                    item = input_img[index]
                    file_path = os.path.join(output_folder, item)
                    if not os.path.exists(file_path):
                        results_pack.append(results[index])
                        os.remove(item)
                    else:
                        with open(file_path, "rb") as fp:
                            b_head = "data:image/" + item.split(
                                ".")[-1] + ";base64"
                            b_body = base64.b64encode(fp.read())
                            b_body = str(b_body).replace("b'", "").replace(
                                "'", "")
                            b_img = b_head + "," + b_body
                            base64_list.append(b_img)
                            results[index]["path"] = results[index]["path"].replace(
                                id + "_", "") if results[index]["path"] != "" \
                                else ""

                            results[index].update({"base64": b_img})
                            results_pack.append(results[index])
                        os.remove(item)
                        os.remove(os.path.join(output_folder, item))
        else:
            results_pack = results

    return gen_result("0", "", str(results_pack))