Spaces:
No application file
No application file
import yolov5 | |
# load model | |
model = yolov5.load('keremberke/yolov5m-garbage') | |
# set model parameters | |
model.conf = 0.25 # NMS confidence threshold | |
model.iou = 0.45 # NMS IoU threshold | |
model.agnostic = False # NMS class-agnostic | |
model.multi_label = False # NMS multiple labels per box | |
model.max_det = 1000 # maximum number of detections per image | |
# set image | |
img = 'sample.jpg' | |
# perform inference | |
results = model(img, size=640) | |
# inference with test time augmentation | |
results = model(img, augment=True) | |
# parse results | |
predictions = results.pred[0] | |
boxes = predictions[:, :4] # x1, y1, x2, y2 | |
scores = predictions[:, 4] | |
categories = predictions[:, 5] | |
# show detection bounding boxes on image | |
results.show() | |
print(results) # results output image 1/1: 720x1280 2 biodegradables, 1 paper | |
sample = str(results).split() | |
# print(type(results)) | |
# regex to match the particular things | |
ct =0 | |
for words in sample: | |
if(words=='paper' or words =='plastic' or words =='rubber'): | |
ct+=1 | |
if(ct >5): | |
print("Image contains garbage") | |
else: | |
print("Image does not contain garbage") | |
# save results into "results/" folder | |
# results.save(save_dir='results/') | |