def main():
dataset = Balls_CF_Detection("./mini_balls/train", 20999)
# Get a single image from the dataset and display it
img, p, pose = dataset.__getitem__(2)
print(img.shape)
print(pose.shape)
show_bboxes(img, pose, COLORS, out_fn='_x.png')
def main():
torch.seed()
dataset_train = Balls_CF_Detection("data\\train\\train\\")
dataset_test = Balls_CF_Detection("data\\train\\val\\")
train_loader = torch.utils.data.DataLoader(dataset_train, batch_size=batch_size, shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset_test, batch_size=test_batch_size)
model = Net().to("cuda")
optimizer = optim.Adam(model.parameters(), lr=0.0025)
scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=1)
for epoch in range(1, epochs):
train(model, train_loader, optimizer, epoch)
test(model, test_loader)
scheduler.step()
def load(name):
dataset = Balls_CF_Detection("valid", 2099)
colormodel = ColorClassifier()
colormodel.load_state_dict(torch.load("colorClassifier.pth")["state_dict"])
colormodel.eval()
colormodel.train(False)
model = BoxClassifier(colormodel)
model.load_state_dict(torch.load(name + ".pth")["state_dict"])
model.eval()
model.train(False)
return model, dataset
def performance(name):
validation = 0
criterion = torch.nn.MSELoss()
dataset = Balls_CF_Detection("valid", 2099)
colormodel = ColorClassifier()
colormodel.load_state_dict(torch.load("colorClassifier.pth")["state_dict"])
colormodel.eval()
colormodel.train(False)
model = BoxClassifier(colormodel)
model.load_state_dict(torch.load(name + ".pth")["state_dict"])
model.eval()
model.train(False)
for i in range(2100):
img, p, pose = dataset.__getitem__(i)
out = model(img.view(1, 3, 100, 100)).view(9, 4)
if out[8, :].sum() > 0:
if out[8, 1] == 0:
out[8, 1] = out[8, 3] - 11
loss = criterion(out, torch.from_numpy(pose))
validation += loss.item()
return validation / 2100
np_bbox: np.array of shape (9,4) and a bbox is of type [x1,y1,x2,y2]
list_colors: list of string of length 9
"""
assert np_bbox.shape[0] == len(list_colors)
r = rgb_array.numpy()
r = np.uint8(r)
r = np.transpose(r, (1, 2, 0))
img_rgb = Image.fromarray(r, 'RGB')
draw = ImageDraw.Draw(img_rgb)
N = np_bbox.shape[0]
for i in range(N):
color = COLORS[i]
x_1, y_1, x_2, y_2 = np_bbox[i]
draw.rectangle(((x_1, y_1), (x_2, y_2)), outline=color, fill=None)
img_rgb.show()
img_rgb.save(out_fn)
if __name__ == "__main__":
dataset = Balls_CF_Detection("data\\train\\train\\")
# Get a single image from the dataset and display it
for i in range(10):
img, p, pose = dataset.__getitem__(i)
print(p)
print(img.shape)
print(pose.shape)
show_bboxes(img, pose, COLORS, out_fn='_x.png')
transforms.Normalize((0.1307,), (0.3081,))])) # mean, std of dataset
valid_loader = torch.utils.data.DataLoader(valid_dataset,
batch_size=BATCHSIZE, shuffle=True)
train_dataset = MNISTDataset ("MNIST-png/training",
transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))])) # mean, std of dataset)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=BATCHSIZE, shuffle=True)
'''
BATCHSIZE = 250
train_dataset = Balls_CF_Detection(
"./mini_balls/train", 0,
16000) #, transforms.Normalize([128, 128, 128], [3, 3, 3]))
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=BATCHSIZE,
shuffle=True)
valid_dataset = Balls_CF_Detection(
"./mini_balls/train", 16000,
21000) #, transforms.Normalize([128, 128, 128], [3, 3, 3]))
valid_loader = torch.utils.data.DataLoader(valid_dataset,
batch_size=BATCHSIZE,
shuffle=True)
class LeNet(torch.nn.Module):
def __init__(self):
import torch
from classifier import ColorClassifier
from dataset_det import Balls_CF_Detection
from torch.utils.tensorboard import SummaryWriter
import numpy as np
# environment settings
device = torch.device("cuda" if torch.cuda.is_available() else 'cpu')
print(torch.cuda.get_device_name(torch.cuda.current_device()))
writer = SummaryWriter("runs/color_classifier.01")
# dataset
train_dataset = Balls_CF_Detection("train", 18899, train=1)
valid_dataset = Balls_CF_Detection("valid", 2099)
# dataloader
training = torch.utils.data.DataLoader(train_dataset,
batch_size=50,
shuffle=True)
validation = torch.utils.data.DataLoader(valid_dataset,
batch_size=50,
shuffle=True)
# model
model = ColorClassifier().to(device)
model.cuda()
criterion = torch.nn.BCELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
batch_size_1 = 378
batch_size_2 = 42
transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))])) # mean, std of dataset
valid_loader = torch.utils.data.DataLoader(valid_dataset,
batch_size=BATCHSIZE, shuffle=True)
train_dataset = MNISTDataset ("MNIST-png/training",
transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))])) # mean, std of dataset)
train_loader = torch.utils.data.DataLoader(train_dataset,
batch_size=BATCHSIZE, shuffle=True)
'''
BATCHSIZE=10
train_dataset = Balls_CF_Detection ("./mini_balls/train", 20999)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=BATCHSIZE, shuffle=True)
'''
valid_dataset = train_dataset
valid_loader = train_loader
'''
class LeNet(torch.nn.Module):
def __init__(self):
super(LeNet, self).__init__()
self.conv1 = torch.nn.Conv2d(3, 12, 5, 1) #(1, 20, 5, 1)
self.conv2 = torch.nn.Conv2d(12, 36, 5, 1) #(20, 50, 5, 1)
self.fc1 = torch.nn.Linear(174240//10, 400) #(4*4*50, 500)
self.fc2 = torch.nn.Linear(400, 9) #(500, 10)
def forward(self, x):
np_bbox: np.array of shape (9,4) and a bbox is of type [x1,y1,x2,y2]
list_colors: list of string of length 9
"""
assert np_bbox.shape[0] == len(list_colors)
r = rgb_array.numpy()
r = np.uint8(r)
r = np.transpose(r, (1, 2, 0))
img_rgb = Image.fromarray(r, 'RGB')
draw = ImageDraw.Draw(img_rgb)
N = np_bbox.shape[0]
for i in range(N):
color = COLORS[i]
x_1, y_1, x_2, y_2 = np_bbox[i]
draw.rectangle(((x_1, y_1), (x_2, y_2)), outline=color, fill=None)
img_rgb.show()
img_rgb.save(out_fn)
if __name__ == "__main__":
dataset = Balls_CF_Detection("train", 20999)
# Get a single image from the dataset and display it
img, p, pose = dataset.__getitem__(9)
print(img.shape)
print(pose.shape)
print(pose)
show_bboxes(img, pose, COLORS, out_fn='_x.png')
""" Show the bounding box on a RGB image
rgb_array: a np.array of shape (H,W,3) - it represents the rgb frame in uint8 type
np_bbox: np.array of shape (9,4) and a bbox is of type [x1,y1,x2,y2]
list_colors: list of string of length 9
"""
assert np_bbox.shape[0] == len(list_colors)
r=rgb_array.numpy()
r = np.uint8(r)
r = np.transpose(r, (1,2,0))
img_rgb = Image.fromarray(r, 'RGB')
draw = ImageDraw.Draw(img_rgb)
N = np_bbox.shape[0]
for i in range(N):
color = COLORS[i]
x_1, y_1, x_2, y_2 = np_bbox[i]
draw.rectangle(((x_1, y_1), (x_2, y_2)), outline=color, fill=None)
img_rgb.show()
img_rgb.save(out_fn)
if __name__ == "__main__":
dataset = Balls_CF_Detection("../mini_balls/train",20999)
# Get a single image from the dataset and display it
img,p,pose = dataset.__getitem__(2)
print (img.shape)
print (pose.shape)
show_bboxes(img, pose, COLORS, out_fn='_x.png')