Update app.py

app.py

@@ -1,63 +1,123 @@
 import gradio as gr
 import torch
 from PIL import Image
-from PIL import ImageDraw
-from PIL import ImageFont
+# from PIL import ImageDraw
+# from PIL import ImageFont
 import cv2
 import numpy as np
+import argparse
 
-model = torch.hub.load('ultralytics/yolov5', 'custom', path='model/yolov5n_rebar_kaggle.pt')
-
-def yolo(im, conf, iou, size=640):
-
-    mask = np.array(im["mask"])
-    mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
-    contours, hierarchy = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
-    if contours:
-      mask = np.zeros(mask.shape, np.uint8)
-      cnt = contours[0]
-      mask = cv2.drawContours(mask, [cnt], 0, 255, -1)
-      im = np.array(im["image"])
-      im = cv2.bitwise_and(im, im, mask=mask)
-      im = Image.fromarray(im)
-    else:
-      im = im["image"]
-    
-    model.conf = conf
-    model.iou = iou
-    results = model(im, size=size)  # custom inference size
-     # inference
-
-    # output_im = Image.fromarray(results.render(labels=False)[0])
-    # output_im = results.render(labels=False)[0]
-    output_im = np.array(im)
-
-    pred = results.pandas().xyxy[0]
-    counting = pred.shape[0]
-    text = f"{counting} objects"
-    for index, row in pred.iterrows():
-      cv2.circle(output_im, (int((row["xmin"] + row["xmax"]) * 0.5), int((row["ymin"] + row["ymax"]) * 0.5)), int((row["xmax"] - row["xmin"]) * 0.5 * 0.6), (255, 0, 0), -1)
- 
-    return Image.fromarray(output_im), text
-
-slider_step = 0.05
-nms_conf = 0.25
-nms_iou = 0.1
-# inputs_image = gr.inputs.Image(type='pil', label="Original Image")
-inputs_image = gr.inputs.Image(tool="sketch", label="Original Image",type="pil")
-inputs_conf = gr.Slider(0, 1, step=slider_step, value=nms_conf, label="Conf Thres")
-inputs_iou = gr.Slider(0, 1, step=slider_step, value=nms_iou, label="IoU Thres")
-
-inputs = [inputs_image, inputs_conf, inputs_iou]
-
-outputs_image = gr.outputs.Image(type="pil", label="Output Image")
-outputs_text = gr.Textbox(label="Number of objects")
-outputs = [outputs_image, outputs_text]
-
-title = "OBJECT COUNTING"
-description = "Object counting demo. Upload an image or click an example image to use. You can select the area to count by drawing a closed area on the input image."
-article = "<p style='text-align: center'>Counting objects in image</a></p>"
-examples = [['./images/S__275668998.jpg'], ['./images/S__271433737.jpg']]
-
-gr.Interface(yolo, inputs, outputs, title=title, description=description, article=article, examples=examples, cache_examples=False, analytics_enabled=False).launch(
-    debug=True)#, share=True)
+def parse_args(known=False):
+  parser = argparse.ArgumentParser(description="Object Counting v1.0")
+  parser.add_argument(
+      "--yolov5_path",
+      "-yp",
+      default="ultralytics/yolov5",
+      type=str,
+      help="yolov5 path",
+  )
+  parser.add_argument(
+      "--model_path",
+      "-mp",
+      default="model/yolov5n_rebar_kaggle.pt",
+      type=str,
+      help="model path",
+  )
+  parser.add_argument(
+      "--nms_conf",
+      "-conf",
+      default=0.25,
+      type=float,
+      help="model NMS confidence threshold",
+  )
+  parser.add_argument(
+      "--nms_iou", 
+      "-iou", 
+      default=0.1, 
+      type=float, 
+      help="model NMS IoU threshold"
+  )
+  parser.add_argument(
+      "--device",
+      "-dev",
+      default="cuda:0",
+      type=str,
+      help="cuda or cpu",
+  )
+  parser.add_argument("--inference_size", "-isz", default=640, type=int, help="model inference size")
+  parser.add_argument("--slider_step", "-ss", default=0.05, type=float, help="slider step")
+
+  args = parser.parse_known_args()[0] if known else parser.parse_args()
+  return args
+
+def image_roi(im):
+  mask = np.array(im["mask"])
+  mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
+  contours, hierarchy = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
+  if contours:
+    mask = np.zeros(mask.shape, np.uint8)
+    cnt = contours[0]
+    mask = cv2.drawContours(mask, [cnt], 0, 255, -1)
+    im = np.array(im["image"])
+    im = cv2.bitwise_and(im, im, mask=mask)
+    im = Image.fromarray(im)
+  else:
+    im = im["image"] 
+  return im
+
+def yolo(im, conf, iou, size):
+
+  global model
+  if im is None:
+    print ("No image")
+    return None, None
+  im = image_roi(im)
+  model.conf = conf
+  model.iou = iou
+  results = model(im, size=size)  # custom inference size
+  output_im = np.array(im)
+
+  pred = results.pandas().xyxy[0]
+  counting = pred.shape[0]
+  text = f"{counting} objects"
+  for index, row in pred.iterrows():
+    cv2.circle(output_im, (int((row["xmin"] + row["xmax"]) * 0.5), int((row["ymin"] + row["ymax"]) * 0.5)), int((row["xmax"] - row["xmin"]) * 0.5 * 0.6), (255, 0, 0), -1)
+
+  return Image.fromarray(output_im), text
+
+
+def main(args):
+  gr.close_all()
+
+  global model
+  yolo_path = args.yolov5_path
+  model_path = args.model_path
+  nms_conf = args.nms_conf
+  nms_iou = args.nms_iou
+  device = args.device
+  inference_size = args.inference_size
+  slider_step = args.slider_step
+
+  model = torch.hub.load(yolo_path, 'custom', path=model_path, device=device)
+  inputs_image = gr.inputs.Image(tool="sketch", label="Original Image",type="pil")
+  inputs_conf = gr.Slider(0, 1, step=slider_step, value=nms_conf, label="Conf Thres")
+  inputs_iou = gr.Slider(0, 1, step=slider_step, value=nms_iou, label="IoU Thres")
+  inputs_size = gr.Slider(384, 1536, step=128, value=inference_size, label="Inference Size")
+  inputs = [inputs_image, inputs_conf, inputs_iou, inputs_size]
+
+  outputs_image = gr.outputs.Image(type="pil", label="Output Image")
+  outputs_text = gr.Textbox(label="Number of objects")
+  outputs = [outputs_image, outputs_text]
+
+  title = "OBJECT COUNTING"
+  description = "Object counting demo. Upload an image or click an example image to use. You can select the area to count by drawing a closed area on the input image."
+  article = "<p style='text-align: center'>Counting objects in image</a></p>"
+  examples = [['./images/S__275668998.jpg'], ['./images/S__275669003.jpg'], ['./images/S__275669004.jpg']]
+
+  gr.Interface(fn=yolo, inputs=inputs, outputs=outputs, title=title, description=description, article=article, examples=examples, cache_examples=False, analytics_enabled=False).launch(
+      debug=True, share=True)
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    main(args)