pyesonekyaw commited on
Commit
29f3834
1 Parent(s): 7f71b29

added comments and explanations

Browse files
Files changed (1) hide show
  1. app.py +94 -63
app.py CHANGED
@@ -4,12 +4,15 @@ from PIL import Image
4
  import json
5
  import numpy as np
6
  import cv2
 
 
7
  week8_model = torch.hub.load(
8
  './', 'custom', path='Weights/Week_8.pt', source='local')
9
  week9_model = torch.hub.load(
10
  './', 'custom', path='Weights/Week_9.pt', source='local')
11
 
12
- def draw_own_bbox(img,x1,y1,x2,y2,label,color=(36,255,12),text_color=(0,0,0)):
 
13
  """
14
  Draw bounding box on the image with text label and save both the raw and annotated image in the 'own_results' folder
15
 
@@ -81,29 +84,47 @@ def draw_own_bbox(img,x1,y1,x2,y2,label,color=(36,255,12),text_color=(0,0,0)):
81
  x2 = int(x2)
82
  y1 = int(y1)
83
  y2 = int(y2)
84
-
85
- # Save the raw image
86
- # img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
87
-
88
  # Draw the bounding box
89
  img = cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
90
  # For the text background, find space required by the text so that we can put a background with that amount of width.
91
  (w, h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 1)
92
- # Print the text
93
  img = cv2.rectangle(img, (x1, y1 - 20), (x1 + w, y1), color, -1)
94
- img = cv2.putText(img, label, (x1, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.6, text_color, 1)
 
95
  return img
96
 
97
- def yolo(img, model, toggles, signal, size=1024):
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
98
  if model == "Week 8":
99
  model = week8_model
100
  else:
101
  model = week9_model
102
 
 
103
  results = model(img)
104
 
105
  # Original output image and results
106
- original_results = json.loads(results.pandas().xyxy[0].to_json(orient="records"))
 
107
  output_image = Image.fromarray(results.render()[0])
108
 
109
  # Convert the results to a pandas dataframe and calculate the height and width of the bounding box and the area of the bounding box
@@ -116,16 +137,17 @@ def yolo(img, model, toggles, signal, size=1024):
116
  df_results = df_results.sort_values('bboxArea', ascending=False)
117
 
118
  # Filter out Bullseye
119
- pred_list = df_results
120
  if 'Ignore Bullseye' in toggles:
121
  pred_list = pred_list[pred_list['name'] != 'Bullseye']
122
 
 
123
  if len(pred_list) == 0:
124
  return [output_image, original_results, output_image, original_results]
125
-
126
  elif len(pred_list) == 1:
127
  pred = pred_list.iloc[0]
128
-
129
  else:
130
  pred_shortlist = []
131
  current_area = pred_list.iloc[0]['bboxArea']
@@ -137,7 +159,7 @@ def yolo(img, model, toggles, signal, size=1024):
137
  pred_shortlist.append(row)
138
  # Update the current area to the area of the prediction
139
  current_area = row['bboxArea']
140
-
141
  # If only 1 prediction remains after filtering by confidence and area
142
  if len(pred_shortlist) == 1:
143
  # Choose that prediction
@@ -145,19 +167,19 @@ def yolo(img, model, toggles, signal, size=1024):
145
 
146
  # If multiple predictions remain after filtering by confidence and area
147
  else:
148
- # Use signal of {signal} to filter further
149
-
150
  # Sort the predictions by xmin
151
  pred_shortlist.sort(key=lambda x: x['xmin'])
152
 
153
  # If signal is 'L', choose the first prediction in the list, i.e. leftmost in the image
154
  if signal == 'L':
155
  pred = pred_shortlist[0]
156
-
157
  # If signal is 'R', choose the last prediction in the list, i.e. rightmost in the image
158
  elif signal == 'R':
159
  pred = pred_shortlist[-1]
160
-
161
  # If signal is 'C', choose the prediction that is central in the image
162
  else:
163
  # Loop through the predictions shortlist
@@ -166,74 +188,83 @@ def yolo(img, model, toggles, signal, size=1024):
166
  if pred_shortlist[i]['xmin'] > 250 and pred_shortlist[i]['xmin'] < 774:
167
  pred = pred_shortlist[i]
168
  break
169
-
170
  # If no prediction is central, choose the one with the largest area
171
- if isinstance(pred,str):
172
  # Choosing one with largest area if none are central
173
- pred_shortlist.sort(key=lambda x: x['bboxArea'])
174
  pred = pred_shortlist[-1]
175
-
176
-
177
- filtered_img = draw_own_bbox(np.array(img), pred['xmin'], pred['ymin'], pred['xmax'], pred['ymax'], pred['name'])
178
  return [output_image, original_results, filtered_img, json.loads(pred.to_json(orient="records"))]
179
 
180
-
181
- inputs = [gr.inputs.Image(type='pil', label="Original Image"),
182
- gr.inputs.Radio(['Week 8', 'Week 9'], type="value", default='Week 8', label='Model Selection'),
183
- gr.CheckboxGroup(["Ignore Bullseye", "Biggest BBox Only and Position-Based Heuristics",], value=["Ignore Bullseye", "Biggest BBox Only and Position-Based Heuristics"], label="Heuristic Toggles"),
184
- gr.inputs.Radio(['Left', 'Center', 'Right', 'Disabled'], type="value", default='Center', label='Position Heuristic'),
 
 
 
185
  ]
186
  outputs = [gr.outputs.Image(type="pil", label="Output Image"),
187
  gr.outputs.JSON(label="Output JSON"),
188
  gr.outputs.Image(type="pil", label="Filtered Output Image"),
189
  gr.outputs.JSON(label="Filtered Output JSON")
190
  ]
191
-
192
  examples = [['Examples/One.jpg'], ['Examples/Two.jpg'], ['Examples/Three.jpg'], ['Examples/1.jpg'], ['Examples/2.jpg'], ['Examples/3.jpg'], ['Examples/4.jpg'], ['Examples/5.jpg'], ['Examples/6.jpg'],
193
  ['Examples/7.jpg'], ['Examples/8.jpg'], ['Examples/9.jpg'], ['Examples/10.jpg'], ['Examples/11.jpg'], ['Examples/12.jpg']]
194
 
195
-
196
  with gr.Blocks(css="#custom_header {min-height: 2rem; text-align: center} #custom_title {min-height: 2rem}") as demo:
197
- gr.Markdown("# YOLOv5 Symbol Recognition for CZ3004/SC2079 Multi-Disciplinary Project", elem_id="custom_header")
198
- gr.Markdown("Gradio Demo for YOLOv5 Symbol Recognition for CZ3004 Multi-Disciplinary Project. To use it, simply upload your image, or click one of the examples to load them.", elem_id="custom_header")
199
- gr.Markdown("CZ3004 is a module in Nanyang Technological University's Computer Science curriculum that involves creating a robot car that can navigate within an arena and around obstacles. Part of the assessment is to go to obstacles and detect alphanumeric symbols pasted on them.", elem_id="custom_title")
200
  gr.Markdown("The two models available, Week 8 and Week 9, are for different subtasks. Week 8 model (as assessment was done in Week 8 of the school semester), \
201
  is able to detect all symbols seen in the first three example images below. Week 9 model is limited to just the bullseye, left and right arrow symbols. \
202
  Additionally, Week 9 model has been further trained on extreme edge cases where there is harsh sunlight behind the symbol/obstacle (seen in some of the examples).", elem_id="custom_title")
203
- gr.Markdown("Heuristics used are based on the task as of AY22-23 Semester 2's edition of MDP. These include ignoring the bullseye symbol, taking only the biggest bounding box, and filtering similar sized detections by the expected position of the symbol based on where the robot is supposed to be relative to the symbol.", elem_id="custom_title")
204
- gr.Markdown("This demo is part of a guide that is currently work-in-progress, for future CZ3004/SC2079 students to refer to.", elem_id="custom_title")
205
-
206
 
207
  with gr.Row():
208
- with gr.Column():
209
- with gr.Box():
210
- gr.Markdown("## Inputs", elem_id="custom_header")
211
- input_image = gr.inputs.Image(type='pil', label="Original Image")
212
- btn = gr.Button(value="Submit")
213
- btn.style(full_width=True)
214
- with gr.Column():
215
- with gr.Box():
216
- gr.Markdown("## Parameters", elem_id="custom_header")
217
- model_selection = gr.inputs.Radio(['Week 8', 'Week 9'], type="value", default='Week 8', label='Model Selection')
218
- toggles = gr.CheckboxGroup(["Ignore Bullseye", "Biggest BBox Only and Position-Based Heuristics",], value=["Ignore Bullseye", "Biggest BBox Only and Position-Based Heuristics"], label="Heuristic Toggles")
219
- radios = gr.inputs.Radio(['Left', 'Center', 'Right', 'Disabled'], type="value", default='Center', label='Position Heuristic')
220
- with gr.Row():
221
  with gr.Box():
222
- with gr.Column():
223
- gr.Markdown("## Raw Outputs", elem_id="custom_header")
224
- output_image = gr.outputs.Image(type="pil", label="Output Image")
225
- output_json = gr.outputs.JSON(label="Output JSON")
 
 
226
  with gr.Box():
227
- with gr.Column():
228
- gr.Markdown("## Filtered Outputs", elem_id="custom_header")
229
- filtered_image = gr.outputs.Image(type="pil", label="Filtered Output Image")
230
- filtered_json = gr.outputs.JSON(label="Filtered Output JSON")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
231
  with gr.Row():
232
  gr.Examples(examples=examples,
233
- inputs=input_image,
234
- outputs=output_image,
235
- fn=yolo,
236
- cache_examples=False)
237
- btn.click(yolo, inputs=[input_image, model_selection,toggles, radios], outputs=[output_image, output_json, filtered_image, filtered_json])
238
-
 
239
  demo.launch(debug=True)
 
4
  import json
5
  import numpy as np
6
  import cv2
7
+
8
+ # Load the models
9
  week8_model = torch.hub.load(
10
  './', 'custom', path='Weights/Week_8.pt', source='local')
11
  week9_model = torch.hub.load(
12
  './', 'custom', path='Weights/Week_9.pt', source='local')
13
 
14
+
15
+ def draw_own_bbox(img, x1, y1, x2, y2, label, color=(36, 255, 12), text_color=(0, 0, 0)):
16
  """
