def __getitem__(self, idx):
if self.use_M == True:
image = self.images[idx]
else:
image, _ = self.datainfo.__getitem__(idx)
if self.aug == True:
self.rv = random.random()
else:
self.rv = 1
if self.rv < 0.9:
# augmenation of img and masks
angle = random.randrange(-15, 15)
# trans img with masks
self.data_trans = mytransforms.Compose([mytransforms.ToPILImage(),
mytransforms.Rotation(angle),
mytransforms.ColorJitter(brightness=random.random(),
contrast=random.random(),
saturation=random.random(),
hue=random.random() / 2),
mytransforms.ToTensor(),
])
self.mask_trans = mytransforms.Compose([mytransforms.ToPILImage(),
mytransforms.Rotation(angle),
mytransforms.ToTensor(),
])
if self.binary == True:
image = perturbator.comb_binary_rec(image, [self.H, self.W])
#image = comb_black_rec(image, [self.H, self.W])
image = self.data_trans(image)
mask = torch.empty(self.mask_num, image.shape[1], image.shape[2], dtype=torch.float)
if self.use_M == True:
for k in range(0, self.mask_num):
X = self.images[idx + (self.data_num * (1 + k))]
mask[k] = self.mask_trans(X)
else:
for k in range(0, self.mask_num):
X, _ = self.datainfo.__getitem__(idx + (self.data_num * (1 + k)))
mask[k] = self.mask_trans(X)
else:
mask = torch.empty(self.mask_num, image.shape[1], image.shape[2], dtype=torch.float)
if self.use_M == True:
for k in range(0, self.mask_num):
X = self.images[idx + (self.data_num * (1 + k))]
mask[k] = X
else:
for k in range(0, self.mask_num):
X, _ = self.datainfo.__getitem__(idx + (self.data_num * (1 + k)))
mask[k] = X
mask = torch.pow(mask, self.pow_n)
mask = mask / mask.max()
return [image, mask]
def get_prediction(model, img_path, cat_names, threshold):
"""
get_prediction
parameters:
- model - the model to be used
- img_path - path of the input image
- cat_names - selected name for each category
- threshold - the confidence interval for making predictions
method:
- Image is obtained from the image path
- the image is converted to image tensor using PyTorch's Transforms
- image is passed through the model to get the predictions
- masks, classes and bounding boxes are obtained from the model and soft masks are made binary(0 or 1) on masks
ie: eg. segment of cat is made 1 and rest of the image is made 0
"""
img = Image.open(img_path)
transform = T.Compose([T.ToTensor()])
img = transform(img)
if use_cuda:
img = img.cuda()
pred = model([img])
pred_score = list(pred[0]['scores'].detach().cpu().numpy())
pred_t = [pred_score.index(x) for x in pred_score if x > threshold][-1]
masks = None
if 'masks' in pred[0]:
masks = (pred[0]['masks'] > 0.5).squeeze().detach().cpu().numpy()
elif 'keypoints' in pred[0]:
masks = (pred[0]['keypoints']).squeeze().detach().cpu().numpy()
pred_class = [cat_names[i] for i in list(pred[0]['labels'].cpu().numpy())]
pred_id = [i for i in list(pred[0]['labels'].cpu().numpy())]
pred_boxes = [[(i[0], i[1]), (i[2], i[3])]
for i in list(pred[0]['boxes'].detach().cpu().numpy())]
masks = masks[:pred_t + 1]
pred_boxes = pred_boxes[:pred_t + 1]
pred_class = pred_class[:pred_t + 1]
pred_id = pred_id[:pred_t + 1]
return masks, pred_boxes, pred_class, pred_id
roll_gen = mytransforms.RandomHorizontalRollGenerator()
flip_gen = mytransforms.RandomHorizontalFlipGenerator()
panostretch_gen = mytransforms.RandomPanoStretchGenerator(max_stretch=2.0)
joint_transform = mytransforms.Compose([
panostretch_gen,
[
mytransforms.RandomPanoStretch(panostretch_gen),
mytransforms.RandomPanoStretch(panostretch_gen),
mytransforms.RandomPanoStretch(panostretch_gen), None
],
flip_gen,
[
mytransforms.RandomHorizontalFlip(flip_gen),
mytransforms.RandomHorizontalFlip(flip_gen),
mytransforms.RandomHorizontalFlip(flip_gen), None
],
[
transforms.ToTensor(),
transforms.ToTensor(),
transforms.ToTensor(), None
],
roll_gen,
[
mytransforms.RandomHorizontalRoll(roll_gen),
mytransforms.RandomHorizontalRoll(roll_gen),
mytransforms.RandomHorizontalRoll(roll_gen), None
],
[transforms.RandomErasing(p=0.5, value=0), None, None, None],
])
trainset = SUN360Dataset(file="traindata.json",
transform=None,
def _train(args):
"""
is_distributed = len(args.hosts) > 1 and args.dist_backend is not None
logger.debug("Distributed training - {}".format(is_distributed))
if is_distributed:
# Initialize the distributed environment.
world_size = len(args.hosts)
os.environ['WORLD_SIZE'] = str(world_size)
host_rank = args.hosts.index(args.current_host)
os.environ['RANK'] = str(host_rank)
dist.init_process_group(backend=args.dist_backend, rank=host_rank, world_size=world_size)
logger.info(
'Initialized the distributed environment: \'{}\' backend on {} nodes. '.format(
args.dist_backend,
dist.get_world_size()) + 'Current host rank is {}. Using cuda: {}. Number of gpus: {}'.format(
dist.get_rank(), torch.cuda.is_available(), args.num_gpus))
"""
device = 'cuda' if torch.cuda.is_available() else 'cpu'
#device = 'cpu'
logger.info("Device Type: {}".format(device))
img_size = [128, 256]
logger.info("Loading SUN360 dataset")
train_transform = transforms.Compose([
transforms.Resize((img_size[0], img_size[1])),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
train_target_transform = transforms.Compose(
[transforms.Resize((img_size[0], img_size[1])),
transforms.ToTensor()])
roll_gen = mytransforms.RandomHorizontalRollGenerator()
flip_gen = mytransforms.RandomHorizontalFlipGenerator()
#panostretch_gen = mytransforms.RandomPanoStretchGenerator(max_stretch = 2.0)
#panostretch_gen,
#[mytransforms.RandomPanoStretch(panostretch_gen), mytransforms.RandomPanoStretch(panostretch_gen), mytransforms.RandomPanoStretch(panostretch_gen), None]
train_joint_transform = mytransforms.Compose([
[
transforms.Resize((img_size[0], img_size[1])),
transforms.Resize((img_size[0], img_size[1])),
transforms.Resize((img_size[0], img_size[1])), None
],
flip_gen,
[
mytransforms.RandomHorizontalFlip(flip_gen, p=0.5),
mytransforms.RandomHorizontalFlip(flip_gen, p=0.5),
mytransforms.RandomHorizontalFlip(flip_gen, p=0.5), None
],
[
transforms.ToTensor(),
transforms.ToTensor(),
transforms.ToTensor(), None
],
[
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]), None, None, None
],
roll_gen,
[
mytransforms.RandomHorizontalRoll(roll_gen, p=0.5),
mytransforms.RandomHorizontalRoll(roll_gen, p=0.5),
mytransforms.RandomHorizontalRoll(roll_gen, p=0.5), None
],
[
transforms.RandomErasing(p=0.5,
scale=(0.01, 0.02),
ratio=(0.3, 3.3),
value=0), None, None, None
],
])
valid_transform = transforms.Compose([
transforms.Resize((img_size[0], img_size[1])),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
])
valid_target_transform = transforms.Compose(
[transforms.Resize((img_size[0], img_size[1])),
transforms.ToTensor()])
"""
#uncomment this block if train/val split is needed
indices = list(range(len(trainvalidset)))
split = int(np.floor(len(trainvalidset)*0.8))
train_idx = indices[:10]
valid_idx = indices[10:]
train = Subset(trainvalidset, train_idx)
valid = Subset(trainvalidset, valid_idx)
trainset = SplitDataset(train, transform = None, target_transform = None, joint_transform=train_joint_transform)
"""
trainset = SUN360Dataset(file="traindata.json",
transform=None,
target_transform=None,
joint_transform=train_joint_transform)
train_loader = DataLoader(trainset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.workers)
#supplement= SUN360Dataset('morethan4corners.json',transform=None,target_transform=None,joint_transform=train_joint_transform)
#suppl_loader = DataLoader(supplement, batch_size=1,
# shuffle=True, num_workers=2)
validset = SUN360Dataset(file="testdata.json",
transform=valid_transform,
target_transform=valid_target_transform,
joint_transform=None)
valid_loader = DataLoader(validset,
batch_size=args.batch_size,
shuffle=False,
num_workers=args.workers)
logger.info("Model loaded")
if args.conv_type == "Std":
model = StdConvsCFL(args.model_name,
conv_type=args.conv_type,
layerdict=None,
offsetdict=None)
elif args.conv_type == "Equi":
layerdict, offsetdict = torch.load('layertrain.pt'), torch.load(
'offsettrain.pt')
model = EquiConvsCFL(args.model_name,
conv_type=args.conv_type,
layerdict=layerdict,
offsetdict=offsetdict)
# if torch.cuda.device_count() > 1:
# logger.info("Gpu count: {}".format(torch.cuda.device_count()))
# model = nn.DataParallel(model)
model = model.to(device)
criterion = CELoss().to(device)
# optimizer = torch.optim.Adam(model.parameters(), lr=args.lr,weight_decay=0.0005)
optimizer = torch.optim.Adam(model.parameters(),
lr=args.lr,
weight_decay=0.0)
LR_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, 0.995)
writer = SummaryWriter(
log_dir="{}".format(args.logdir),
comment="visualising losses of training and validation")
# for epoch in progressbar(range(1, args.epochs+1),redirect_stdout=True):
for epoch in tqdm.tqdm(range(1, args.epochs + 1)):
epochtime1 = time.time()
# training phase
phase = 'train'
running_loss = 0.0
for i, data in enumerate(train_loader):
# get the inputs
inputs, EM, CM = data
"""
'''this code block is to add one example of a room with
more than 4 floor-ceiling corner pairs to each batch '''
RGBsup,EMsup,CMsup = next(itertools.cycle(suppl_loader))
inputs = torch.cat([inputs,RGBsup],dim=0)
EM = torch.cat([EM,EMsup],dim=0)
CM = torch.cat([CM,CMsup],dim=0)
"""
inputs, EM, CM = inputs.to(device), EM.to(device), CM.to(device)
# zero the parameter gradients
optimizer.zero_grad()
# forward + backward + optimize
model.train()
outputs = model(inputs)
if (epoch % 10 == 0 and i == 0):
convert_to_images(outputs, epoch, phase)
EMLoss, CMLoss = map_loss(outputs, EM, CM, criterion)
loss = EMLoss + CMLoss
loss.backward()
optimizer.step()
# print statistics
running_loss += loss.item()
"""
if i % 1 == 0: # print every 1 mini-batches
print('[%d, %5d] loss: %.3f' %
(epoch, i + 1, running_loss / args.batch_size))
running_loss = 0.0
"""
epoch_loss = running_loss / len(trainset)
print("epoch: {}".format(epoch), ", train_loss: %.3f" % (epoch_loss))
writer.add_scalar("training_loss", epoch_loss, epoch)
# validation phase
if (epoch % 1 == 0):
phase = 'val'
with torch.no_grad():
running_loss = 0.0
for i, data in enumerate(valid_loader):
# get the inputs
inputs, EM, CM = data
inputs, EM, CM = inputs.to(device), EM.to(device), CM.to(
device)
model.eval()
outputs = model(inputs)
if (epoch % 10 == 0 and i == 0):
convert_to_images(outputs, epoch, phase)
EMLoss, CMLoss = map_loss(outputs, EM, CM, criterion)
loss = EMLoss + CMLoss
# print statistics
running_loss += loss.item()
#map_predict(outputs,EM,CM)
epoch_loss = running_loss / len(validset)
print("epoch: {}".format(epoch),
", valid_loss: %.3f" % (epoch_loss))
writer.add_scalar("validation loss", epoch_loss, epoch)
if (epoch % 100 == 0):
_save_model(model, args.model_dir, epoch)
LR_scheduler.step()
epochtime2 = time.time()
epochdiff = epochtime2 - epochtime1
writer.close()
print("time for 1 complete epoch: ", epochdiff)
print('Finished Training')
from Libs.pytorch_utils import utils
from torchvision.models.detection.rpn import AnchorGenerator
from create_detection_dataset import porpoise_dataset
IMG_RESIZE = 800
BATCH_SIZE = 8
NUM_WORKERS = 0
DATA_PATH = "porpoise_detection_data"
TRAIN_SPLIT = 0.1
TRANSFORM_TRAIN = T.Compose([
T.ToTensor(),
T.Resize(IMG_RESIZE),
T.RandomVerticalFlip(0.5),
T.RandomHorizontalFlip(0.5),
T.RandomColor(0.4, 0.2, 0.3, 0.1),
T.AddRandomNoise(0.02, 0.5),
#T.ShowImg(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
TRANSFORM_VAL = T.Compose([
T.ToTensor(),
T.Resize(IMG_RESIZE),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
def main():
# train on the GPU or on the CPU
device = torch.device(
def vis():
coco = CocoFolder(
file_dir, 8,
mytransforms.Compose([
mytransforms.RandomResized(),
mytransforms.RandomRotate(40),
mytransforms.RandomCrop(368),
mytransforms.RandomHorizontalFlip(),
]))
image, heatmaps, vecmap, mask = coco[9]
print(image.shape, heatmaps.shape, vecmap.shape, mask.shape)
image = image.numpy().transpose(1, 2, 0)
image *= 255
image += 128
image /= 255
print(mask.shape)
# for line in mask[0]:
# print(line)
mask = mask.numpy().transpose(1, 2, 0)
mask = cv2.resize(mask, (368, 368))
mask = mask.reshape((368, 368))
# new_img = image * mask
plt.imshow(image)
plt.imshow(mask, alpha=0.5)
plt.show()
plt.close()
heatmaps = heatmaps.numpy().transpose(1, 2, 0)
heatmaps = cv2.resize(heatmaps, (368, 368))
for j in range(19):
print(j)
heatmap = heatmaps[:, :, j]
for line in heatmap:
print(line)
heatmap = heatmap.reshape((368, 368))
heatmap *= 255
heatmap /= 255
result = np.expand_dims(heatmap, 2).repeat(3, axis=2) * 0.5
plt.imshow(image)
plt.imshow(heatmap, alpha=0.5)
plt.show()
plt.imsave('img{0}.jpg'.format(j), result)
plt.close()
print(vecmap.shape)
vecs = vecmap.numpy()
vecs = vecs.transpose(1, 2, 0)
vecs = cv2.resize(vecs, (368, 368))
for j in range(0, 38):
vec = np.abs(vecs[:, :, j])
# vec += np.abs(vecs[:, :, j + 1])
vec[vec > 1] = 1
vec = vec.reshape((368, 368, 1))
# vec[vec > 0] = 1
vec *= 255
# vec = cv2.applyColorMap(vec, cv2.COLORMAP_JET)
vec = vec.reshape((368, 368))
vec /= 255
plt.imshow(image)
# result = vec * 0.4 + img * 0.5
plt.imshow(vec, alpha=0.5)
plt.show()
plt.close()
import json
from PIL import Image
import matplotlib.pyplot as plt
import cv2
import mytransforms
image_dir = 'data/train2017/'
mask_dir = 'data/mask/'
file_path = 'data/data_json/train/filelist.txt'
mask_path = 'data/data_json/train/masklist.txt'
json_path = 'data/data_json/train/train2017.json'
file_dir = [file_path, mask_path, json_path]
transfor = mytransforms.Compose([
mytransforms.RandomResized(),
mytransforms.RandomRotate(40),
mytransforms.RandomCrop(368),
mytransforms.RandomHorizontalFlip(),
])
def read_data_file(file_dir, root_dir):
lists = []
with open(file_dir, 'r') as fp:
line = fp.readline()
while line:
path = line.strip()
lists.append(root_dir + path.split('/')[-1])
line = fp.readline()
return lists