Spaces:
Runtime error
Runtime error
import gradio as gr | |
import pandas as pd | |
from ultralytics import YOLO | |
from datetime import datetime | |
categories = { | |
0: "Beans", | |
1: "Butter", | |
2: "Cheese", | |
3: "Drink", | |
4: "Jogurt", | |
5: "Milk", | |
6: "Pasta", | |
7: "Rice", | |
8: "Snacks", | |
9: "Tuna" | |
} | |
# inventory_separate_entries = pd.DataFrame(columns=["Product", "Time added"]) | |
inventory_separate_entries = pd.read_excel("file.xlsx", sheet_name='History of entries') | |
print(inventory_separate_entries) | |
model = YOLO("../object_detection/runs/detect/train7/weights/best.pt") | |
def get_value_counts_df(inventory_separate_entries): | |
value_counts = inventory_separate_entries['Product'].value_counts(dropna=True, sort=True) | |
df_val_counts = pd.DataFrame(value_counts) | |
df_value_counts_reset = df_val_counts.reset_index() | |
df_value_counts_reset.columns = ['Product', 'Quantity'] | |
return df_value_counts_reset | |
def add_item(item_name): | |
inventory_separate_entries.loc[-1] = [item_name, datetime.today().strftime('%Y-%m-%d %H:%M:%S')] | |
inventory_separate_entries.index = inventory_separate_entries.index + 1 # shifting index | |
inventory_separate_entries.sort_index(inplace=True) | |
def detect_objects(input_img): | |
result_img = model([input_img]) | |
print("result", result_img[0].boxes) | |
detected_categories = [] | |
for r in result_img[0].boxes: | |
for element in r.data: | |
detected_category = int(element[-1]) | |
confidence = element[-2] | |
print(r.data) | |
detected_categories.append((detected_category, confidence)) | |
print(detected_categories) | |
sorted_detected_categories = sorted(detected_categories, key=lambda x: x[1], reverse=True) | |
if len(sorted_detected_categories) > 0: | |
detected_category = categories[sorted_detected_categories[0][0]] | |
else: | |
detected_category = "Nothing detected" | |
label = gr.Label(detected_category, visible=True) | |
result_img = result_img[0].plot() | |
return gr.Image(result_img, visible=True), label#, add_item_btn # detection_interface | |
def refresh_inventory(): | |
value_counts_df = get_value_counts_df(inventory_separate_entries) | |
with pd.ExcelWriter("file.xlsx") as writer: | |
value_counts_df.to_excel(writer, sheet_name='Total inventory', index=False) | |
inventory_separate_entries.to_excel(writer, sheet_name='History of entries', index=False) | |
return value_counts_df, inventory_separate_entries | |
with gr.Blocks() as inventory_interface: | |
with gr.Row(): | |
refresh_btn = gr.Button("Refresh Inventory") | |
export_btn = gr.DownloadButton("Export Inventory", value="/file=file.xlsx") | |
value_counts_df = get_value_counts_df(inventory_separate_entries) | |
inventory_title = gr.Label("Inventory Summary", color="grey") | |
inventory_table = gr.DataFrame(value_counts_df) | |
history_title = gr.Label("History of added items", color="grey") | |
history_table = gr.DataFrame(inventory_separate_entries) | |
output_file = gr.File(visible=False) | |
refresh_btn.click(refresh_inventory, outputs=[inventory_table, history_table]) | |
with gr.Blocks() as scan_interface: | |
with gr.Row(equal_height=False): | |
input = gr.Image() | |
greet_btn = gr.Button("Detect items") | |
with gr.Row() as row: | |
output = gr.Image() | |
with gr.Column(): | |
label = gr.Label("Nothing detected", visible=True) | |
add_item_btn = gr.Button("Add item", visible=True) | |
add_item_btn.click(add_item, inputs=label) | |
greet_btn.click(detect_objects, inputs=input, outputs=[output, label]) | |
demo = gr.TabbedInterface([scan_interface, inventory_interface], ["Scan items", "Inventory"]) | |
demo.launch(share=True, allowed_paths=["file.xlsx"]) |