def test_hublike_predict(self):
import yolov5
from PIL import Image
# init yolov5s model
model_path = TestConstants.YOLOV5S_MODEL_PATH
model = yolov5.load(model_path)
# prepare image
image_path = TestConstants.ZIDANE_IMAGE_PATH
image = Image.open(image_path)
# perform inference
results = model(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
model = yolov5.load(model_path)
# prepare image
image_path = TestConstants.BUS_IMAGE_PATH
image = Image.open(image_path)
# perform inference
results = model(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
model = yolov5.load(model_path)
# 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 = model([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 find_table(dirpath):
print('test', dirpath)
print('hello')
model = yolov5.load(os.getcwd() + '/api_test/best.pt')
file_list = os.listdir(dirpath + '/img')
print(file_list)
for file_name in file_list:
img = cv2.imread(dirpath + '/img/' + file_name)
img_ori = img.copy()
img_count = 0
results = model(img)
predictions = results.pred[0]
predictions = predictions.cpu().numpy()
#boxes = predictions[:, :4] # x1, x2, y1, y2
#scores = predictions[:, 4]
#categories = predictions[:, 5]
for box in predictions:
if box[4] > 0.8:
img_result = cv2.rectangle(img, (int(box[0]), int(box[1])),
(int(box[2]), int(box[3])),
(0, 255, 120), 2)
img_count += 1
img_crop = img_ori[int(box[1]):int(box[3]) + 3,
int(box[0]):int(box[2]) + 3]
save_img = cv2.imwrite(
dirpath + '/crops/' + file_name.split('.')[0] +
str(img_count) + '.jpg', img_crop)
def test_hublike_load_model(self):
import yolov5
# init model
model_path = TestConstants.YOLOV5S_MODEL_PATH
model = yolov5.load(model_path)
# check if loaded
self.assertNotEqual(model, None)
def __init__(self, model_path: str, device: Optional[str] = None):
if device is not None and "cuda" in device and not torch.cuda.is_available(
):
raise Exception(
"Selected device='cuda', but cuda is not available to Pytorch."
)
# automatically set device if its None
elif device is None:
device = "cuda:0" if torch.cuda.is_available() else "cpu"
# load model
self.model = yolov5.load(model_path, device=device)
def load_model(self):
"""
Detection model is initialized and set to self.model.
"""
try:
import yolov5
except ImportError:
raise ImportError('Please run "pip install -U yolov5" ' "to install YOLOv5 first for YOLOv5 inference.")
# set model
try:
model = yolov5.load(self.model_path, device=self.device)
self.model = model
except Exception as e:
TypeError("model_path is not a valid yolov5 model path: ", e)
# set category_mapping
if not self.category_mapping:
category_mapping = {str(ind): category_name for ind, category_name in enumerate(self.category_names)}
self.category_mapping = category_mapping