Spaces:
Runtime error
Runtime error
| from transformers import DetrFeatureExtractor, DetrForObjectDetection | |
| from transformers import pipeline | |
| from PIL import Image, ImageDraw, ImageFont | |
| import gradio as gr | |
| # Draw bounding box definition | |
| def draw_bounding_box(im, score, label, xmin, ymin, xmax, ymax, index, num_boxes): | |
| """ Draw a bounding box. """ | |
| # Draw the actual bounding box | |
| outline = 'blue' | |
| im_with_rectangle = ImageDraw.Draw(im) | |
| im_with_rectangle.rounded_rectangle((xmin, ymin, xmax, ymax), outline = outline, width = 2, radius = 7) | |
| # Return the result | |
| return im | |
| def object_classify(img): | |
| feature_extractor = DetrFeatureExtractor.from_pretrained('facebook/detr-resnet-50') | |
| model = DetrForObjectDetection.from_pretrained('facebook/detr-resnet-50') | |
| object_detector = pipeline("object-detection", model = model, feature_extractor = feature_extractor) | |
| bboxes = object_detector(img) | |
| price_total = 0 | |
| total_items_in_cart = 0 | |
| # Iteration elements | |
| num_boxes = len(bboxes) | |
| index = 0 | |
| # Draw bounding box for each result and count the price | |
| for i in bboxes: | |
| if i['label'] == 'apple': | |
| price_total += 25 #pesos? dunno | |
| elif i['label'] == 'bottle': | |
| price_total += 15 | |
| elif i['label'] == 'broccoli': | |
| price_total += 100 | |
| elif i['label'] == 'orange': | |
| price_total += 20 | |
| elif i['label'] == 'banana': | |
| price_total += 50 | |
| box = i['box'] | |
| #Draw the bounding box | |
| output_image = draw_bounding_box(img, i['score'],i['label'], | |
| box['xmin'], box['ymin'], | |
| box['xmax'], box['ymax'], | |
| index, num_boxes) | |
| index += 1 | |
| total_items_in_cart += 1 | |
| return output_image, str(price_total), str(total_items_in_cart) | |
| TITLE = 'Object Detection for Effective Self-Checkout in Grocery Shopping [Work In Progress]' | |
| DESCRIPTION = 'A deep learning application using DETR model to reimagine self-checkout stores.' | |
| EXAMPLES = ['ex1.jpg'] | |
| interface=gr.Interface(object_classify, | |
| gr.inputs.Image(type = 'pil'),outputs = [gr.outputs.Image(), gr.outputs.Textbox(label='Total Price: '), gr.outputs.Textbox(label='Total items in cart: ')], | |
| examples = EXAMPLES,title = TITLE, description=DESCRIPTION, allow_flagging="never") | |
| interface.launch() |