import threading
from huggingface_hub import hf_hub_download
from ultralytics import YOLO as yolo
from supervision import Detections
from PyQt5 import *
from PIL import Image
import random
import time
import cv2
model = yolo("yolov3.pt")
frame_count = 0
results = []
videoCap = cv2.VideoCapture(0)
ret, frame = videoCap.read()
train_folder = "dataset_craters/train"
test_folder= "dataset_craters/test"
valid_folder="dataset_craters/valid"
yaml_file= "dataset_craters/data.yaml"
def getColours(cls_num):
"""Generate unique colors for each class ID"""
random.seed(cls_num)
return tuple(random.randint(0, 255) for _ in range(3))
def detect():
global results, frame
while True:
output = model(frame)
results = Detections.from_ultralytics(output[0])
time.sleep(0.0017)
t1 = threading.Thread(target=detect)
t1.start()
safe = True
while True and safe:
ret, frame = videoCap.read()
for result in results:
class_names = result[5]
if "person" in class_names['class_name']:
safe = True
print("person detected")
conf = result[2]
for box in result[0]:
if conf > 0.6:
x1, y1, x2, y2 = map(int, result[0])
# cls = int(box.cls[0])
class_name = class_names['class_name']
# conf = float(box.conf[0])
colour = getColours(class_name)
cv2.rectangle(frame, (x1, y1), (x2, y2), colour, 2)
cv2.putText(frame, f"{class_name} {conf:.2f}",
(x1, max(y1 - 10, 20)), cv2.FONT_HERSHEY_SIMPLEX,
0.6, colour, 2)
time.sleep(0.025)
if frame is not None:
cv2.imshow("Video", frame)
if cv2.waitKey(1) == ord('q'):
break
frame_count += 1
t1.join()
exit()
videoCap.release()
cv2.destroyAllWindows()
"""
Принимаем видео с вебкамеры, и пытаемся сегментировать его с помощью YOLOv8.
Выводим в окне маски объектов.
Потребуется установить пакеты следующей командой:
pip install ultralytics
"""
from sys import argv
from pathlib import Path
import numpy
import cv2
import torch
COLORS = [
(64, 128, 128),
(128, 64, 64),
(64, 128, 64),
(64, 64, 128),
(128, 64, 128),
]
VIDEO_SOURCE = 0 # 0 - вебкамера. Строка - имя файла или URL видео потока.
script_dir = Path(argv[0]).parent.resolve()
# определяем, на каком устройстве будет выполняться модель
device = (
'cuda' if torch.cuda.is_available() else
'mps' if torch.backends.mps.is_available() else
'cpu'
)
import ultralytics
def text_at( # рисует текст по центру
img: numpy.ndarray, # на чём
text: str, # что
pos: tuple[int, int], # где центр
font=cv2.FONT_HERSHEY_COMPLEX, # каким шрифтом
font_scale: float = 1.0, # в каком масштабе (не кегль!)
font_thickness: int = 2, # насколько толстые линии
text_color=(255, 0, 0), # каким цветом
bg_color=None # на каком фоне
):
x, y = pos
(text_w, text_h), _ = cv2.getTextSize(text, font, font_scale, font_thickness)
if bg_color is not None:
cv2.rectangle(img,
(x - text_w // 2, y - text_h // 2),
(x + text_w // 2, y + text_h // 2),
bg_color, -1)
cv2.putText(img,
text,
(x - text_w // 2, y + text_h // 2 - 1),
font, font_scale, text_color, font_thickness)
model = ultralytics.YOLO(
script_dir / 'yolov8m-seg.pt' # имя файла модели указывает на её тип
)
# цикл работы с видео
video = cv2.VideoCapture(VIDEO_SOURCE)
if not video.isOpened():
raise SystemExit('Не удалось открыть источник видео')
overlay_image: numpy.ndarray = numpy.zeros( # изображение-оверлей для отображения разметки кадра
(
int(video.get(cv2.CAP_PROP_FRAME_HEIGHT)),
int(video.get(cv2.CAP_PROP_FRAME_WIDTH)),
3
),
numpy.uint8
)
try:
while True:
success, frame = video.read() # читаем кадр
if not success:
raise SystemExit('Видео закончилось')
# для YOLO предобработка не требуется
outputs = model.predict( # получаем отклик сети
source=frame,
conf=0.15, # минимальная степень уверенности в объекте
device=device, # устройство
verbose=False # чтобы не спамило логом в консоль
)
output = outputs[0] # модель всегда возвращает список результатов, даже для одного изображения
if output.masks is not None: # если были получены маски
masks = output.masks.cpu().xy # список масок как координат вершин многоугольников (контуров)
boxes = output.boxes.cpu() # ограничивающие прямоугольники
image_output = [] # список аннотаций, которые нужно будет вывести
for box, mask in zip(boxes, masks): # формируем список аннотаций
class_id = int(box.cls.item()) # номер класса, соответствующего очередной области
class_name = model.names[class_id] # имя класса, соответствующее области
pts = numpy.array(mask).astype(numpy.int32) # массив координат точек контура Nx2
image_output.append((class_id, class_name, pts))
overlay_image.fill(0) # очищаем изображение-оверлей
for (class_id, class_name, pts) in image_output: # закрашиваем каждую маску цветом
cv2.fillPoly(
img=overlay_image,
pts=pts[numpy.newaxis, ...], # для работы fillPoly() массив должен быть вида 1xNx2
color=COLORS[class_id % len(COLORS)], # цвет массива определяем по номеру класса для стабильности
)
for (class_id, class_name, pts) in image_output: # выводим надписи отдельным циклом, чтобы их не закрасило
M = cv2.moments(pts) # вычисляем моменты контура
# используем их для расчёта координат центроида контура
cX = int(M["m10"] / M["m00"])
cY = int(M["m01"] / M["m00"])
# выводим имя класса
text_at(overlay_image, class_name,
(cX, cY),
text_color=(255, 255, 255),
bg_color=(1, 1, 1))
# накладываем оверлей с разметкой на кадр
alpha = 0.8 # вес оверлея, чем меньше - тем он прозрачнее.
apply = overlay_image > 0 # изменяем только те части кадра, где есть оверлей
frame[apply] = alpha * overlay_image[apply] + (1 - alpha) * frame[apply] # обожаю numpy
cv2.imshow('Press Esc to exit', frame) # показываем результат
if cv2.waitKey(10) == 27: # если пользователь нажал Esc, выходим
break
finally:
video.release() # не забываем отключиться от источника видео в итоге
Средний
Простой
Простой
Простой
Простой
Средний