Create main.py

main.py

@@ -0,0 +1,55 @@
+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/')