17
  Draw bounding box on the image with text label and save both the raw and annotated image in the 'own_results' folder
18
 
 
84
  x2 = int(x2)
85
  y1 = int(y1)
86
  y2 = int(y2)
 
 
 
 
87
  # Draw the bounding box
88
  img = cv2.rectangle(img, (x1, y1), (x2, y2), color, 2)
89
  # For the text background, find space required by the text so that we can put a background with that amount of width.
90
  (w, h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 1)
91
+ # Print the text
92
  img = cv2.rectangle(img, (x1, y1 - 20), (x1 + w, y1), color, -1)
93
+ img = cv2.putText(img, label, (x1, y1 - 5),
94
+ cv2.FONT_HERSHEY_SIMPLEX, 0.6, text_color, 1)
95
  return img
96
 
97
+
98
+ def yolo(img, model, toggles, signal):
99
+ """
100
+ Run YOLOv5 on the image and return the results
101
+
102
+ Inputs
103
+ ------
104
+ img: numpy.ndarray - image on which the YOLOv5 model is to be run
105
+ model: str - name of the model to be used
106
+ toggles: dict - dictionary containing the toggles for the model
107
+ signal: str - signal for position heuristic
108
+
109
+ Returns
110
+ -------
111
+ output_image: PIL.Image - image with bounding boxes drawn on it
112
+ original_results: json - json containing the original results
113
+ filtered_image: PIL.Image - image with bounding boxes drawn on it after filtering
114
+ filtered_results: json - json containing the filtered results
115
+ """
116
+ # Load the model based on the model name
117
  if model == "Week 8":
