File size: 1,311 Bytes
d11c72e 37bbbbc 8c43771 d11c72e 8c43771 d11c72e 8c43771 d11c72e 8c43771 d11c72e 8c43771 d11c72e 37bbbbc d11c72e 37bbbbc d11c72e |
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 |
# app.py for Hugging Face Spaces
import gradio as gr
from ultralyticsplus import YOLO
from PIL import Image
import numpy as np
# Load the model
model = YOLO('foduucom/product-detection-in-shelf-yolov8')
# Set model parameters
model.overrides['conf'] = 0.25 # NMS confidence threshold
model.overrides['iou'] = 0.45 # NMS IoU threshold
model.overrides['agnostic_nms'] = False # NMS class-agnostic
model.overrides['max_det'] = 1000 # maximum number of detections per image
def detect_skus(image):
# Convert Gradio image to numpy array if needed, but YOLO accepts PIL
results = model(image)
# Extract unique SKU names (assuming classes are SKU names)
sku_set = set()
for result in results:
for box in result.boxes:
class_id = int(box.cls)
sku_name = result.names[class_id]
sku_set.add(sku_name)
sku_list = list(sku_set)
return "\n".join(sku_list) if sku_list else "No SKUs detected."
# Create Gradio interface
iface = gr.Interface(
fn=detect_skus,
inputs=gr.Image(type="pil", label="Upload Shelf Image"),
outputs=gr.Textbox(label="Detected SKUs"),
title="Shelf SKU Detector",
description="Upload an image of a cigarette shelf to detect and list the SKUs."
)
if __name__ == "__main__":
iface.launch() |