Spaces:
Runtime error
Runtime error
File size: 2,708 Bytes
c3155e8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 | import sys
from pathlib import Path
import torch
import yolov5
'''Class for loading the Yolo-v5 inference_models
'''
class InferenceModel:
def __init__(self, model_name):
self.model_name = model_name
# path to inference_models
self.model_path = Path('./data/inference_models/{}'.format(model_name))
print(self.model_name + ' loaded')
print('cuda available: ' + str(torch.cuda.is_available()))
if torch.cuda.is_available():
print('running GPU inference..')
device_memory = {}
# get gpu with the highest memory
for i in range(torch.cuda.device_count()):
props = torch.cuda.get_device_properties(i)
device_memory[i] = props.total_memory
device_idx = max(device_memory, key=device_memory.get)
cuda = torch.device('cuda:{}'.format(device_idx))
# load inference_models into memory
try:
self.model = yolov5.load(str(self.model_path), device=str(cuda))
except Exception as e:
print(e)
print('Could not load model')
sys.exit(-1)
else:
print('running CPU inference..')
try:
self.model = yolov5.load(str(self.model_path), device='cpu')
except Exception as e:
print(e)
print('Could not load model')
sys.exit(-1)
# inference_models properties
self.model.conf = 0.50 # NMS confidence threshold
self.model.iou = 0.50 # NMS IoU threshold
self.model.classes = [0, 1] # Only show these classes
self.model.agnostic = False # NMS class-agnostic
self.model.multi_label = False # NMS multiple labels per box
self.model.max_det = 1 # maximum number of detections per image
self.model.amp = True # Automatic Mixed Precision (AMP) inference
# return prediction
def predict(self, image):
return self.model(image)
# extract items from results
@staticmethod
def get_results(results):
(bbox_x1, bbox_y1, bbox_x2, bbox_y2, class_name, confidence) = None, None, None, None, None, None
results = results.pandas().xyxy[0].to_dict(orient="records")
if results:
for result in results:
confidence = result['confidence']
class_name = result['class']
bbox_x1 = int(result['xmin'])
bbox_y1 = int(result['ymin'])
bbox_x2 = int(result['xmax'])
bbox_y2 = int(result['ymax'])
return bbox_x1, bbox_y1, bbox_x2, bbox_y2, class_name, confidence
|