118
  model = week8_model
119
  else:
120
  model = week9_model
121
 
122
+ # Run the model on the image
123
  results = model(img)
124
 
125
  # Original output image and results
126
+ original_results = json.loads(
127
+ results.pandas().xyxy[0].to_json(orient="records"))
128
  output_image = Image.fromarray(results.render()[0])
129
 
130
  # Convert the results to a pandas dataframe and calculate the height and width of the bounding box and the area of the bounding box
 
137
  df_results = df_results.sort_values('bboxArea', ascending=False)
138
 
139
  # Filter out Bullseye
140
+ pred_list = df_results
141
  if 'Ignore Bullseye' in toggles:
142
  pred_list = pred_list[pred_list['name'] != 'Bullseye']
143
 
144
+ # If no predictions, return the empty results
145
  if len(pred_list) == 0:
146
  return [output_image, original_results, output_image, original_results]
147
+ # If only one prediction, no need to filter
148
  elif len(pred_list) == 1:
149
  pred = pred_list.iloc[0]
150
+ # If more than one prediction, filter the predictions
151
  else:
152
  pred_shortlist = []
153
  current_area = pred_list.iloc[0]['bboxArea']
 
159
  pred_shortlist.append(row)
160
  # Update the current area to the area of the prediction
161
  current_area = row['bboxArea']
