def 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)
img = transforms.ToTensor()(img)
return img, boxes, labels
def collate_fn(batch):
imgs = [x[0] for x in batch] # [PIL.Image]
boxes = [x[1] for x in batch]
labels = [x[2] for x in batch]
num_imgs = len(imgs)
max_w = max([im.size[0] for im in imgs])
max_h = max([im.size[1] for im in imgs])
inputs = torch.zeros(num_imgs, 3, max_h, max_w)
loc_targets = []
cls_targets = []
for i in range(num_imgs):
img = imgs[i]
boxes_i = boxes[i].clone()
labels_i = labels[i].clone()
img, boxes_i = random_paste(img, boxes_i, (max_w, max_h), True)
inputs[i] = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])(img)
loc_target, cls_target = box_coder.encode(boxes_i, labels_i,
(max_w, max_h))
loc_targets.append(loc_target)
cls_targets.append(cls_target)
return inputs, torch.stack(loc_targets), torch.stack(cls_targets)
def transform_train_target(img, boxes, labels, img_size):
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
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=(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)
# print ("labels: ")
# print (labels.size())
boxes, labels = box_coder.encode(boxes, labels)
# print (labels.size())
return img, boxes, labels
def 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=(opt.img_size, opt.img_size),
random_interpolation=True)
img, boxes = random_flip(img, boxes)
img = transforms.Compose([transforms.ToTensor(), caffe_normalize])(img)
boxes, labels = box_coder.encode(boxes, labels)
return img, boxes, labels
def transform_train(img, boxes, labels, switch):
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)
if switch == 1:
while((0 not in labels) or (len(boxes_) == 1 and labels_[0]==0)) :
img_, boxes_, labels_ = random_crop(img, boxes, labels)
img, boxes, labels = 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)
boxes_tmp = list(boxes.clone().data.cpu().numpy())
labels_tmp = list(labels.clone().data.cpu().numpy())
att_box = []
if switch == 1:
new_boxes = []
new_labels = []
for ii in range(len(labels_tmp)):
if labels_tmp[ii] != 0:
new_boxes.append(boxes_tmp[ii])
new_labels.append(labels_tmp[ii])
else:
att_box.append(boxes_tmp[ii])
boxes = torch.from_numpy(np.array(new_boxes))
labels = torch.from_numpy(np.array(new_labels))
boxes, labels = box_coder.encode(boxes, labels)
return img, boxes, labels, att_box #boxes_tmp, labels_tmp
