holmeshoo/beans_sorting

1import gradio as gr2import prediction as prd3import multiple_prdiction as mprd4import pandas as pd5import cv26import numpy as np7import copy8 9model_label = []10model_label.append(prd.loadModelAndLabel("./model/white_beans.h5", "./model/white_beans.json"))11model_label.append(prd.loadModelAndLabel("./model/red_beans.h5", "./model/red_beans.json"))12model_label.append(prd.loadModelAndLabel("./model/black_beans.h5", "./model/black_beans.json"))13print("model loaded")14beans_variety = [15    "white",16    "red",17    "black"18]19 20 21def my_find(l, x):22    if x in l:23        return l.index(x)24    else:25        return -126 27 28def areaTest(file_name, block_size=33, area_size=200):29    if block_size % 2 == 0:30        block_size += 131    img = cv2.imread(file_name)32    point = mprd.getArea(img, block_size=block_size, area_size=area_size)33 34    img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)35    img = cv2.adaptiveThreshold(img, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, block_size, 6)36    img = 255 - img37    img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)38    img = mprd.printArea(img, point)39    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)40    img = np.array(img)41    return img42 43 44def multiPredeiction(file_name, block_size, area_size, n):45    (model, label) = model_label[n]46    img = cv2.imread(file_name)47    point = mprd.getArea(img, block_size, area_size)48    result = mprd.multiPredict(model=model, label=label, img=img, point=point)49    img = mprd.printResultImg(img, point, result)50    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)51    img = np.array(img)52    sum = {}53    for key in label:54        sum[key] = 055    for i in result:56        max_label = ""57        max_val = 058        for n_label, n_val in i.items():59            if (max_val < n_val):60                max_val = n_val61                max_label = n_label62        sum[max_label] += 163    sum_for_txt = copy.deepcopy(sum)64    for l, v in sum.items():65        sum[l] /= len(result)66    keys = "|"67    values = "|"68    for k, v in sum_for_txt.items():69        keys += str(k).upper() + "|"70        values += str(v) + "|"71    table_text = keys + "\n" + "|:-:" * len(label) + "|\n" + values + "\n"72    # print(table_text)73    return img, sum, table_text74 75 76def predictFile(file_name, n):77    (model, label) = model_label[n]78    result = prd.predictFile(model, label, file_name)79 80    keys = "|"81    values = "|"82 83    for k, v in result.items():84        keys += str(k).upper() + "|"85        values += format(v, '.2f') + "|"86    table_text = keys + "\n" + "|:-:" * len(label) + "|\n" + values + "\n"87 88    return (gr.Markdown(table_text), result)89 90 91with gr.Blocks(theme=gr.themes.Soft(primary_hue="orange")) as demo:92    gr.Markdown("# Beans Sorting")93    with gr.Tab("single detection"):94        #  デフォルトはwhite95        dropdown = gr.Dropdown(beans_variety, label="beans variety", value=beans_variety[0])96        picture = gr.Image(type='filepath')97        model_number = gr.Number(value=0, visible=False)98 99        @dropdown.change(inputs=dropdown, outputs=model_number)100        def changeModel(choice):101            return gr.Number(value=my_find(beans_variety, choice))102 103        button = gr.Button("Run")104 105        with gr.Row():106            text = gr.Markdown("")107            label = gr.Label()108 109        button.click(fn=predictFile, inputs=[picture, model_number], outputs=[text, label])110    with gr.Tab("Multiple detection"):111        dropdown = gr.Dropdown(beans_variety, label="beans variety", value=beans_variety[0])112        picture = gr.Image(type='filepath')113        #  デフォルトはwhite114        model_number = gr.Number(value=0, visible=False)115 116        block_size_slider = gr.Slider(value=45, minimum=1, step=2, maximum=100, label="block_size(Perform convolution with this value)")117        area_size_slider = gr.Slider(value=2000, minimum=0, step=1, maximum=10000, label="area size slider(If the number of pixels inside the circle is less than this, it will not be detected.)")118 119        @dropdown.change(inputs=dropdown, outputs=model_number)120        def changeModel(choice):121            return gr.Number(value=my_find(beans_variety, choice))122        with gr.Row():123            check_button = gr.Button("Check")124            run_button = gr.Button("Run")125        with gr.Row():126            text = gr.Markdown("")127            label = gr.Label("")128        output = gr.Image(type="numpy")129 130        check_button.click(fn=areaTest, inputs=[picture, block_size_slider, area_size_slider], outputs=[output])131        run_button.click(fn=multiPredeiction, inputs=[picture, block_size_slider, area_size_slider, model_number], outputs=[output, label, text])132 133 134if __name__ == "__main__":135    demo.launch()136