import math # from YoloVal import DetectionValidatorEnsemble from argparse import ArgumentParser from collections import deque import cv2 import numpy as np import torch from torch import nn from ultralytics import YOLO from ultralytics.engine.results import Results from ultralytics.models.yolo.detect import DetectionValidator from ultralytics.nn.autobackend import AutoBackend from ultralytics.utils import ops, nms def do_rectangles_overlap(rect1, rect2, overlap_threshold=0.5): # Rect1 coords x1_min, y1_min, x1_max, y1_max = rect1 # Rect2 coords x2_min, y2_min, x2_max, y2_max = rect2 # Check if one rectangle is to the left of the other if x1_max < x2_min or x2_max < x1_min: return False # Check if one rectangle is above the other if y1_max < y2_min or y2_max < y1_min: return False # Find the area of the first rectangle area_rect1 = (x1_max - x1_min) * (y1_max - y1_min) area_rect2 = (x2_max - x2_min) * (y2_max - y2_min) # Find the coordinates of the intersection rectangle inter_x_min = max(x1_min, x2_min) inter_x_max = min(x1_max, x2_max) inter_y_min = max(y1_min, y2_min) inter_y_max = min(y1_max, y2_max) # Check if there is no intersection if inter_x_max <= inter_x_min or inter_y_max <= inter_y_min: return False # Calculate the area of the intersection rectangle inter_area = (inter_x_max - inter_x_min) * (inter_y_max - inter_y_min) # Calculate the percentage of overlap relative to both rectangles overlap_percentage_1 = inter_area / area_rect1 overlap_percentage_2 = inter_area / area_rect2 # Check for complete containment contained = ((x1_min <= x2_min <= x1_max and x1_min <= x2_max <= x1_max) and (y1_min <= y2_min <= y1_max and y1_min <= y2_max <= y1_max)) or \ ((x2_min <= x1_min <= x2_max and x2_min <= x1_max <= x2_max) and (y2_min <= y1_min <= y2_max and y2_min <= y1_max <= y2_max)) # Return True if the overlap meets the threshold return overlap_percentage_1 >= overlap_threshold or overlap_percentage_2 >= overlap_threshold or contained import spaces class YoloEnsemble: def __init__(self, weights: list[str]): self.models = [YOLO(weight) for weight in weights] @spaces.GPU(duration=10) def predict(self, img_path: str, conf: float = 0.25, verbose: bool = True): import torch import numpy as np import random seed = 42 torch.manual_seed(seed) np.random.seed(seed) random.seed(seed) torch.cuda.manual_seed_all(seed) # if you are using multi-GPU. # For full reproducibility, you might also need this torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False predictions = [_model(img_path, conf=conf, verbose=verbose) for _model in self.models] if len(self.models) > 1: return self.ensemble(predictions) return predictions[0] def ensemble(self, predictions: list): hits = None orig_shape = None names = None orig_img = None path = None speed = None for results in predictions: for result in results: _hits = result.boxes.data.unsqueeze(dim=0) if hits is None: hits = _hits else: hits = torch.cat((hits, _hits), dim=1) if orig_shape is None: orig_shape = result.orig_shape names = result.names orig_img = result.orig_img path = result.path speed = result.speed # hits = hits.unsqueeze(dim=0) nms_hits = nms.non_max_suppression(hits, conf_thres=0.25, classes=[0, 1, 2, 3, 4, 5, 6]) boxes = deque(nms_hits[0].tolist()) non_overlapping_boxes = [] while len(boxes) > 0: box = boxes.popleft() overlappers = [box] rem = [] for i, b in enumerate(boxes): if do_rectangles_overlap(box[:4], b[:4]): overlappers.append(b) rem.append(i) for _i, _ in enumerate(rem): del boxes[_ - _i] keep_box = sorted(overlappers, key=lambda x: x[4], reverse=True)[0] non_overlapping_boxes.append(keep_box) if len(non_overlapping_boxes) == 0: return [Results(names=names, orig_img=orig_img, path=path, speed=speed)] # result = Results(boxes=torch.Tensor(non_overlapping_boxes).to(nms_hits[0].get_device()), names=names, orig_img=orig_img, path=path, speed=speed) return [Results(boxes=torch.Tensor(non_overlapping_boxes), #.to(nms_hits[0].get_device()), names=names, orig_img=orig_img, path=path, speed=speed)] class YoloEnsembleAutoBackend: def __init__(self, weights: list[str], val=False, **kwargs): if isinstance(weights, list): self.models = [ # AutoBackend( # weights=weight, # device=kwargs.get('device', None), # dnn=kwargs.get('dnn', False), # data=kwargs.get('data', None), # fp16=kwargs.get('fp16', False), # ) for weight in weights YOLO(weight) for weight in weights ] else: self.models = [ AutoBackend( weights=weights, device=kwargs.get('device', None), dnn=kwargs.get('dnn', False), data=kwargs.get('data', None), fp16=kwargs.get('fp16', False), ) ] model = AutoBackend( weights=weights[0], device=kwargs.get('device', None), dnn=kwargs.get('dnn', False), data=kwargs.get('data', None), fp16=kwargs.get('fp16', False), ) # self.models[0].val() self.device = kwargs.get('device', None) self.fp16 = model.fp16 self.stride = model.stride self.pt = model.pt self.jit = model.jit self.engine = model.engine self.val = val self.names = model.names def warmup(self, imgsz=(1, 3, 640, 640)): pass def eval(self): for model in self.models: model.eval() def predict(self, imgs, conf=0.25, verbose=True): predictions = [_model(imgs, conf=conf, verbose=verbose) for _model in self.models] predictions = [list(x) for x in zip(*predictions)] if len(self.models) > 1: # return self.ensemble([torch.cat([p[0] for p in predictions], 1)]) return self.ensemble2(predictions) if not self.val: return predictions[0] return predictions[0] def ensemble(self, predictions: list): final_preds = [] device = None for ip, results in enumerate(predictions): for ir, result in enumerate(results): device = result.device _array = deque(result.cpu().tolist()) non_overlapping_boxes = [] while len(_array) > 0: box = _array.popleft() overlappers = [box] rem = [] for i, b in enumerate(_array): if do_rectangles_overlap(box[:4], b[:4]): overlappers.append(b) rem.append(i) for _i, _ in enumerate(rem): del _array[_ - _i] keep_box = sorted(overlappers, key=lambda x: x[4], reverse=True)[0] non_overlapping_boxes.append(keep_box) repeat = int(math.ceil(300 / len(non_overlapping_boxes))) non_overlapping_boxes = non_overlapping_boxes * repeat final_preds.append(non_overlapping_boxes[:300]) _new_preds = torch.tensor(final_preds, device=device) return _new_preds def ensemble2(self, predictions: list): final_preds = [] device = None preds = [] for ip, prediction in enumerate(predictions): # for image i model_preds = [] for ir, result in enumerate(prediction): # for model r's prediction on image i if not device: device = result.boxes.xyxy.device boxes = np.array(result.boxes.xyxy.cpu().tolist()) if len(boxes) == 0: continue _cls = np.array(result.boxes.cls.cpu().tolist()) _cls = _cls.reshape(-1, 1) _conf = np.array(result.boxes.conf.cpu().tolist()) _conf = _conf.reshape(-1, 1) try: np.hstack((boxes, _conf, _cls)) except: breakpoint() boxes = np.hstack((boxes, _conf, _cls)) boxes = boxes.tolist() model_preds.extend(boxes) preds.append(model_preds) for ip, pred in enumerate(preds): # for image i _array = deque(pred) non_overlapping_boxes = [] while len(_array) > 0: box = _array.popleft() overlappers = [box] rem = [] for i, b in enumerate(_array): if do_rectangles_overlap(box[:4], b[:4]): overlappers.append(b) rem.append(i) for _i, _ in enumerate(rem): del _array[_ - _i] keep_box = sorted(overlappers, key=lambda x: x[4], reverse=True)[0] non_overlapping_boxes.append(keep_box) # increase to 100 if len(non_overlapping_boxes) != 0: repeat = int(math.ceil(100 / len(non_overlapping_boxes))) non_overlapping_boxes = non_overlapping_boxes * repeat final_preds.append(non_overlapping_boxes[:100]) _new_preds = torch.tensor(final_preds, device=device) # for ip, results in enumerate(predictions): # per_img_preds = [] # for ir, result in enumerate(results): # device = result.boxes.xyxy.device # boxes = np.array(result.boxes.xyxy.cpu().tolist()) # _cls = np.array(result.boxes.cls.cpu().tolist()) # _cls = _cls.reshape(-1, 1) # _conf = np.array(result.boxes.conf.cpu().tolist()) # _conf = _conf.reshape(-1, 1) # # boxes = np.hstack((boxes, _conf, _cls)) # boxes = boxes.tolist() # # _array = deque(boxes) # non_overlapping_boxes = [] # while len(_array) > 0: # box = _array.popleft() # overlappers = [box] # rem = [] # for i, b in enumerate(_array): # if do_rectangles_overlap(box[:4], b[:4]): # overlappers.append(b) # rem.append(i) # for _i, _ in enumerate(rem): # del _array[_ - _i] # keep_box = sorted(overlappers, key=lambda x: x[4], reverse=True)[0] # non_overlapping_boxes.append(keep_box) # # # repeat = int(math.ceil(100 / len(non_overlapping_boxes))) # # non_overlapping_boxes = non_overlapping_boxes * repeat # # final_preds.append(non_overlapping_boxes[:100]) # per_img_preds.extend(non_overlapping_boxes) # # # _new_preds = torch.tensor(final_preds, device=device) return _new_preds class YoloPreprocess(nn.Module): def __init__(self): super(YoloPreprocess, self).__init__() def pre_transform(self, img: np.ndarray): img = img shape = len(img), len(img[0]) new_shape = [1280, 1280] r = min(new_shape[0] / shape[0], new_shape[1] / shape[1]) new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r)) dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # wh padding dw, dh = np.mod(dw, 32), np.mod(dh, 32) dw /= 2 dh /= 2 if shape[::-1] != new_unpad: # resize img = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR) top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1)) left, right = int(round(dw - 0.1)), int(round(dw + 0.1)) img = cv2.copyMakeBorder( img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(114, 114, 114) ) return [img] def forward(self, im): im = np.stack(self.pre_transform(im)) im = im[..., ::-1].transpose((0, 3, 1, 2)) # BGR to RGB, BHWC to BCHW, (n, 3, h, w) im = np.ascontiguousarray(im) # contiguous # im = torch.from_numpy(im) _im = im / 255 return _im if __name__ == '__main__': parser = ArgumentParser() parser.add_argument('--weights', nargs='+', help="Model paths", required=True) args = parser.parse_args() # img = cv2.imread('askubuntu2.png') # x = YoloPreprocess() # x(img) # model = YoloEnsemble(args.weights) # model = YOLO('./train16.pt').to('cuda') # results = model.predict(['askubuntu2.png'], conf=0.7) # for result in results: # boxes = result.boxes # Boxes object for bounding box outputs # masks = result.masks # Masks object for segmentation masks outputs # keypoints = result.keypoints # Keypoints object for pose outputs # probs = result.probs # Probs object for classification outputs # # result.show() # display to screen # result.save(filename='result.jpg') args = dict(model='./train16.pt', data='dataset/data.yaml') validator = DetectionValidator(args=args) validator()