0

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

Markus
  • 5,976
  • 5
  • 6
  • 21

0 Answers0