nolenfelten commited on
Commit
003ba4f
·
verified ·
1 Parent(s): 5e7ae62

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +48 -60
app.py CHANGED
@@ -283,81 +283,69 @@ def object_detection_density_edge(image, conf_threshold=0.25, iou_threshold=0.45
283
  summary = json.dumps({"object_count": int(object_count)}, indent=4)
284
  return image_with_density, json_response, summary
285
 
286
- def procedure(image_input, yolov7_confidence_threshold_input, yolov7_IOU_Threshold_input, roboflow_confidence_threshold_input, roboflow_IOU_Threshold_input, roboflow_labels_input, roboflow_stroke_width_input):
287
- '''
288
- This function takes in an image and applies both YOLOv7 and Roboflow object detection models to it.
289
- It then returns the images and JSON results.
290
- '''
291
- print("Begin Roboflow inferences.")
292
- roboflow_inference = roboflow(image=image_input, confidence=roboflow_confidence_threshold_input, overlap=roboflow_IOU_Threshold_input, stroke_width=roboflow_stroke_width_input, labels=roboflow_labels_input)
293
-
294
- if roboflow_inference["image"] is None:
295
- raise ValueError("Roboflow API did not return a valid image.")
296
-
297
- roboflow_image = roboflow_inference["image"]
298
- roboflow_json = json.dumps(roboflow_inference["json"], indent=4)
299
-
300
- return None, None, roboflow_image, roboflow_json
301
 
302
- # Uploaded image.
303
- image_input = gr.Image(label="Upload Image")
304
-
305
- # YOLOv7 Confidence Threshold input.
306
- yolov7_confidence_threshold_input = gr.Slider(minimum=0.0, maximum=1.0, value=0.45, step=0.01, label="YOLOv7 Confidence Threshold")
307
-
308
- # YOLOv7 IOU Threshold.
309
- yolov7_IOU_Threshold_input = gr.Slider(minimum=0.0, maximum=1.0, value=0.45, step=0.01, label="YOLOv7 IOU Threshold")
310
-
311
- # Roboflow Confidence Threshold input.
312
- roboflow_confidence_threshold_input = gr.Slider(minimum=0.0, maximum=1.0, value=0.45, step=0.01, label="Roboflow Confidence Threshold")
313
-
314
- # Roboflow IOU Threshold.
315
- roboflow_IOU_Threshold_input = gr.Slider(minimum=0.0, maximum=1.0, value=0.45, step=0.01, label="Roboflow IOU Threshold")
316
 
317
- # Roboflow Labels.
318
- roboflow_labels_input = gr.Checkbox(label="Roboflow Labels")
319
 
320
- # Roboflow Stroke Width.
321
- roboflow_stroke_width_input = gr.Radio([1, 2, 5, 10], label="Stroke Width")
 
 
 
 
 
322
 
323
- # YOLOv7 Image Output.
324
- yolov7_image_output = gr.Image(type="pil", label="YOLOv7 Output Image")
 
 
 
 
 
 
 
 
325
 
326
- # YOLOv7 JSON Output.
327
- yolov7_json_output = gr.Textbox(label="YOLOv7 Bounding Boxes JSON")
328
 
329
- # Roboflow Image Output.
330
- roboflow_image_output = gr.Image(type="pil", label="Roboflow Output Image")
 
 
 
 
331
 
332
- # Roboflow JSON Output.
333
- roboflow_json_output = gr.Textbox(label="Roboflow Bounding Boxes JSON")
334
 
335
- # Gradio Interface Definitions
336
- inputs = [
337
- image_input,
338
- yolov7_confidence_threshold_input,
339
- yolov7_IOU_Threshold_input,
340
- roboflow_confidence_threshold_input,
341
- roboflow_IOU_Threshold_input,
342
- roboflow_labels_input,
343
- roboflow_stroke_width_input,
344
- ]
345
 
346
- outputs = [
347
- yolov7_image_output,
348
- yolov7_json_output,
349
- roboflow_image_output,
350
- roboflow_json_output,
351
- ]
352
 
353
  title = "<center>Cigarette Pack Counter</center>"
354
- description = "<center><a href='http://counttek.online'><img width='25%' src='https://mvp-83056e96f7ab.herokuapp.com/static/countteklogo2.png'></a><br><a href='https://nolenfelten.github.io'>Project by Nolen Felten</a></center>"
355
  footer = ("<center><b>Item Classes it will detect (Total 140 Classes)</b></center>")
356
 
357
  interface = gr.Interface(
358
  fn=procedure,
359
- inputs=inputs,
360
- outputs=outputs,
 
 
 
 
 
 
 
 
 
 
 
 
 
361
  title=title,
362
  description=description,
363
  article=footer,
 
283
  summary = json.dumps({"object_count": int(object_count)}, indent=4)
284
  return image_with_density, json_response, summary
285
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
286
 
287
+ # Function to resize and encode an image
288
+ def resize_image(image, max_size=1500):
289
+ if max(image.size) > max_size:
290
+ ratio = max_size / float(max(image.size))
291
+ new_size = tuple([int(x * ratio) for x in image.size])
292
+ image = image.resize(new_size, Image.LANCZOS)
293
+ buffer = io.BytesIO()
294
+ image.save(buffer, format="PNG")
295
+ return buffer.getvalue()
 
 
 
 
 
296
 
 
 
297
 
298
+ def infer(image, model, version, api_key, confidence=0.45, overlap=0.45, format="json", labels=False, stroke=1):
299
+ base_url = f"https://detect.roboflow.com/{model}/{version}?api_key={api_key}&confidence={confidence}&overlap={overlap}&format={format}"
300
+
301
+ if format == "image":
302
+ if labels:
303
+ base_url += "&labels=on"
304
+ base_url += f"&stroke={stroke}"
305
 
306
+ image_data = resize_image(image)
307
+ encoded_image = base64.b64encode(image_data).decode("utf-8")
308
+
309
+ headers = {"Content-Type": "application/x-www-form-urlencoded"}
310
+ response = requests.post(base_url, data=encoded_image, headers=headers)
311
+
312
+ if format == "json":
313
+ return json.dumps(response.json(), indent=4)
314
+ elif format == "image":
315
+ return Image.open(io.BytesIO(response.content))
316
 
 
 
317
 
318
+ def gradio_infer(image, model, version, api_key, confidence, overlap, format, labels, stroke):
319
+ result = infer(image, model, version, api_key, confidence, overlap, format, labels, stroke)
320
+ if format == "json":
321
+ return result, None
322
+ else:
323
+ return None, result
324
 
 
 
325
 
 
 
 
 
 
 
 
 
 
 
326
 
 
 
 
 
 
 
327
 
328
  title = "<center>Cigarette Pack Counter</center>"
329
+ description = "<center><a href='http://counttek.online'><img width='25%' height='25%' src='https://mvp-83056e96f7ab.herokuapp.com/static/countteklogo2.png'></a><br><a href='https://nolenfelten.github.io'>Project by Nolen Felten</a></center>"
330
  footer = ("<center><b>Item Classes it will detect (Total 140 Classes)</b></center>")
331
 
332
  interface = gr.Interface(
333
  fn=procedure,
334
+ inputs=[
335
+ gr.Image(type="pil", label="Input Image"),
336
+ gr.Textbox(value="sku-110k", label="Model Name"),
337
+ gr.Textbox(value="2", label="Model Version"),
338
+ gr.Textbox(value="gHiUgOSq9GqTnRy5mErk", label="API Key"),
339
+ gr.Slider(0.0, 1.0, value=0.45, label="Confidence Threshold"),
340
+ gr.Slider(0.0, 1.0, value=0.45, label="Overlap Threshold"),
341
+ gr.Radio(["json", "image"], value="image", label="Output Format"),
342
+ gr.Checkbox(False, label="Include Labels"),
343
+ gr.Slider(1, 10, value=1, step=1, label="Stroke Width"),
344
+ ],
345
+ outputs=[
346
+ gr.Textbox(label="JSON Result"),
347
+ gr.Image(label="Output Image"),
348
+ ],
349
  title=title,
350
  description=description,
351
  article=footer,