162
+
163
  # If only 1 prediction remains after filtering by confidence and area
164
  if len(pred_shortlist) == 1:
165
  # Choose that prediction
 
167
 
168
  # If multiple predictions remain after filtering by confidence and area
169
  else:
170
+ # Use signal of {signal} to filter further
171
+
172
  # Sort the predictions by xmin
173
  pred_shortlist.sort(key=lambda x: x['xmin'])
174
 
175
  # If signal is 'L', choose the first prediction in the list, i.e. leftmost in the image
176
  if signal == 'L':
177
  pred = pred_shortlist[0]
178
+
179
  # If signal is 'R', choose the last prediction in the list, i.e. rightmost in the image
180
  elif signal == 'R':
181
  pred = pred_shortlist[-1]
182
+
183
  # If signal is 'C', choose the prediction that is central in the image
184
  else:
185
  # Loop through the predictions shortlist
 
188
  if pred_shortlist[i]['xmin'] > 250 and pred_shortlist[i]['xmin'] < 774:
189
  pred = pred_shortlist[i]
190
  break
191
+
192
  # If no prediction is central, choose the one with the largest area
193
+ if isinstance(pred, str):
194
  # Choosing one with largest area if none are central
195
+ pred_shortlist.sort(key=lambda x: x['bboxArea'])
196
  pred = pred_shortlist[-1]
197
+ # Draw the bounding box on the image
198
+ filtered_img = draw_own_bbox(np.array(
199
+ img), pred['xmin'], pred['ymin'], pred['xmax'], pred['ymax'], pred['name'])
200
  return [output_image, original_results, filtered_img, json.loads(pred.to_json(orient="records"))]
201
 
202
+ # Define the interface
203
+ inputs = [gr.inputs.Image(type='pil', label="Original Image"),
204
+ gr.inputs.Radio(['Week 8', 'Week 9'], type="value",
205
+ default='Week 8', label='Model Selection'),
206
+ gr.CheckboxGroup(["Ignore Bullseye", "Biggest BBox Only and Position-Based Heuristics",], value=[
207
+ "Ignore Bullseye", "Biggest BBox Only and Position-Based Heuristics"], label="Heuristic Toggles"),
208
+ gr.inputs.Radio(['Left', 'Center', 'Right', 'Disabled'],
209
+ type="value", default='Center', label='Position Heuristic'),
210
  ]
211
  outputs = [gr.outputs.Image(type="pil", label="Output Image"),
212
  gr.outputs.JSON(label="Output JSON"),
213
  gr.outputs.Image(type="pil", label="Filtered Output Image"),
214
  gr.outputs.JSON(label="Filtered Output JSON")
215
  ]
216
+ # Define the examples
217
  examples = [['Examples/One.jpg'], ['Examples/Two.jpg'], ['Examples/Three.jpg'], ['Examples/1.jpg'], ['Examples/2.jpg'], ['Examples/3.jpg'], ['Examples/4.jpg'], ['Examples/5.jpg'], ['Examples/6.jpg'],
218
  ['Examples/7.jpg'], ['Examples/8.jpg'], ['Examples/9.jpg'], ['Examples/10.jpg'], ['Examples/11.jpg'], ['Examples/12.jpg']]
