https://github.com/Sharpiless/yolov5-deepsort
Используется YOLOv5+Deepsort для отслеживания и подсчета автомобилей и пешеходов. Код упакован в класс Detector, что облегчает его встраивание в собственные проекты.
Адрес кода (ваше звездание приветствуется):
https://github.com/Sharpiless/Yolov5-deepsort-inference
Конечный результат:
class Detector(baseDet):
def __init__(self):
super(Detector, self).__init__()
self.init_model()
self.build_config()
def init_model(self):
self.weights = 'weights/yolov5m.pt'
self.device = '0' if torch.cuda.is_available() else 'cpu'
self.device = select_device(self.device)
model = attempt_load(self.weights, map_location=self.device)
model.to(self.device).eval()
model.half()
# torch.save(model, 'test.pt')
self.m = model
self.names = model.module.names if hasattr(
model, 'module') else model.names
def preprocess(self, img):
img0 = img.copy()
img = letterbox(img, new_shape=self.img_size)[0]
img = img[:, :, ::-1].transpose(2, 0, 1)
img = np.ascontiguousarray(img)
img = torch.from_numpy(img).to(self.device)
img = img.half() # half-precision
img /= 255.0 # image normalization
if img.ndimension() == 3:
img = img.unsqueeze(0)
``` return img0, img
def detect(self, im):
im0, img = self.preprocess(im)
pred = self.m(img, augment=False)[0]
pred = pred.float()
pred = non_max_suppression(pred, self.threshold, 0.4)
pred_boxes = []
for det in pred:
if det is not None and len(det):
det[:, :4] = scale_coords(
img.shape[2:], det[:, :4], im0.shape).round()
# Трекер DeepSort:
```python
deepsort = DeepSort(cfg.DEEPSORT.REID_CKPT,
max_dist=cfg.DEEPSORT.MAX_DIST, min_confidence=cfg.DEEPSORT.MIN_CONFIDENCE,
nms_max_overlap=cfg.DEEPSORT.NMS_MAX_OVERLAP, max_iou_distance=cfg.DEEPSORT.MAX_IOU_DISTANCE,
max_age=cfg.DEEPSORT.MAX_AGE, n_init=cfg.DEEPSORT.N_INIT, nn_budget=cfg.DEEPSORT.NN_BUDGET,
use_cuda=True)
Вызовите метод self.update
для обновления результатов трекинга.
python demo.py
См. мою другую статью:
Обучение собственной модели YOLOv5 (от настройки окружения Windows до развертывания модели)
После обучения сохраните модель в папке weights
.
from AIDetector_pytorch import Detector
det = Detector()
func_status = {}
func_status['headpose'] = None
result = det.feedCap(im, func_status)
Здесь im
- это изображение в формате BGR.
Возвращаемый result
- это словарь, result['frame']
возвращает отображаемое изображение.
Если вас заинтересовало, подпишитесь на мой канал - "Доступный глубокий обучение":
Bilibili: https://space.bilibili.com/470550823
AI Studio: https://aistudio.baidu.com/aistudio/personalcenter/thirdview/67156
Github: https://github.com/Sharpiless
При использовании следует указать лицензию GNU General Public License v3.0. Оригинальная часть кода для цели детекции взята из: https://github.com/ultralytics/yolov5/
Вы можете оставить комментарий после Вход в систему
Неприемлемый контент может быть отображен здесь и не будет показан на странице. Вы можете проверить и изменить его с помощью соответствующей функции редактирования.
Если вы подтверждаете, что содержание не содержит непристойной лексики/перенаправления на рекламу/насилия/вульгарной порнографии/нарушений/пиратства/ложного/незначительного или незаконного контента, связанного с национальными законами и предписаниями, вы можете нажать «Отправить» для подачи апелляции, и мы обработаем ее как можно скорее.
Комментарии ( 0 )