Agnostic NMS when loading from PyTorch Hub

[c-axel]

Is there a way to use agnostic NMS when loading a custom model from PyTorch Hub? Trying to replicate what I get when using detect.py with --agnostic.

I can copy the approach in utils.general.non_max_suppression(agnostic=True) in my script and apply it on each result, but curious if there is a more elegant way.

[glenn-jocher]

@c-axel this is not available currently, but if you'd like to submit a PR you can introduce this in the autoShape() class:

yolov5/models/common.py

Lines 168 to 229 in e931b9d

	class autoShape(nn.Module):
	# input-robust model wrapper for passing cv2/np/PIL/torch inputs. Includes preprocessing, inference and NMS
	img_size = 640 # inference size (pixels)
	conf = 0.25 # NMS confidence threshold
	iou = 0.45 # NMS IoU threshold
	classes = None # (optional list) filter by class
	def __init__(self, model):
	super(autoShape, self).__init__()
	self.model = model.eval()
	def autoshape(self):
	print('autoShape already enabled, skipping... ') # model already converted to model.autoshape()
	return self
	def forward(self, imgs, size=640, augment=False, profile=False):
	# Inference from various sources. For height=720, width=1280, RGB images example inputs are:
	# filename: imgs = 'data/samples/zidane.jpg'
	# URI: = 'https://github.com/ultralytics/yolov5/releases/download/v1.0/zidane.jpg'
	# OpenCV: = cv2.imread('image.jpg')[:,:,::-1] # HWC BGR to RGB x(720,1280,3)
	# PIL: = Image.open('image.jpg') # HWC x(720,1280,3)
	# numpy: = np.zeros((720,1280,3)) # HWC
	# torch: = torch.zeros(16,3,720,1280) # BCHW
	# multiple: = [Image.open('image1.jpg'), Image.open('image2.jpg'), ...] # list of images
	p = next(self.model.parameters()) # for device and type
	if isinstance(imgs, torch.Tensor): # torch
	return self.model(imgs.to(p.device).type_as(p), augment, profile) # inference
	# Pre-process
	n, imgs = (len(imgs), imgs) if isinstance(imgs, list) else (1, [imgs]) # number of images, list of images
	shape0, shape1, files = [], [], [] # image and inference shapes, filenames
	for i, im in enumerate(imgs):
	if isinstance(im, str): # filename or uri
	im, f = Image.open(requests.get(im, stream=True).raw if im.startswith('http') else im), im # open
	im.filename = f # for uri
	files.append(Path(im.filename).with_suffix('.jpg').name if isinstance(im, Image.Image) else f'image{i}.jpg')
	im = np.array(im) # to numpy
	if im.shape[0] < 5: # image in CHW
	im = im.transpose((1, 2, 0)) # reverse dataloader .transpose(2, 0, 1)
	im = im[:, :, :3] if im.ndim == 3 else np.tile(im[:, :, None], 3) # enforce 3ch input
	s = im.shape[:2] # HWC
	shape0.append(s) # image shape
	g = (size / max(s)) # gain
	shape1.append([y * g for y in s])
	imgs[i] = im # update
	shape1 = [make_divisible(x, int(self.stride.max())) for x in np.stack(shape1, 0).max(0)] # inference shape
	x = [letterbox(im, new_shape=shape1, auto=False)[0] for im in imgs] # pad
	x = np.stack(x, 0) if n > 1 else x[0][None] # stack
	x = np.ascontiguousarray(x.transpose((0, 3, 1, 2))) # BHWC to BCHW
	x = torch.from_numpy(x).to(p.device).type_as(p) / 255. # uint8 to fp16/32
	# Inference
	with torch.no_grad():
	y = self.model(x, augment, profile)[0] # forward
	y = non_max_suppression(y, conf_thres=self.conf, iou_thres=self.iou, classes=self.classes) # NMS
	# Post-process
	for i in range(n):
	scale_coords(shape1, y[i][:, :4], shape0[i])
	return Detections(imgs, y, files, self.names)

[c-axel]

Understood, thanks for the quick reply and thank you for yolov5 it is awesome.

[glenn-jocher]

Thanks, we've tried to make it better by listening to feedback and implementing many small ideas just like this one!