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}")
