def __init__(self):
self._net = YOLOv5(model_path, device) #jetson.inference.detectNet("ssd-mobilenet-v2")
self.img = None
self.width = None
self.height = None
self.need_cam_info = True
self.camera_model = PinholeCameraModel()
self.marker = Marker()
self.prev_time = 0
self.marker_pub = rospy.Publisher("visualization_markers", Marker, queue_size=10)
self.camera_info = rospy.Subscriber('/camera/depth/camera_info', CameraInfo, self.info_callback)
def test_load_model(self):
from yolov5 import YOLOv5
# init model
model_path = TestConstants.YOLOV5S_MODEL_PATH
device = "cpu"
yolov5 = YOLOv5(model_path, device, load_on_init=False)
yolov5.load_model()
# check if loaded
self.assertNotEqual(yolov5.model, None)
import time
import cv2
from yolov5 import YOLOv5
import torch
import torchvision.transforms as T
from PIL import Image
#import random
helmet_model_path = os.path.join(os.path.dirname(os.path.realpath(__file__)),
"yolov5_m_hat.pt")
mask_model_path = os.path.join(os.path.dirname(os.path.realpath(__file__)),
"mask_classification_model.pth")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
helmet_model = YOLOv5(helmet_model_path, device) # сеть поиска касок
mtcnn = MTCNN(keep_all=True, device=device) # сеть поиска лиц
mask_model = torch.load(
mask_model_path, map_location=device) # классификация лиц в маске или без
mask_model.eval()
gis = ImageSignature()
def add_new_face_frame(frame):
"""
функция для добавления лиц по кадрам
:param frame:
:return:
"""
_, width = frame.shape[1], frame.shape[0]
k = 900 / width
from PIL import Image
from yolov5 import YOLOv5
# set model params
model_path = "/Users/saskia/Dropbox/ELLA/yolov5_fluffy_tello/yolov5/weights/best.pt" # it automatically downloads yolov5s model to given path
device = "cpu" # or "cuda"
# init yolov5 model
yolov5 = YOLOv5(model_path, device)
# load images
image1 = Image.open(
"/Users/saskia/Dropbox/ELLA/yolov5_fluffy_tello/ella_and_fredi.jpg")
# image2 = Image.open("yolov5/data/images/zidane.jpg")
# perform inference
results = yolov5.predict(image1)
# perform inference with higher input size
results = yolov5.predict(image1, size=1280)
# perform inference with test time augmentation
results = yolov5.predict(image1, augment=True)
# perform inference on multiple images
# results = yolov5.predict([image1, image2], size=1280, augment=True)
def test_predict(self):
from PIL import Image
from yolov5 import YOLOv5
# init yolov5s model
model_path = TestConstants.YOLOV5S_MODEL_PATH
device = "cpu"
yolov5 = YOLOv5(model_path, device, load_on_init=True)
# prepare image
image_path = TestConstants.ZIDANE_IMAGE_PATH
image = Image.open(image_path)
# perform inference
results = yolov5.predict(image, size=640, augment=False)
# compare
self.assertEqual(results.n, 1)
self.assertEqual(len(results.names), 80)
self.assertEqual(len(results.pred[0]), 4)
# prepare image
image = image_path
# perform inference
results = yolov5.predict(image, size=640, augment=False)
# compare
self.assertEqual(results.n, 1)
self.assertEqual(len(results.names), 80)
self.assertEqual(len(results.pred[0]), 4)
# init yolov5l model
model_path = TestConstants.YOLOV5L_MODEL_PATH
device = "cpu"
yolov5 = YOLOv5(model_path, device, load_on_init=True)
# prepare image
image_path = TestConstants.BUS_IMAGE_PATH
image = Image.open(image_path)
# perform inference
results = yolov5.predict(image, size=1280, augment=False)
# compare
self.assertEqual(results.n, 1)
self.assertEqual(len(results.names), 80)
self.assertEqual(len(results.pred[0]), 8)
# prepare image
image = image_path
# perform inference
results = yolov5.predict(image, size=1280, augment=False)
# compare
self.assertEqual(results.n, 1)
self.assertEqual(len(results.names), 80)
self.assertEqual(len(results.pred[0]), 8)
# init yolov5s model
model_path = TestConstants.YOLOV5S_MODEL_PATH
device = "cpu"
yolov5 = YOLOv5(model_path, device, load_on_init=True)
# prepare images
image_path1 = TestConstants.ZIDANE_IMAGE_PATH
image_path2 = TestConstants.BUS_IMAGE_PATH
image1 = Image.open(image_path1)
image2 = Image.open(image_path2)
# perform inference with multiple images and test augmentation
results = yolov5.predict([image1, image2], size=1280, augment=True)
# compare
self.assertEqual(results.n, 2)
self.assertEqual(len(results.names), 80)
self.assertEqual(len(results.pred[0]), 4)
self.assertEqual(len(results.pred[1]), 8)
# prepare image
image_path1 = TestConstants.ZIDANE_IMAGE_PATH
image_path2 = TestConstants.BUS_IMAGE_PATH
image1 = image_path1
image2 = image_path2
# perform inference
results = yolov5.predict([image1, image2], size=1280, augment=True)
# compare
self.assertEqual(results.n, 2)
self.assertEqual(len(results.names), 80)
self.assertEqual(len(results.pred[0]), 4)
self.assertEqual(len(results.pred[1]), 8)
def __init__(self):
# self.model_path = torch.hub.load('ultralytics/yolov5', 'yolov5x', pretrained=True, force_reload=False)
self.model_path = 'weights/yolov5/yolov5x.pt'
self.device = 'cpu'
self.model = YOLOv5(self.model_path, self.device)