Spaces:
Runtime error
Runtime error
| 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 | |
| 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 | |