219
 
220
+ # Define the gradio app
221
  with gr.Blocks(css="#custom_header {min-height: 2rem; text-align: center} #custom_title {min-height: 2rem}") as demo:
222
+ gr.Markdown("# YOLOv5 Symbol Recognition for CZ3004/SC2079 Multi-Disciplinary Project",
223
+ elem_id="custom_header")
224
+ gr.Markdown("Gradio Demo for YOLOv5 Symbol Recognition for CZ3004 Multi-Disciplinary Project. To use it, simply upload your image, or click one of the examples to load them. CZ3004 is a module in Nanyang Technological University's Computer Science curriculum that involves creating a robot car that can navigate within an arena and around obstacles. Part of the assessment is to go to obstacles and detect alphanumeric symbols pasted on them.", elem_id="custom_title")
225
  gr.Markdown("The two models available, Week 8 and Week 9, are for different subtasks. Week 8 model (as assessment was done in Week 8 of the school semester), \
226
  is able to detect all symbols seen in the first three example images below. Week 9 model is limited to just the bullseye, left and right arrow symbols. \
227
  Additionally, Week 9 model has been further trained on extreme edge cases where there is harsh sunlight behind the symbol/obstacle (seen in some of the examples).", elem_id="custom_title")
228
+ gr.Markdown("Heuristics used are based on AY22-23 Semester 2's edition of MDP. These include ignoring the bullseye symbol, taking only the biggest bounding box, and filtering similar sized detections by the expected position of the symbol based on where the robot is supposed to be relative to the symbol.", elem_id="custom_title")
229
+ gr.Markdown("This demo is part of a guide that is currently work-in-progress, for future CZ3004/SC2079 students to refer to. On a local environment, inference should be around 100ms at worst, and can be made faster with a GPU and/or conversion to a more optimized model format.", elem_id="custom_title")
 
230
 
231
  with gr.Row():
232
+ with gr.Column():
 
 
 
 
 
 
 
 
 
 
 
 
233
  with gr.Box():
234
+ gr.Markdown("## Inputs", elem_id="custom_header")
235
+ input_image = gr.inputs.Image(
236
+ type='pil', label="Original Image")
237
+ btn = gr.Button(value="Submit")
238
+ btn.style(full_width=True)
239
+ with gr.Column():
240
  with gr.Box():
241
+ gr.Markdown("## Parameters", elem_id="custom_header")
242
+ model_selection = gr.inputs.Radio(
243
+ ['Week 8', 'Week 9'], type="value", default='Week 8', label='Model Selection')
244
+ toggles = gr.CheckboxGroup(["Ignore Bullseye", "Biggest BBox Only and Position-Based Heuristics",], value=[
245
+ "Ignore Bullseye", "Biggest BBox Only and Position-Based Heuristics"], label="Heuristic Toggles")
246
+ radios = gr.inputs.Radio(['Left', 'Center', 'Right', 'Disabled'],
247
+ type="value", default='Center', label='Position Heuristic')
248
+ with gr.Row():
249
+ with gr.Box():
250
+ with gr.Column():
251
+ gr.Markdown("## Raw Outputs", elem_id="custom_header")
252
+ output_image = gr.outputs.Image(
253
+ type="pil", label="Output Image")
254
+ output_json = gr.outputs.JSON(label="Output JSON")
255
+ with gr.Box():
256
+ with gr.Column():
257
+ gr.Markdown("## Filtered Outputs", elem_id="custom_header")
258
+ filtered_image = gr.outputs.Image(
259
+ type="pil", label="Filtered Output Image")
260
+ filtered_json = gr.outputs.JSON(label="Filtered Output JSON")
261
  with gr.Row():
262
  gr.Examples(examples=examples,
263
+ inputs=input_image,
264
+ outputs=output_image,
265
+ fn=yolo,
266
+ cache_examples=False)
267
+ btn.click(yolo, inputs=[input_image, model_selection, toggles, radios], outputs=[
268
+ output_image, output_json, filtered_image, filtered_json])
269
+ # Run the gradio app
270
  demo.launch(debug=True)