File size: 2,327 Bytes
f7edb17
709f5ad
ed4a7a1
c5c5b54
 
ee7ed58
ed4a7a1
 
 
 
 
 
709f5ad
ed4a7a1
 
 
 
c5c5b54
 
 
 
ed4a7a1
 
 
ee7ed58
 
58d7863
ed4a7a1
8932c35
ed4a7a1
 
dc3de32
c4f0af6
dc3de32
ee7ed58
 
 
 
 
 
dc3de32
 
 
 
 
f3b0d7b
ed4a7a1
cd643ec
ed4a7a1
 
 
709f5ad
ee7ed58
 
ed4a7a1
07a88b5
ee7ed58
07a88b5
d72fe67
dc3de32
93d0dcd
c5c5b54
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
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()