File size: 3,366 Bytes
bcce310 9d67f5e bcce310 9d67f5e bcce310 9d67f5e bcce310 02348b1 bcce310 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 |
import gradio as gr
import numpy as np
import pandas as pd
from app_utils import annotate_planogram_compliance, do_sorting, xml_to_csv
from inference import run
import json
import os
from tempfile import NamedTemporaryFile
# Target names list
target_names = [
"Bottle,100PLUS ACTIVE 1.5L",
"Bottle,100PLUS ACTIVE 500ML",
"Bottle,100PLUS LEMON LIME 1.5L",
# Add all other target names here
]
# Define the function to run planogram compliance check
def planogram_compliance_check(planogram_image, master_planogram_image, annotation_file):
# Convert uploaded images to numpy arrays
planogram_img = np.array(planogram_image)
master_planogram_img = np.array(master_planogram_image)
# Perform inference on planogram image
result_list = run(
weights="base_line_best_model_exp5.pt",
source=planogram_img,
imgsz=[640, 640],
conf_thres=0.6,
iou_thres=0.6,
)
# Load annotation file and convert to DataFrame
if annotation_file is not None:
annotation_df = xml_to_csv(annotation_file)
sorted_xml_df = do_sorting(annotation_df)
else:
sorted_xml_df = None
# Run planogram compliance check
compliance_score, annotated_image = run_compliance_check(
planogram_img, master_planogram_img, sorted_xml_df, result_list
)
return compliance_score, annotated_image
def run_compliance_check(planogram_img, master_planogram_img, sorted_xml_df, result_list):
# Placeholder for actual score calculation
compliance_score = 0.0
# Placeholder for annotated image
annotated_image = planogram_img.copy()
if sorted_xml_df is not None:
annotate_df = sorted_xml_df[["xmin", "ymin", "xmax", "ymax", "line_number", "cls"]].astype(int)
else:
annotate_df = None
mask = master_table != non_null_product
m_detected_table = np.ma.masked_array(master_table, mask=mask)
m_annotated_table = np.ma.masked_array(detected_table, mask=mask)
# wrong_indexes = np.ravel_multi_index(master_table*mask != detected_table*mask, master_table.shape)
wrong_indexes = np.where(master_table != detected_table)
correct_indexes = np.where(master_table == detected_table)
# Annotate planogram compliance on the image
annotated_image = annotate_planogram_compliance(
annotated_image, annotate_df, correct_indexes, wrong_indexes, target_names
)
# Calculate compliance score
correct_matches = (np.ma.masked_equal(master_table, non_null_product) == detected_table).sum()
total_products = (master_table != non_null_product).sum()
if total_products != 0:
compliance_score = correct_matches / total_products
return compliance_score, annotated_image
# Gradio interface
planogram_check_interface = gr.Interface(
fn=planogram_compliance_check,
inputs=[
gr.inputs.Image(label="Planogram Image"),
gr.inputs.Image(label="Master Planogram Image"),
gr.inputs.Dataframe(label="Annotation File (XML)")
],
outputs=[
gr.outputs.Textbox(label="Compliance Score"),
gr.outputs.Image(label="Annotated Planogram Image"),
],
title="Planogram Compliance Checker",
description="Upload planogram image, master planogram image, and annotation file (if available) to check compliance."
)
planogram_check_interface.launch()
|