Finding issue with object tracking and counting individual objects in yolov5.
I have tried yolov5 object tracking with custom data where I should count every individual defect but its not working with what I expected. This is my code for object detection counting defects in each frame:
import torch
import cv2
import numpy as np
import pandas as pd
import yolodetector
from deep_sort_realtime.deepsort_tracker import DeepSort
model = torch.hub.load('yolov5', 'custom', path="demo.pt", force_reload=True, source='local')
#model = torch.hub.load('ultralytics/yolov5', 'custom', 'G:\ola inference\olav2.onnx')
cap = cv2.VideoCapture("demo.avi")
def get_results():
total_defect= 0
D1=0
D2=0
D3=0
while cap.isOpened():
ret, frame = cap.read()
if ret == True:
# Make detections
results = model(frame)
# Remove confidence scores from labels
labels, cord = results.xyxyn[0][:, -1].to('cpu').numpy(), results.xyxyn[0][:, :-1].to('cpu').numpy()
n = len(labels)
total_defect = total_defect+n
x_shape, y_shape = frame.shape[1], frame.shape[0]
for i in range(n):
row = cord[i]
x1, y1, x2, y2 = int(row[0] * x_shape), int(row[1] * y_shape), int(row[2] * x_shape), int(row[3] * y_shape)
label = int(labels[i])
if label == 2:
label = "D1"
D2 = D1+1
elif label == 1:
label = "D2"
D2 = D2+1
else:
label = "D3"
D3 = D3+1
print("loading.....")
#cv2.imshow('YOLO', np.squeeze(results.render()))
else:
break
final_count={"totalcount" : [total_defect],
"D1": [D1],
"D2":[D2],
"D3":[D3]}
df = pd.DataFrame(final_count)
df.to_csv('file1.csv')
return df
I have tried many ways but I couldn't find a better solution. I need to add an object tracking with counting individual defects in a video (not in frames)
Thanks in advance