def display(self, pprint=False, show=False, save=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]} ' if pred is not None: for c in pred[:, -1].unique(): n = (pred[:, -1] == c).sum() # detections per class str += f'{n} {self.names[int(c)]}s, ' # add to string if show or save: img = Image.fromarray(img.astype(np.uint8)) if isinstance( img, np.ndarray) else img # from np for *box, conf, cls in pred: # xyxy, confidence, class # str += '%s %.2f, ' % (names[int(cls)], conf) # label ImageDraw.Draw(img).rectangle( box, width=4, outline=colors[int(cls) % 10]) # plot if save: f = f'results{i}.jpg' str += f"saved to '{f}'" img.save(f) # save if show: img.show(f'Image {i}') # show if pprint: temp = [] for *box, conf, cls in pred: x1 = int(box[0].cpu().numpy()) y1 = int(box[1].cpu().numpy()) x2 = int(box[2].cpu().numpy()) y2 = int(box[3].cpu().numpy()) cls = int(cls.cpu().numpy()) temp.append([x1, y1, x2, y2, cls]) print(temp) return temp
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]} ' if pred is not None: for c in pred[:, -1].unique(): n = (pred[:, -1] == c).sum() # detections per class str += f"{n} {self.names[int(c)]}{'s' * (n > 1)}, " # add to string if show or save or render: for *box, conf, cls in pred: # xyxy, confidence, class label = f'{self.names[int(cls)]} {conf:.2f}' plot_one_box(box, img, label=label, color=colors[int(cls) % 10]) img = Image.fromarray(img.astype(np.uint8)) if isinstance( img, np.ndarray) else img # from np if pprint: print(str.rstrip(', ')) if show: img.show(self.files[i]) # show if save: f = self.files[i] img.save(Path(save_dir) / f) # save print(f"{'Saved' * (i == 0)} {f}", end=',' if i < self.n - 1 else f' to {save_dir}\n') if render: self.imgs[i] = np.asarray(img)
def display(self, pprint=False, show=False, save=False, render=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]} ' if pred is not None: for c in pred[:, -1].unique(): n = (pred[:, -1] == c).sum() # detections per class str += f"{n} {self.names[int(c)]}{'s' * (n > 1)}, " # add to string if show or save or render: img = Image.fromarray(img.astype(np.uint8)) if isinstance( img, np.ndarray) else img # from np for *box, conf, cls in pred: # xyxy, confidence, class # str += '%s %.2f, ' % (names[int(cls)], conf) # label ImageDraw.Draw(img).rectangle( box, width=4, outline=colors[int(cls) % 10]) # plot if pprint: print(str.rstrip(', ')) if show: img.show(f'image {i}') # show if save: f = f'results{i}.jpg' img.save(f) # save print(f"{'Saving' * (i == 0)} {f},", end='' if i < self.n - 1 else ' done.\n') if render: self.imgs[i] = np.asarray(img)
def display(self, pprint=False, show=False, save=False, render=False, save_txt=False, save_dir=''): colors = color_list() txt = "" count = 0 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]}, ' if pred is not None: for c in pred[:, -1].unique(): n = (pred[:, -1] == c).sum() # detections per class str += f"{n} {self.names[int(c)]}{'s' * (n > 1)}, " + "\n" if i == count: txt += f"{i}: {n}" + "\n" else: txt += f"{count}: 0" + "\n" txt += f"{i}: {n}" + "\n" count += 1 count += 1 nlines = txt.count("\n") if nlines != len(self.pred) and count >= len(self.pred): txt += f"{count}: 0" + "\n" if show or save or render: for *box, conf, cls in pred: # xyxy, confidence, class label = f'{self.names[int(cls)]} {conf:.2f}' plot_one_box(box, img, label=label, color=colors[int(cls) % 10]) img = Image.fromarray(img.astype(np.uint8)) if isinstance( img, np.ndarray) else img # from np if pprint: print(str.rstrip(', ')) if show: img.show(self.files[i]) # show if save: 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: self.imgs[i] = np.asarray(img) if save_txt: # nlines = txt.count("\n") # if nlines != len(self.pred): # txt += f"{i}: 0" + "\n" text_file = open("output.txt", "w") text_file.write(txt) text_file.close()
def annotate(save_img=False): images, label_path, imgsz = opt.images, opt.labels, opt.img_size if not os.path.exists(images): print('{} not found!'.format(images)) sys.exit() if not os.path.exists(label_path): print('{} not found!'.format(label_path)) sys.exit() # get label filter (if any) filt = eval(opt.filter) filt = np.array(filt) if len(filt) > 0 else None if filt is not None: prefilter = filt # get crop param (if any) crop = eval(opt.crop) crop = np.array(crop) if len(crop) > 0 else None if crop is not None: prefilter = crop # filter images BEFORE loading them if prefilter is not None: lab_files = glob.glob(f'{label_path}/*.txt') img_files = glob.glob(f'{images}/*.[jJ][pP][gG]') images = [] for lab_file in lab_files: labels = load_labels(lab_file) if len(labels) == 0: continue y = np.array(labels)[:, 0] idx = np.in1d(y, prefilter) if not idx.any(): continue _, name = path_parts(lab_file) images.extend([f for f in img_files if '.'.join(name[:-1]) in f]) random.shuffle(images) # Directories save_dir = Path( increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok)) # increment run save_dir.mkdir(parents=True, exist_ok=True) # make dir # Initialize set_logging() # Set Dataloader save_img = True dataset = LoadImages(images, img_size=imgsz) #, verbose=False) # Get names and colors names = load_list(opt.names) colors = color_list() ## process all images... n, ct = 0, 0 for path, _, img, _ in dataset: ct += 1 if ct % 10 == 0: print(f'{ct}/{len(dataset)}') _, name = path_parts(path) lab_file = os.path.join(label_path, '{}.txt'.format('.'.join(name[:-1]))) if not os.path.exists(lab_file): print('{} not found!'.format(lab_file)) continue labels = load_labels(lab_file) if len(labels) == 0: continue if filt is not None: y = np.array(labels)[:, 0] idx = np.in1d(y, filt) if not idx.any(): continue labels = labels[idx] if crop is not None: y = np.array(labels)[:, 0] idx = np.in1d(y, crop) if not idx.any(): continue labels = labels[idx] p = Path(path) # to Path save_path = str(save_dir / p.name) # img.jpg ## resize image (unless cropping)... if crop is None: q = max(img.shape[:2]) / imgsz h, w = int(img.shape[0] / q), int(img.shape[1] / q) img = cv2.resize(img, dsize=(w, h), interpolation=cv2.INTER_CUBIC) ## process all labels... first = names.copy() ## only print class name the first time for i, det in enumerate(labels): # detections per image cls, xywh = det[0], det[1:] cls_id = int(cls) xyxy = xywh2xyxy(xywh, img.shape[1], img.shape[0]) if crop is not None: n += 1 crop_path = f'{save_dir}/{cls_id}-{i}_{n}_{p.name}' crop_path = f'{save_dir}/{cls_id}-{i}_{n}+{p.name}' img_crop = crop_one_box(xyxy, img) cv2.imwrite(crop_path, img_crop) continue label = True if first and first[cls_id]: first[cls_id] = 0 else: label = False label = f'{names[cls_id]}' if label else None plot_one_box(xyxy, img, label=label, color=colors[cls_id % len(colors)], line_thickness=1) #3 # Save results (image with detections) if crop is None: cv2.imwrite(save_path, img) # s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else '' print(f"Results saved to {save_dir}")