Spaces:
Running
Running
File size: 2,473 Bytes
146664f cf8f583 146664f a0c06ea 146664f d163e81 cc4ccdd d163e81 f34dd16 |
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 |
import gradio as gr
import torch
from PIL import Image
from PIL import ImageDraw
from PIL import ImageFont
import cv2
import numpy as np
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, live=True).launch(
debug=True, share=True) |