Bienqq
/
car-detection-bayes
2
1
Fork 0
Code Issues Pull Requests Releases Wiki Activity

cleanup for #1114

This commit is contained in:
Glenn Jocher 2020-04-30 14:53:57 -07:00
parent fb1b5e09b2
commit 0ffbf5534e
2 changed files with 49 additions and 80 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
test.py
View File

@ -165,9 +165,9 @@ def test(cfg,
# Plot images
if batch_i < 1:
f = 'test_batch%g_gt.jpg' % batch_i # filename
plot_images(images=imgs, targets=targets, paths=paths, names=names, fname=f) # ground truth
f = 'test_batch%g_pred.jpg' % batch_i # filename
plot_images(images=imgs, targets=output_to_target(output, width, height), paths=paths, names=names, fname=f) # predictions
plot_images(imgs, targets, paths=paths, names=names, fname=f) # ground truth
f = 'test_batch%g_pred.jpg' % batch_i
plot_images(imgs, output_to_target(output, width, height), paths=paths, names=names, fname=f) # predictions
# Compute statistics
stats = [np.concatenate(x, 0) for x in zip(*stats)] # to numpy

77
utils/utils.py
View File

@ -893,7 +893,9 @@ def plot_wh_methods(): # from utils.utils import *; plot_wh_methods()
fig.savefig('comparison.png', dpi=200)
def plot_images(images, targets, paths=None, fname='images.jpg', names=None, class_labels=True, confidence_labels=True, max_size=640, max_subplots=16):
def plot_images(images, targets, paths=None, fname='images.jpg', names=None, max_size=640, max_subplots=16):
tl = 3 # line thickness
tf = max(tl - 1, 1) # font thickness
if isinstance(images, torch.Tensor):
images = images.cpu().numpy()
@ -910,17 +912,13 @@ def plot_images(images, targets, paths=None, fname='images.jpg', names=None, cla
ns = np.ceil(bs ** 0.5) # number of subplots (square)
# Check if we should resize
should_resize = False
if w > max_size or h > max_size:
scale_factor = max_size / max(h, w)
if scale_factor < 1:
h = math.ceil(scale_factor * h)
w = math.ceil(scale_factor * w)
should_resize=True
# Empty array for output
mosaic_width = int(ns*w)
mosaic_height = int(ns*h)
mosaic = 255*np.ones((mosaic_height, mosaic_width, 3), dtype=np.uint8)
mosaic = np.full((int(ns * h), int(ns * w), 3), 255, dtype=np.uint8)
# Fix class - colour map
prop_cycle = plt.rcParams['axes.prop_cycle']
@ -928,72 +926,43 @@ def plot_images(images, targets, paths=None, fname='images.jpg', names=None, cla
hex2rgb = lambda h: tuple(int(h[1 + i:1 + i + 2], 16) for i in (0, 2, 4))
color_lut = [hex2rgb(h) for h in prop_cycle.by_key()['color']]
for i, image in enumerate(images):
# e.g. if the last batch has fewer images than we expect
if i == max_subplots:
for i, img in enumerate(images):
if i == max_subplots: # if last batch has fewer images than we expect
break
block_x = int(w * (i // ns))
block_y = int(h * (i % ns))
image = image.transpose(1,2,0)
if should_resize:
image = cv2.resize(image, (w, h))
mosaic[block_y:block_y+h, block_x:block_x+w,:] = image
img = img.transpose(1, 2, 0)
if scale_factor < 1:
img = cv2.resize(img, (w, h))
mosaic[block_y:block_y + h, block_x:block_x + w, :] = img
if targets is not None:
image_targets = targets[targets[:, 0] == i]
boxes = xywh2xyxy(image_targets[:, 2:6]).T
classes = image_targets[:, 1].astype('int')
# Check if we have object confidences (gt vs pred)
confidences = None
if image_targets.shape[1] > 6:
confidences = image_targets[:,6]
gt = image_targets.shape[1] == 6 # ground truth if no conf column
conf = None if gt else image_targets[:, 6] # check for confidence presence (gt vs pred)
boxes[[0, 2]] *= w
boxes[[0, 2]] += block_x
boxes[[1, 3]] *= h
boxes[[1, 3]] += block_y
for j, box in enumerate(boxes.T):
color = color_lut[int(classes[j]) % len(color_lut)]
box = box.astype(int)
cv2.rectangle(mosaic, (box[0], box[1]), (box[2], box[3]), color, thickness=2)
# Draw class label
if class_labels and max_size > 250:
label = str(classes[j]) if names is None else names[classes[j]]
if confidences is not None and confidence_labels:
label += " {:1.2f}".format(confidences[j])
font_scale = 0.4/10 * min(20, h * 0.05)
font_thickness = 2 if max(w, h) > 320 else 1
label_size, _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, font_scale, font_thickness)
cv2.rectangle(mosaic, (box[0], box[1]), (box[0]+label_size[0], box[1]-label_size[1]), color, thickness=-1)
cv2.putText(mosaic, label, (box[0], box[1]), fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=font_scale, thickness=font_thickness, color=(255,255,255))
cls = int(classes[j])
color = color_lut[cls % len(color_lut)]
cls = names[cls] if names else cls
if gt or conf[j] > 0.3: # 0.3 conf thresh
label = '%s' % cls if gt else '%s %.1f' % (cls, conf[j])
plot_one_box(box, mosaic, label=label, color=color, line_thickness=tl)
# Draw image filename labels
if paths is not None:
# Trim to 40 chars
label = os.path.basename(paths[i])[:40]
# Empirical calculation to fit label
# 0.4 is at most (13, 10) px per char at thickness = 1
# Fit label to 20px high, or shrink if it would be too big
max_font_scale = (w/len(label))*(0.4/8)
font_scale = min(0.4 * 20/8.5, max_font_scale)
font_thickness = 1
label_size, baseline = cv2.getTextSize(label, cv2.FONT_HERSHEY_DUPLEX, font_scale, font_thickness)
cv2.rectangle(mosaic, (block_x+5, block_y+label_size[1]+baseline+5), (block_x+label_size[0]+5, block_y), 0, thickness=-1)
cv2.putText(mosaic, label, (block_x+5, block_y+label_size[1]+5), cv2.FONT_HERSHEY_DUPLEX, font_scale, (255,255,255), font_thickness)
label = os.path.basename(paths[i])[:40] # trim to 40 char
t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
cv2.putText(mosaic, label, (block_x + 5, block_y + t_size[1] + 5), 0, tl / 3, [220, 220, 220], thickness=tf,
lineType=cv2.LINE_AA)
# Image border
cv2.rectangle(mosaic, (block_x, block_y), (block_x + w, block_y + h), (255, 255, 255), thickness=3)