apoorva/greenhouse
2
1
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Initial commit

Browse Source
This commit is contained in:
Glenn Jocher
2018-08-26 10:51:39 +02:00
commit c3731591af
20 changed files with 7285 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
utils/datasets.py
Executable
+284
View File
@@ -0,0 +1,284 @@
import glob
import math
import os
import random
import cv2
import numpy as np
import torch
# from torch.utils.data import Dataset
from utils.utils import xyxy2xywh
class ImageFolder(): # for eval-only
def __init__(self, path, batch_size=1, img_size=416):
if os.path.isdir(path):
self.files = sorted(glob.glob('%s/*.*' % path))
elif os.path.isfile(path):
self.files = [path]
self.nF = len(self.files) # number of image files
self.nB = math.ceil(self.nF / batch_size) # number of batches
self.batch_size = batch_size
self.height = img_size
assert self.nF > 0, 'No images found in path %s' % path
# RGB normalization values
# self.rgb_mean = np.array([60.134, 49.697, 40.746], dtype=np.float32).reshape((3, 1, 1))
# self.rgb_std = np.array([29.99, 24.498, 22.046], dtype=np.float32).reshape((3, 1, 1))
def __iter__(self):
self.count = -1
return self
def __next__(self):
self.count += 1
if self.count == self.nB:
raise StopIteration
img_path = self.files[self.count]
# Read image
img = cv2.imread(img_path) # BGR
# Padded resize
img, _, _, _ = resize_square(img, height=self.height, color=(127.5, 127.5, 127.5))
# Normalize RGB
img = img[:, :, ::-1].transpose(2, 0, 1)
img = np.ascontiguousarray(img, dtype=np.float32)
# img -= self.rgb_mean
# img /= self.rgb_std
img /= 255.0
return [img_path], img
def __len__(self):
return self.nB # number of batches
class ListDataset(): # for training
def __init__(self, path, batch_size=1, img_size=608):
self.path = path
#self.img_files = sorted(glob.glob('%s/*.*' % path))
with open(path, 'r') as file:
self.img_files = file.readlines()
self.img_files = [path.replace('\n', '').replace('/images','/Users/glennjocher/Downloads/DATA/coco/images') for path in self.img_files]
self.label_files = [path.replace('images', 'labels').replace('.png', '.txt').replace('.jpg', '.txt') for path in
self.img_files]
self.nF = len(self.img_files) # number of image files
self.nB = math.ceil(self.nF / batch_size) # number of batches
self.batch_size = batch_size
#assert self.nB > 0, 'No images found in path %s' % path
self.height = img_size
# RGB normalization values
# self.rgb_mean = np.array([60.134, 49.697, 40.746], dtype=np.float32).reshape((1, 3, 1, 1))
# self.rgb_std = np.array([29.99, 24.498, 22.046], dtype=np.float32).reshape((1, 3, 1, 1))
def __iter__(self):
self.count = -1
# self.shuffled_vector = np.random.permutation(self.nF) # shuffled vector
self.shuffled_vector = np.arange(self.nF)
return self
def __next__(self):
self.count += 1
if self.count == self.nB:
raise StopIteration
ia = self.count * self.batch_size
ib = min((self.count + 1) * self.batch_size, self.nF)
height = self.height
img_all = []
labels_all = []
for index, files_index in enumerate(range(ia, ib)):
img_path = self.img_files[self.shuffled_vector[files_index]]
label_path = self.label_files[self.shuffled_vector[files_index]]
img = cv2.imread(img_path) # BGR
if img is None:
continue
augment_hsv = False
if augment_hsv:
# SV augmentation by 50%
fraction = 0.50
img_hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
S = img_hsv[:, :, 1].astype(np.float32)
V = img_hsv[:, :, 2].astype(np.float32)
a = (random.random() * 2 - 1) * fraction + 1
S *= a
if a > 1:
np.clip(S, a_min=0, a_max=255, out=S)
a = (random.random() * 2 - 1) * fraction + 1
V *= a
if a > 1:
np.clip(V, a_min=0, a_max=255, out=V)
img_hsv[:, :, 1] = S.astype(np.uint8)
img_hsv[:, :, 2] = V.astype(np.uint8)
cv2.cvtColor(img_hsv, cv2.COLOR_HSV2BGR, dst=img)
h, w, _ = img.shape
img, ratio, padw, padh = resize_square(img, height=height, color=(127.5, 127.5, 127.5))
# Load labels
if os.path.isfile(label_path):
labels0 = np.loadtxt(label_path, dtype=np.float32).reshape(-1, 5)
# Normalized xywh to pixel xyxy format
labels = labels0.copy()
labels[:, 1] = ratio * w * (labels0[:, 1] - labels0[:, 3] / 2) + padw
labels[:, 2] = ratio * h * (labels0[:, 2] - labels0[:, 4] / 2) + padh
labels[:, 3] = ratio * w * (labels0[:, 1] + labels0[:, 3] / 2) + padw
labels[:, 4] = ratio * h * (labels0[:, 2] + labels0[:, 4] / 2) + padh
else:
labels = np.array([])
# Augment image and labels
# img, labels, M = random_affine(img, targets=labels, degrees=(-5, 5), translate=(0.1, 0.1), scale=(0.8, 1.2)) # RGB
plotFlag = False
if plotFlag:
import matplotlib.pyplot as plt
plt.subplot(4, 4, index + 1).imshow(img[:, :, ::-1])
plt.plot(labels[:, [1, 3, 3, 1, 1]].T, labels[:, [2, 2, 4, 4, 2]].T, '.-')
nL = len(labels)
if nL > 0:
# convert xyxy to xywh
labels[:, 1:5] = xyxy2xywh(labels[:, 1:5].copy()) / height
# random left-right flip
lr_flip = False
if lr_flip & (random.random() > 0.5):
img = np.fliplr(img)
if nL > 0:
labels[:, 1] = 1 - labels[:, 1]
# random up-down flip
ud_flip = False
if ud_flip & (random.random() > 0.5):
img = np.flipud(img)
if nL > 0:
labels[:, 2] = 1 - labels[:, 2]
img_all.append(img)
labels_all.append(torch.from_numpy(labels))
# Normalize
img_all = np.stack(img_all)[:, :, :, ::-1].transpose(0, 3, 1, 2) # BGR to RGB and cv2 to pytorch
img_all = np.ascontiguousarray(img_all, dtype=np.float32)
# img_all -= self.rgb_mean
# img_all /= self.rgb_std
img_all /= 255.0
return torch.from_numpy(img_all), labels_all
def __len__(self):
return self.nB # number of batches
def resize_square(img, height=416, color=(0, 0, 0)): # resize a rectangular image to a padded square
shape = img.shape[:2] # shape = [height, width]
ratio = float(height) / max(shape)
new_shape = [round(shape[0] * ratio), round(shape[1] * ratio)]
dw = height - new_shape[1] # width padding
dh = height - new_shape[0] # height padding
top, bottom = dh // 2, dh - (dh // 2)
left, right = dw // 2, dw - (dw // 2)
img = cv2.resize(img, (new_shape[1], new_shape[0]), interpolation=cv2.INTER_AREA)
return cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color), ratio, dw // 2, dh // 2
def random_affine(img, targets=None, degrees=(-10, 10), translate=(.1, .1), scale=(.9, 1.1), shear=(-3, 3),
borderValue=(0, 0, 0)):
# torchvision.transforms.RandomAffine(degrees=(-10, 10), translate=(.1, .1), scale=(.9, 1.1), shear=(-10, 10))
# https://medium.com/uruvideo/dataset-augmentation-with-random-homographies-a8f4b44830d4
border = 0 # width of added border (optional)
height = max(img.shape[0], img.shape[1]) + border * 2
# Rotation and Scale
R = np.eye(3)
a = random.random() * (degrees[1] - degrees[0]) + degrees[0]
# a += random.choice([-180, -90, 0, 90]) # random 90deg rotations added to small rotations
s = random.random() * (scale[1] - scale[0]) + scale[0]
R[:2] = cv2.getRotationMatrix2D(angle=a, center=(img.shape[1] / 2, img.shape[0] / 2), scale=s)
# Translation
T = np.eye(3)
T[0, 2] = (random.random() * 2 - 1) * translate[0] * img.shape[0] + border # x translation (pixels)
T[1, 2] = (random.random() * 2 - 1) * translate[1] * img.shape[1] + border # y translation (pixels)
# Shear
S = np.eye(3)
S[0, 1] = math.tan((random.random() * (shear[1] - shear[0]) + shear[0]) * math.pi / 180) # x shear (deg)
S[1, 0] = math.tan((random.random() * (shear[1] - shear[0]) + shear[0]) * math.pi / 180) # y shear (deg)
M = S @ T @ R # ORDER IS IMPORTANT HERE!!
imw = cv2.warpPerspective(img, M, dsize=(height, height), flags=cv2.INTER_LINEAR,
borderValue=borderValue) # BGR order (YUV-equalized BGR means)
# Return warped points also
if targets is not None:
if len(targets) > 0:
n = targets.shape[0]
points = targets[:, 1:5].copy()
area0 = (points[:, 2] - points[:, 0]) * (points[:, 3] - points[:, 1])
# warp points
xy = np.ones((n * 4, 3))
xy[:, :2] = points[:, [0, 1, 2, 3, 0, 3, 2, 1]].reshape(n * 4, 2) # x1y1, x2y2, x1y2, x2y1
xy = (xy @ M.T)[:, :2].reshape(n, 8)
# create new boxes
x = xy[:, [0, 2, 4, 6]]
y = xy[:, [1, 3, 5, 7]]
xy = np.concatenate((x.min(1), y.min(1), x.max(1), y.max(1))).reshape(4, n).T
# apply angle-based reduction
radians = a * math.pi / 180
reduction = max(abs(math.sin(radians)), abs(math.cos(radians))) ** 0.5
x = (xy[:, 2] + xy[:, 0]) / 2
y = (xy[:, 3] + xy[:, 1]) / 2
w = (xy[:, 2] - xy[:, 0]) * reduction
h = (xy[:, 3] - xy[:, 1]) * reduction
xy = np.concatenate((x - w / 2, y - h / 2, x + w / 2, y + h / 2)).reshape(4, n).T
# reject warped points outside of image
np.clip(xy, 0, height, out=xy)
w = xy[:, 2] - xy[:, 0]
h = xy[:, 3] - xy[:, 1]
area = w * h
ar = np.maximum(w / (h + 1e-16), h / (w + 1e-16))
i = (w > 4) & (h > 4) & (area / area0 > 0.1) & (ar < 10)
targets = targets[i]
targets[:, 1:5] = xy[i]
return imw, targets, M
else:
return imw
def convert_tif2bmp(p='/Users/glennjocher/Downloads/DATA/xview/val_images_bmp'):
import glob
import cv2
files = sorted(glob.glob('%s/*.tif' % p))
for i, f in enumerate(files):
print('%g/%g' % (i + 1, len(files)))
img = cv2.imread(f)
cv2.imwrite(f.replace('.tif', '.bmp'), img)
os.system('rm -rf ' + f)
utils/gcp.sh
+12
View File
@@ -0,0 +1,12 @@
#!/usr/bin/env bash
# Start
sudo rm -rf yolov3 && git clone https://github.com/ultralytics/yolov3 && cd yolov3 && python3 train.py -img_size 416 -epochs 999
# Resume
cd yolov3 && python3 train.py -img_size 416 -resume 1
# Detect
gsutil cp gs://ultralytics/fresh9_5_e201.pt yolov3/checkpoints
cd yolov3 && python3 detect.py
utils/parse_config.py
+36
View File
@@ -0,0 +1,36 @@
def parse_model_config(path):
"""Parses the yolo-v3 layer configuration file and returns module definitions"""
file = open(path, 'r')
lines = file.read().split('\n')
lines = [x for x in lines if x and not x.startswith('#')]
lines = [x.rstrip().lstrip() for x in lines] # get rid of fringe whitespaces
module_defs = []
for line in lines:
if line.startswith('['): # This marks the start of a new block
module_defs.append({})
module_defs[-1]['type'] = line[1:-1].rstrip()
if module_defs[-1]['type'] == 'convolutional':
module_defs[-1]['batch_normalize'] = 0
else:
key, value = line.split("=")
value = value.strip()
module_defs[-1][key.rstrip()] = value.strip()
return module_defs
def parse_data_config(path):
"""Parses the data configuration file"""
options = dict()
options['gpus'] = '0,1,2,3'
options['num_workers'] = '10'
with open(path, 'r') as fp:
lines = fp.readlines()
for line in lines:
line = line.strip()
if line == '' or line.startswith('#'):
continue
key, value = line.split('=')
options[key.strip()] = value.strip()
return options
utils/utils.py
Executable
+372
View File
@@ -0,0 +1,372 @@
import random
import cv2
import numpy as np
import torch
import torch.nn.functional as F
# set printoptions
torch.set_printoptions(linewidth=1320, precision=5, profile='long')
np.set_printoptions(linewidth=320, formatter={'float_kind': '{11.5g}'.format}) # format short g, %precision=5
def load_classes(path):
"""
Loads class labels at 'path'
"""
fp = open(path, "r")
names = fp.read().split("\n")[:-1]
return names
def modelinfo(model):
nparams = sum(x.numel() for x in model.parameters())
ngradients = sum(x.numel() for x in model.parameters() if x.requires_grad)
print('\n%4s %70s %9s %12s %20s %12s %12s' % ('', 'name', 'gradient', 'parameters', 'shape', 'mu', 'sigma'))
for i, (name, p) in enumerate(model.named_parameters()):
name = name.replace('module_list.', '')
print('%4g %70s %9s %12g %20s %12g %12g' % (
i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std()))
print('\n%g layers, %g parameters, %g gradients' % (i + 1, nparams, ngradients))
def xview_class_weights(indices): # weights of each class in the training set, normalized to mu = 1
weights = 1 / torch.FloatTensor(
[74, 364, 713, 71, 2925, 209767, 6925, 1101, 3612, 12134, 5871, 3640, 860, 4062, 895, 149, 174, 17, 1624, 1846,
125, 122, 124, 662, 1452, 697, 222, 190, 786, 200, 450, 295, 79, 205, 156, 181, 70, 64, 337, 1352, 336, 78,
628, 841, 287, 83, 702, 1177, 313865, 195, 1081, 882, 1059, 4175, 123, 1700, 2317, 1579, 368, 85])
weights /= weights.sum()
return weights[indices]
def plot_one_box(x, im, color=None, label=None, line_thickness=None):
tl = line_thickness or round(0.003 * max(im.shape[0:2])) # line thickness
color = color or [random.randint(0, 255) for _ in range(3)]
c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))
cv2.rectangle(im, c1, c2, color, thickness=tl)
if label:
tf = max(tl - 1, 1) # font thickness
t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]
c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
cv2.rectangle(im, c1, c2, color, -1) # filled
cv2.putText(im, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)
def weights_init_normal(m):
classname = m.__class__.__name__
if classname.find('Conv') != -1:
torch.nn.init.normal_(m.weight.data, 0.0, 0.03)
elif classname.find('BatchNorm2d') != -1:
torch.nn.init.normal_(m.weight.data, 1.0, 0.03)
torch.nn.init.constant_(m.bias.data, 0.0)
def xyxy2xywh(box):
xywh = np.zeros(box.shape)
xywh[:, 0] = (box[:, 0] + box[:, 2]) / 2
xywh[:, 1] = (box[:, 1] + box[:, 3]) / 2
xywh[:, 2] = box[:, 2] - box[:, 0]
xywh[:, 3] = box[:, 3] - box[:, 1]
return xywh
def compute_ap(recall, precision):
""" Compute the average precision, given the recall and precision curves.
Code originally from https://github.com/rbgirshick/py-faster-rcnn.
# Arguments
recall: The recall curve (list).
precision: The precision curve (list).
# Returns
The average precision as computed in py-faster-rcnn.
"""
# correct AP calculation
# first append sentinel values at the end
mrec = np.concatenate(([0.], recall, [1.]))
mpre = np.concatenate(([0.], precision, [0.]))
# compute the precision envelope
for i in range(mpre.size - 1, 0, -1):
mpre[i - 1] = np.maximum(mpre[i - 1], mpre[i])
# to calculate area under PR curve, look for points
# where X axis (recall) changes value
i = np.where(mrec[1:] != mrec[:-1])[0]
# and sum (\Delta recall) * prec
ap = np.sum((mrec[i + 1] - mrec[i]) * mpre[i + 1])
return ap
def bbox_iou(box1, box2, x1y1x2y2=True):
# if len(box1.shape) == 1:
# box1 = box1.reshape(1, 4)
"""
Returns the IoU of two bounding boxes
"""
if x1y1x2y2:
# Get the coordinates of bounding boxes
b1_x1, b1_y1, b1_x2, b1_y2 = box1[:, 0], box1[:, 1], box1[:, 2], box1[:, 3]
b2_x1, b2_y1, b2_x2, b2_y2 = box2[:, 0], box2[:, 1], box2[:, 2], box2[:, 3]
else:
# Transform from center and width to exact coordinates
b1_x1, b1_x2 = box1[:, 0] - box1[:, 2] / 2, box1[:, 0] + box1[:, 2] / 2
b1_y1, b1_y2 = box1[:, 1] - box1[:, 3] / 2, box1[:, 1] + box1[:, 3] / 2
b2_x1, b2_x2 = box2[:, 0] - box2[:, 2] / 2, box2[:, 0] + box2[:, 2] / 2
b2_y1, b2_y2 = box2[:, 1] - box2[:, 3] / 2, box2[:, 1] + box2[:, 3] / 2
# get the corrdinates of the intersection rectangle
inter_rect_x1 = torch.max(b1_x1, b2_x1)
inter_rect_y1 = torch.max(b1_y1, b2_y1)
inter_rect_x2 = torch.min(b1_x2, b2_x2)
inter_rect_y2 = torch.min(b1_y2, b2_y2)
# Intersection area
inter_area = torch.clamp(inter_rect_x2 - inter_rect_x1, 0) * torch.clamp(inter_rect_y2 - inter_rect_y1, 0)
# Union Area
b1_area = (b1_x2 - b1_x1) * (b1_y2 - b1_y1)
b2_area = (b2_x2 - b2_x1) * (b2_y2 - b2_y1)
return inter_area / (b1_area + b2_area - inter_area + 1e-16)
def build_targets(pred_boxes, pred_conf, pred_cls, target, anchor_wh, nA, nC, nG, requestPrecision):
"""
returns nGT, nCorrect, tx, ty, tw, th, tconf, tcls
"""
nB = len(target) # target.shape[0]
nT = [len(x) for x in target] # torch.argmin(target[:, :, 4], 1) # targets per image
tx = torch.zeros(nB, nA, nG, nG) # batch size (4), number of anchors (3), number of grid points (13)
ty = torch.zeros(nB, nA, nG, nG)
tw = torch.zeros(nB, nA, nG, nG)
th = torch.zeros(nB, nA, nG, nG)
tconf = torch.ByteTensor(nB, nA, nG, nG).fill_(0)
tcls = torch.ByteTensor(nB, nA, nG, nG, nC).fill_(0) # nC = number of classes
TP = torch.ByteTensor(nB, max(nT)).fill_(0)
FP = torch.ByteTensor(nB, max(nT)).fill_(0)
FN = torch.ByteTensor(nB, max(nT)).fill_(0)
TC = torch.ShortTensor(nB, max(nT)).fill_(-1) # target category
for b in range(nB):
nTb = nT[b] # number of targets
if nTb == 0:
continue
t = target[b]
FN[b, :nTb] = 1
# Convert to position relative to box
TC[b, :nTb], gx, gy, gw, gh = t[:, 0].long(), t[:, 1] * nG, t[:, 2] * nG, t[:, 3] * nG, t[:, 4] * nG
# Get grid box indices and prevent overflows (i.e. 13.01 on 13 anchors)
gi = torch.clamp(gx.long(), min=0, max=nG - 1)
gj = torch.clamp(gy.long(), min=0, max=nG - 1)
# iou of targets-anchors (using wh only)
box1 = t[:, 3:5] * nG
# box2 = anchor_grid_wh[:, gj, gi]
box2 = anchor_wh.unsqueeze(1).repeat(1, nTb, 1)
inter_area = torch.min(box1, box2).prod(2)
iou_anch = inter_area / (gw * gh + box2.prod(2) - inter_area + 1e-16)
# Select best iou_pred and anchor
iou_anch_best, a = iou_anch.max(0) # best anchor [0-2] for each target
# Two targets can not claim the same anchor
if nTb > 1:
iou_order = np.argsort(-iou_anch_best) # best to worst
# u = torch.cat((gi, gj, a), 0).view(3, -1).numpy()
# _, first_unique = np.unique(u[:, iou_order], axis=1, return_index=True) # first unique indices
u = gi.float() * 0.4361538773074043 + gj.float() * 0.28012496588736746 + a.float() * 0.6627147212460307
_, first_unique = np.unique(u[iou_order], return_index=True) # first unique indices
# print(((np.sort(first_unique) - np.sort(first_unique2)) ** 2).sum())
i = iou_order[first_unique]
# best anchor must share significant commonality (iou) with target
i = i[iou_anch_best[i] > 0.10]
if len(i) == 0:
continue
a, gj, gi, t = a[i], gj[i], gi[i], t[i]
if len(t.shape) == 1:
t = t.view(1, 5)
else:
if iou_anch_best < 0.10:
continue
i = 0
tc, gx, gy, gw, gh = t[:, 0].long(), t[:, 1] * nG, t[:, 2] * nG, t[:, 3] * nG, t[:, 4] * nG
# Coordinates
tx[b, a, gj, gi] = gx - gi.float()
ty[b, a, gj, gi] = gy - gj.float()
# Width and height (sqrt method)
# tw[b, a, gj, gi] = torch.sqrt(gw / anchor_wh[a, 0]) / 2
# th[b, a, gj, gi] = torch.sqrt(gh / anchor_wh[a, 1]) / 2
# Width and height (yolov3 method)
tw[b, a, gj, gi] = torch.log(gw / anchor_wh[a, 0] + 1e-16)
th[b, a, gj, gi] = torch.log(gh / anchor_wh[a, 1] + 1e-16)
# One-hot encoding of label
tcls[b, a, gj, gi, tc] = 1
tconf[b, a, gj, gi] = 1
if requestPrecision:
# predicted classes and confidence
tb = torch.cat((gx - gw / 2, gy - gh / 2, gx + gw / 2, gy + gh / 2)).view(4, -1).t() # target boxes
pcls = torch.argmax(pred_cls[b, a, gj, gi], 1).cpu()
pconf = torch.sigmoid(pred_conf[b, a, gj, gi]).cpu()
iou_pred = bbox_iou(tb, pred_boxes[b, a, gj, gi].cpu())
TP[b, i] = (pconf > 0.99) & (iou_pred > 0.5) & (pcls == tc)
FP[b, i] = (pconf > 0.99) & (TP[b, i] == 0) # coordinates or class are wrong
FN[b, i] = pconf <= 0.99 # confidence score is too low (set to zero)
return tx, ty, tw, th, tconf, tcls, TP, FP, FN, TC
def non_max_suppression(prediction, conf_thres=0.5, nms_thres=0.4):
prediction = prediction.cpu()
"""
Removes detections with lower object confidence score than 'conf_thres' and performs
Non-Maximum Suppression to further filter detections.
Returns detections with shape:
(x1, y1, x2, y2, object_conf, class_score, class_pred)
"""
output = [None for _ in range(len(prediction))]
for image_i, pred in enumerate(prediction):
# Filter out confidence scores below threshold
# Get score and class with highest confidence
# cross-class NMS
cross_class_nms = False
if cross_class_nms:
thresh = 0.85
a = pred.clone()
a = a[np.argsort(-a[:, 4])] # sort best to worst
radius = 30 # area to search for cross-class ious
for i in range(len(a)):
if i >= len(a) - 1:
break
close = (np.abs(a[i, 0] - a[i + 1:, 0]) < radius) & (np.abs(a[i, 1] - a[i + 1:, 1]) < radius)
close = close.nonzero()
if len(close) > 0:
close = close + i + 1
iou = bbox_iou(a[i:i + 1, :4], a[close.squeeze(), :4].reshape(-1, 4), x1y1x2y2=False)
bad = close[iou > thresh]
if len(bad) > 0:
mask = torch.ones(len(a)).type(torch.ByteTensor)
mask[bad] = 0
a = a[mask]
pred = a
x, y, w, h = pred[:, 0].numpy(), pred[:, 1].numpy(), pred[:, 2].numpy(), pred[:, 3].numpy()
a = w * h # area
ar = w / (h + 1e-16) # aspect ratio
log_w, log_h, log_a, log_ar = np.log(w), np.log(h), np.log(a), np.log(ar)
# n = len(w)
# shape_likelihood = np.zeros((n, 60), dtype=np.float32)
# x = np.concatenate((log_w.reshape(-1, 1), log_h.reshape(-1, 1)), 1)
# from scipy.stats import multivariate_normal
# for c in range(60):
# shape_likelihood[:, c] = multivariate_normal.pdf(x, mean=mat['class_mu'][c, :2], cov=mat['class_cov'][c, :2, :2])
class_prob, class_pred = torch.max(F.softmax(pred[:, 5:], 1), 1)
v = ((pred[:, 4] > conf_thres) & (class_prob > .3)).numpy()
v = v.nonzero()
pred = pred[v]
class_prob = class_prob[v]
class_pred = class_pred[v]
# If none are remaining => process next image
nP = pred.shape[0]
if not nP:
continue
# From (center x, center y, width, height) to (x1, y1, x2, y2)
box_corner = pred.new(nP, 4)
xy = pred[:, 0:2]
wh = pred[:, 2:4] / 2
box_corner[:, 0:2] = xy - wh
box_corner[:, 2:4] = xy + wh
pred[:, :4] = box_corner
# Detections ordered as (x1, y1, x2, y2, obj_conf, class_prob, class_pred)
detections = torch.cat((pred[:, :5], class_prob.float().unsqueeze(1), class_pred.float().unsqueeze(1)), 1)
# Iterate through all predicted classes
unique_labels = detections[:, -1].cpu().unique()
if prediction.is_cuda:
unique_labels = unique_labels.cuda()
nms_style = 'OR' # 'AND' or 'OR' (classical)
for c in unique_labels:
# Get the detections with the particular class
detections_class = detections[detections[:, -1] == c]
# Sort the detections by maximum objectness confidence
_, conf_sort_index = torch.sort(detections_class[:, 4], descending=True)
detections_class = detections_class[conf_sort_index]
# Perform non-maximum suppression
max_detections = []
if nms_style == 'OR': # Classical NMS
while detections_class.shape[0]:
# Get detection with highest confidence and save as max detection
max_detections.append(detections_class[0].unsqueeze(0))
# Stop if we're at the last detection
if len(detections_class) == 1:
break
# Get the IOUs for all boxes with lower confidence
ious = bbox_iou(max_detections[-1], detections_class[1:])
# Remove detections with IoU >= NMS threshold
detections_class = detections_class[1:][ious < nms_thres]
elif nms_style == 'AND': # 'AND'-style NMS, at least two boxes must share commonality to pass, single boxes erased
while detections_class.shape[0]:
if len(detections_class) == 1:
break
ious = bbox_iou(detections_class[:1], detections_class[1:])
if ious.max() > 0.5:
max_detections.append(detections_class[0].unsqueeze(0))
# Remove detections with IoU >= NMS threshold
detections_class = detections_class[1:][ious < nms_thres]
if len(max_detections) > 0:
max_detections = torch.cat(max_detections).data
# Add max detections to outputs
output[image_i] = max_detections if output[image_i] is None else torch.cat(
(output[image_i], max_detections))
return output
def strip_optimizer_from_checkpoint(filename='checkpoints/best.pt'):
# Strip optimizer from *.pt files for lighter files (reduced by 2/3 size)
import torch
a = torch.load(filename, map_location='cpu')
a['optimizer'] = []
torch.save(a, filename.replace('.pt', '_lite.pt'))
def plotResults():
# Plot YOLO training results file "results.txt"
import numpy as np
import matplotlib.pyplot as plt
plt.figure(figsize=(18, 9))
s = ['x', 'y', 'w', 'h', 'conf', 'cls', 'loss', 'prec', 'recall']
for f in ('results.txt',):
results = np.loadtxt(f, usecols=[2, 3, 4, 5, 6, 7, 8, 9, 10]).T
for i in range(9):
plt.subplot(2, 5, i + 1)
plt.plot(results[i, :3000], marker='.', label=f)
plt.title(s[i])
plt.legend()