def __getitem__(self, idx):
"""
"""
fname = self.fnames[idx]
img = Image.open(os.path.join(self.root, fname))
if img.mode != 'RGB':
img = img.convert('RGB')
boxes = self.boxes[idx].clone()
# converting ratios into numbers
w, h = img.size
boxes = boxes * torch.Tensor([w, h, w, h])
labels = self.labels[idx].clone()
size = self.input_size
#Data Augmentation
if self.train:
img, boxes = random_flip(img, boxes)
img, boxes = random_crop(img, boxes)
img, boxes = resize(img, boxes, (size, size))
else:
img, boxes = resize(img, boxes, (size, size))
#img, boxes = center_crop(img, boxes, (size,size))
img = self.transform(img)
return img, boxes, labels
def train_transform(img, boxes, labels):
img, boxes = random_paste(img, boxes, max_ratio=4, fill=(123, 116, 103))
img, boxes, labels = random_crop(img, boxes, labels)
img, boxes = resize(img, boxes, size=600, random_interpolation=True)
img, boxes = random_flip(img, boxes)
print(img.size)
img = transforms.ToTensor()(img)
print(boxes)
return img, boxes, labels
def transform_image_w_bbox(img, boxes, labels, img_size=224):
assert torch.is_tensor(boxes), 'type(boxes) : {}'.format(type(boxes))
img = random_distort(img)
# This slows down learning too much...
# if random.random() < 0.5:
# img, boxes = random_paste(img, boxes, max_ratio=4, fill=(123, 116, 103))
img, boxes, labels = random_crop(img, boxes, labels)
img, boxes = resize(img,
boxes,
size=(img_size, img_size),
random_interpolation=True)
img, boxes = random_flip(img, boxes)
img = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])(img)
return img, boxes, labels
def transform_train(img, boxes, labels):
img = random_distort(img)
if random.random() < 0.5:
img, boxes = random_paste(img,
boxes,
max_ratio=4,
fill=(123, 116, 103))
img, boxes, labels = random_crop(img, boxes, labels)
img, boxes = resize(img,
boxes,
size=(self.img_size, self.img_size),
random_interpolation=False)
img, boxes = random_flip(img, boxes)
img = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406),
(0.229, 0.224, 0.225))
])(img)
boxes, labels = self.box_coder.encode(boxes, labels)
return img, boxes, labels
import dataloader as dataloader
import transforms as t
import numpy as np
import matplotlib.pyplot as plt
from hcat.unet import Unet_Constructor as GUnet
from hcat.loss import dice_loss
import torch
data = dataloader.stack(
path='/home/chris/Desktop/ColorImages',
joint_transforms=[
t.to_float(),
t.reshape(),
t.random_crop([512, 512, 30]),
],
image_transforms=[t.normalize([0.5, 0.5, 0.5, 0.5], [0.5, 0.5, 0.5, 0.5])])
if torch.cuda.is_available():
device = 'cuda:0'
else:
device = 'cpu'
print('Initalizing Unet: ', end='')
unet = GUnet(image_dimensions=3,
in_channels=4,
out_channels=1,
feature_sizes=[16, 32, 64, 128],
kernel={
'conv1': (3, 3, 2),
'conv2': (3, 3, 1)
if True:
print('test random_horizontal_flip')
for i in range(5):
ct = c.copy()
dt = d.copy()
ct, dt = random_horizontal_flip(ct, dt)
show_detection(ct, dt)
if True:
print('test random_crop')
for i in range(7):
ct = c.copy()
dt = d.copy()
ce = e.copy()
ct, dt, ce = random_crop(ct, dt, ce)
show_detection(ct, dt)
if True:
print('test photometric_distort')
for i in range(5):
ct = c.copy()
dt = d.copy()
ct = photometric_distort(ct)
show_detection(ct, dt)
if True:
print('test channel_swap')
for i in range(5):
ct = c.copy()
dt = d.copy()