Update app.py

app.py

@@ -6,8 +6,58 @@ from PIL import ImageFont
 import cv2
 import numpy as np
 
-def greet(name):
-    return "Hello " + name + "!!"
+model = torch.hub.load('ultralytics/yolov5', 'custom', path='yolov5n_rebar_kaggle.pt')
 
-iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-iface.launch()
+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 = [['S__275668998.jpg']]
+
+gr.Interface(yolo, inputs, outputs, title=title, description=description, article=article, examples=examples, analytics_enabled=False).launch(
+    debug=True)#, share=True)