From 435eec746770b675d3b7fe54223b9e42739f8e0a Mon Sep 17 00:00:00 2001 From: Daniel Khromov Date: Thu, 11 Feb 2021 20:54:40 +0000 Subject: [PATCH 1/4] save results with name --- models/common.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/models/common.py b/models/common.py index 7cfea01f223e..7c77193c6580 100644 --- a/models/common.py +++ b/models/common.py @@ -242,7 +242,7 @@ def __init__(self, imgs, pred, names=None): self.xywhn = [x / g for x, g in zip(self.xywh, gn)] # xywh normalized self.n = len(self.pred) - def display(self, pprint=False, show=False, save=False, render=False, save_dir=''): + def display(self, pprint=False, show=False, save=False, render=False, save_dir='', save_orig_name=False): colors = color_list() for i, (img, pred) in enumerate(zip(self.imgs, self.pred)): str = f'image {i + 1}/{len(self.pred)}: {img.shape[0]}x{img.shape[1]} ' @@ -260,7 +260,10 @@ def display(self, pprint=False, show=False, save=False, render=False, save_dir=' if show: img.show(f'image {i}') # show if save: - f = Path(save_dir) / f'results{i}.jpg' + if save_orig_name: + f = Path(save_dir) / f'results{i}.jpg' + else: + f = Path(save_dir) / f'results_{img.name}.jpg' img.save(f) # save print(f"{'Saving' * (i == 0)} {f},", end='' if i < self.n - 1 else ' done.\n') if render: @@ -272,8 +275,8 @@ def print(self): def show(self): self.display(show=True) # show results - def save(self, save_dir=''): - self.display(save=True, save_dir=save_dir) # save results + def save(self, save_dir='', save_orig_name=False): + self.display(save=True, save_dir=save_dir, save_orig_name=save_orig_name) # save results def render(self): self.display(render=True) # render results From 7e779c39e789cd56027b1cd15246cd229517678c Mon Sep 17 00:00:00 2001 From: Daniel Khromov Date: Thu, 11 Feb 2021 21:00:46 +0000 Subject: [PATCH 2/4] debug --- models/common.py | 1 + 1 file changed, 1 insertion(+) diff --git a/models/common.py b/models/common.py index 7c77193c6580..8ef0ab6f72fe 100644 --- a/models/common.py +++ b/models/common.py @@ -264,6 +264,7 @@ def display(self, pprint=False, show=False, save=False, render=False, save_dir=' f = Path(save_dir) / f'results{i}.jpg' else: f = Path(save_dir) / f'results_{img.name}.jpg' + print(f) img.save(f) # save print(f"{'Saving' * (i == 0)} {f},", end='' if i < self.n - 1 else ' done.\n') if render: From 7869b426dda0a3eba8c17692e4da797ff7dc74e5 Mon Sep 17 00:00:00 2001 From: Daniel Khromov Date: Fri, 12 Feb 2021 09:02:23 +0000 Subject: [PATCH 3/4] save original imgs names --- models/common.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/models/common.py b/models/common.py index 8ef0ab6f72fe..1e4f152b30a4 100644 --- a/models/common.py +++ b/models/common.py @@ -189,11 +189,10 @@ def forward(self, imgs, size=640, augment=False, profile=False): # 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 - + orig_imgs = imgs 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 = [], [] # image and inference shapes @@ -224,16 +223,17 @@ def forward(self, imgs, size=640, augment=False, profile=False): for i in range(n): scale_coords(shape1, y[i][:, :4], shape0[i]) - return Detections(imgs, y, self.names) + return Detections(imgs, y, self.names, orig_imgs) class Detections: # detections class for YOLOv5 inference results - def __init__(self, imgs, pred, names=None): + def __init__(self, imgs, pred, names=None, orig_imgs=None): super(Detections, self).__init__() d = pred[0].device # device gn = [torch.tensor([*[im.shape[i] for i in [1, 0, 1, 0]], 1., 1.], device=d) for im in imgs] # normalizations self.imgs = imgs # list of images as numpy arrays + self.orig_imgs = imgs # list of original names self.pred = pred # list of tensors pred[0] = (xyxy, conf, cls) self.names = names # class names self.xyxy = pred # xyxy pixels @@ -261,9 +261,9 @@ def display(self, pprint=False, show=False, save=False, render=False, save_dir=' img.show(f'image {i}') # show if save: if save_orig_name: - f = Path(save_dir) / f'results{i}.jpg' + f = Path(save_dir) / f'results{orig_imgs[i]}.jpg' else: - f = Path(save_dir) / f'results_{img.name}.jpg' + f = Path(save_dir) / f'results_{i}.jpg' print(f) img.save(f) # save print(f"{'Saving' * (i == 0)} {f},", end='' if i < self.n - 1 else ' done.\n') From 1e56487d6156eb0cb2e6db6885ebe03ef843f21a Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Fri, 12 Feb 2021 14:10:26 -0800 Subject: [PATCH 4/4] Update common.py --- models/common.py | 27 +++++++++++++-------------- 1 file changed, 13 insertions(+), 14 deletions(-) diff --git a/models/common.py b/models/common.py index 1e4f152b30a4..4f4f331da583 100644 --- a/models/common.py +++ b/models/common.py @@ -189,16 +189,18 @@ def forward(self, imgs, size=640, augment=False, profile=False): # 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 - orig_imgs = imgs + 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 = [], [] # image and inference shapes + shape0, shape1, files = [], [], [] # image and inference shapes, filenames for i, im in enumerate(imgs): if isinstance(im, str): # filename or uri im = Image.open(requests.get(im, stream=True).raw if im.startswith('http') else im) # open + 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) @@ -223,26 +225,26 @@ def forward(self, imgs, size=640, augment=False, profile=False): for i in range(n): scale_coords(shape1, y[i][:, :4], shape0[i]) - return Detections(imgs, y, self.names, orig_imgs) + return Detections(imgs, y, files, self.names) class Detections: # detections class for YOLOv5 inference results - def __init__(self, imgs, pred, names=None, orig_imgs=None): + def __init__(self, imgs, pred, files, names=None): super(Detections, self).__init__() d = pred[0].device # device gn = [torch.tensor([*[im.shape[i] for i in [1, 0, 1, 0]], 1., 1.], device=d) for im in imgs] # normalizations self.imgs = imgs # list of images as numpy arrays - self.orig_imgs = imgs # list of original names self.pred = pred # list of tensors pred[0] = (xyxy, conf, cls) self.names = names # class names + self.files = files # image filenames self.xyxy = pred # xyxy pixels self.xywh = [xyxy2xywh(x) for x in pred] # xywh pixels self.xyxyn = [x / g for x, g in zip(self.xyxy, gn)] # xyxy normalized self.xywhn = [x / g for x, g in zip(self.xywh, gn)] # xywh normalized self.n = len(self.pred) - def display(self, pprint=False, show=False, save=False, render=False, save_dir='', save_orig_name=False): + def display(self, pprint=False, show=False, save=False, render=False, save_dir=''): colors = color_list() for i, (img, pred) in enumerate(zip(self.imgs, self.pred)): str = f'image {i + 1}/{len(self.pred)}: {img.shape[0]}x{img.shape[1]} ' @@ -258,13 +260,9 @@ def display(self, pprint=False, show=False, save=False, render=False, save_dir=' if pprint: print(str.rstrip(', ')) if show: - img.show(f'image {i}') # show + img.show(self.files[i]) # show if save: - if save_orig_name: - f = Path(save_dir) / f'results{orig_imgs[i]}.jpg' - else: - f = Path(save_dir) / f'results_{i}.jpg' - print(f) + f = Path(save_dir) / self.files[i] img.save(f) # save print(f"{'Saving' * (i == 0)} {f},", end='' if i < self.n - 1 else ' done.\n') if render: @@ -276,8 +274,9 @@ def print(self): def show(self): self.display(show=True) # show results - def save(self, save_dir='', save_orig_name=False): - self.display(save=True, save_dir=save_dir, save_orig_name=save_orig_name) # save results + def save(self, save_dir='results/'): + Path(save_dir).mkdir(exist_ok=True) + self.display(save=True, save_dir=save_dir) # save results def render(self): self.display(render=True) # render results