File size: 5,653 Bytes
b2fbe3d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
import os
import tensorflow as tf
import tensorflow_hub as hub
# Load compressed models from tensorflow_hub
os.environ['TFHUB_MODEL_LOAD_FORMAT'] = 'COMPRESSED'

import matplotlib.pyplot as plt
import matplotlib as mpl

# For drawing onto the image.
import numpy as np
from tensorflow.python.ops.numpy_ops import np_config
np_config.enable_numpy_behavior()
from PIL import Image
from PIL import ImageColor
from PIL import ImageDraw
from PIL import ImageFont
import time

import streamlit as st

# For measuring the inference time.
import time


class ObjectDetector:

    def __init__(self):
        # Load Tokenizer & Model
        # hub_location = 'cardiffnlp/twitter-roberta-base-sentiment'
        # self.tokenizer = AutoTokenizer.from_pretrained(hub_location)
        # self.model = AutoModelForSequenceClassification.from_pretrained(hub_location)

        # Change model labels in config
        # self.model.config.id2label[0] = "Negative"
        # self.model.config.id2label[1] = "Neutral"
        # self.model.config.id2label[2] = "Positive"
        # self.model.config.label2id["Negative"] = self.model.config.label2id.pop("LABEL_0")
        # self.model.config.label2id["Neutral"] = self.model.config.label2id.pop("LABEL_1")
        # self.model.config.label2id["Positive"] = self.model.config.label2id.pop("LABEL_2")

        # Instantiate explainer
        # self.explainer = SequenceClassificationExplainer(self.model, self.tokenizer)
        
        # module_handle = "https://tfhub.dev/google/faster_rcnn/openimages_v4/inception_resnet_v2/1" 
        module_handle = "https://tfhub.dev/google/openimages_v4/ssd/mobilenet_v2/1"
        self.detector = hub.load(module_handle).signatures['default']

    def run_detector(self, path):
      img = path
      
      converted_img  = tf.image.convert_image_dtype(img, tf.float32)[tf.newaxis, ...]
     
      start_time = time.time()
      result = self.detector(converted_img)
      end_time = time.time()
    
      result = {key:value.numpy() for key,value in result.items()}
    
      primer = format(result["detection_class_entities"][0]) + ' ' + format(round(result["detection_scores"][0]*100)) + '%'
    
      image_with_boxes = self.draw_boxes(
        img, result["detection_boxes"],
        result["detection_class_entities"], result["detection_scores"])
    
      # display_image(image_with_boxes)
      return image_with_boxes, primer
  
    def display_image(self, image):
      fig = plt.figure(figsize=(20, 15))
      plt.grid(False)
      plt.imshow(image)
    
    def draw_bounding_box_on_image(self, image,
                                   ymin,
                                   xmin,
                                   ymax,
                                   xmax,
                                   color,
                                   font,
                                   thickness=4,
                                   display_str_list=()):
      """Adds a bounding box to an image."""
      draw = ImageDraw.Draw(image)
      im_width, im_height = image.size
      (left, right, top, bottom) = (xmin * im_width, xmax * im_width,
                                    ymin * im_height, ymax * im_height)
      draw.line([(left, top), (left, bottom), (right, bottom), (right, top),
                 (left, top)],
                width=thickness,
                fill=color)
    
      # If the total height of the display strings added to the top of the bounding
      # box exceeds the top of the image, stack the strings below the bounding box
      # instead of above.
      display_str_heights = [font.getsize(ds)[1] for ds in display_str_list]
      # Each display_str has a top and bottom margin of 0.05x.
      total_display_str_height = (1 + 2 * 0.05) * sum(display_str_heights)
    
      if top > total_display_str_height:
        text_bottom = top
      else:
        text_bottom = top + total_display_str_height
      # Reverse list and print from bottom to top.
      for display_str in display_str_list[::-1]:
        text_width, text_height = font.getsize(display_str)
        margin = np.ceil(0.05 * text_height)
        draw.rectangle([(left, text_bottom - text_height - 2 * margin),
                        (left + text_width, text_bottom)],
                       fill=color)
        draw.text((left + margin, text_bottom - text_height - margin),
                  display_str,
                  fill="black",
                  font=font)
        text_bottom -= text_height - 2 * margin
    
    def draw_boxes(self, image, boxes, class_names, scores, max_boxes=10, min_score=0.4):
      """Overlay labeled boxes on an image with formatted scores and label names."""
      colors = list(ImageColor.colormap.values())
    
      try:
        font = ImageFont.truetype("./Roboto-Light.ttf", 24)
          
      except IOError:
        print("Font not found, using default font.")
        font = ImageFont.load_default()
    
      for i in range(min(boxes.shape[0], max_boxes)):
        if scores[i] >= min_score:
          ymin, xmin, ymax, xmax = tuple(boxes[i])
          display_str = "{}: {}%".format(class_names[i].decode("ascii"),
                                         int(100 * scores[i]))
          color = colors[hash(class_names[i]) % len(colors)]
          image_pil = Image.fromarray(np.uint8(image)).convert("RGB")
          self.draw_bounding_box_on_image(
              image_pil,
              ymin,
              xmin,
              ymax,
              xmax,
              color,
              font,
              display_str_list=[display_str])
          np.copyto(image, np.array(image_pil))
      return image