local-cv

Распознавание и визуализация

Структура папок проекта:

корень
      |____ weights/
      |____ src/
      |____ img/

Как использовать:

Создаем конфигурационный файл .yaml, сохраняем в папку /weights

- weights: VisBig_270922.pt
  classes:
    0: human
    1: car
    2: truck
    3: bike
  description: 'Модель для распознавания людей и автомобилей'
  menu_name: 'Люди и автомобили'
  
  # необязательные параметры
  # пиктограммы в папке ./img, рекомендуемый размер 32px
  picts: ['human.png', 'car.png', 'truck.png', 'bike.png']
  # цвет для рамки каждого класса
  cmap: ['#d62728', '#2ca02c', '#1f77b4', '#9467bd']

Пример рабочего кода

import cv2

# загружаем модули из папки src
import src.detector as Detector
import src.visualizer as Visualizer

# класс DetectorYolo содержит все необходимые методы для загрузки весов и распознавания
det = Detector.DetectorYolo()

# загружаем веса из папки weights. 
# загруженная модель будет храниться в объекте det
det.loadWeights('weights-test.yaml')

# класс Annotator служит для разметки изображений
vis = Visualizer.Annotator()

# готовим пиктограммы (метод преобразует их в маску)
# вариант 1 - если пиктограммы были указаны в конфиге и уже загружены
vis.setPicts(pict_files=det.models[0].picts)
# вариант 2 - если не были указаны в конфиге, или надо изменить
vis.setPicts(pict_files=['human.png', 'car.png', 'truck.png', 'bike.png'])

# загружаем изображение для проверки обнаружения
image = cv2.imread('img/image.jpg')

# делаем обнаружение, метод вернет numpy массив bounding box с абсолютными координатам углов x1y1-x2y2, уверенностью и класс объекта, см. пример
# [[     314.13      542.76      354.79       641.6     0.90917           0]
# [     118.47      608.15      161.96      702.94     0.86808           0]]
res = det.getDetect(image)

# тип разметки - bbox_types:
# 1 - обычная рамка
# 2 - окружность
# 3 - скобки
# picts - нарисовать иконку, conf_text - отобразить уверенность
image = vis.drawDetections(image, res[:], bbox_type=3, lw=2, picts = True, conf_text = True)

# выводим результат
cv2.imshow('Image',image)

cv2.waitKey(0)