Spaces:
Running
on
T4
Running
on
T4
File size: 5,246 Bytes
a277bb8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 |
import PIL #version 1.2.0
import torch
import os
import torchvision.transforms.functional as F
import numpy as np
import random
def intersect(boxes1, boxes2):
'''
Find intersection of every box combination between two sets of box
boxes1: bounding boxes 1, a tensor of dimensions (n1, 4)
boxes2: bounding boxes 2, a tensor of dimensions (n2, 4)
Out: Intersection each of boxes1 with respect to each of boxes2,
a tensor of dimensions (n1, n2)
'''
n1 = boxes1.size(0)
n2 = boxes2.size(0)
max_xy = torch.min(boxes1[:, 2:].unsqueeze(1).expand(n1, n2, 2),
boxes2[:, 2:].unsqueeze(0).expand(n1, n2, 2))
min_xy = torch.max(boxes1[:, :2].unsqueeze(1).expand(n1, n2, 2),
boxes2[:, :2].unsqueeze(0).expand(n1, n2, 2))
inter = torch.clamp(max_xy - min_xy , min=0) # (n1, n2, 2)
return inter[:, :, 0] * inter[:, :, 1] #(n1, n2)
def find_IoU(boxes1, boxes2):
'''
Find IoU between every boxes set of boxes
boxes1: a tensor of dimensions (n1, 4) (left, top, right , bottom)
boxes2: a tensor of dimensions (n2, 4)
Out: IoU each of boxes1 with respect to each of boxes2, a tensor of
dimensions (n1, n2)
Formula:
(box1 ∩ box2) / (box1 u box2) = (box1 ∩ box2) / (area(box1) + area(box2) - (box1 ∩ box2 ))
'''
inter = intersect(boxes1, boxes2)
area_boxes1 = (boxes1[:, 2] - boxes1[:, 0]) * (boxes1[:, 3] - boxes1[:, 1])
area_boxes2 = (boxes2[:, 2] - boxes2[:, 0]) * (boxes2[:, 3] - boxes2[:, 1])
area_boxes1 = area_boxes1.unsqueeze(1).expand_as(inter) #(n1, n2)
area_boxes2 = area_boxes2.unsqueeze(0).expand_as(inter) #(n1, n2)
union = (area_boxes1 + area_boxes2 - inter)
return inter / union
def random_crop(image, boxes, labels, difficulties=None):
'''
image: A PIL image
boxes: Bounding boxes, a tensor of dimensions (#objects, 4)
labels: labels of object, a tensor of dimensions (#objects)
difficulties: difficulties of detect object, a tensor of dimensions (#objects)
Out: cropped image , new boxes, new labels, new difficulties
'''
if type(image) == PIL.Image.Image:
image = F.to_tensor(image)
original_h = image.size(1)
original_w = image.size(2)
while True:
mode = random.choice([0.1, 0.3, 0.5, 0.9, None])
if mode is None:
return F.to_pil_image(image), boxes, labels, difficulties
new_image = image
new_boxes = boxes
new_difficulties = difficulties
new_labels = labels
for _ in range(50):
# Crop dimensions: [0.3, 1] of original dimensions
new_h = random.uniform(0.3*original_h, original_h)
new_w = random.uniform(0.3*original_w, original_w)
# Aspect ratio constraint b/t .5 & 2
if new_h/new_w < 0.5 or new_h/new_w > 2:
continue
#Crop coordinate
left = random.uniform(0, original_w - new_w)
right = left + new_w
top = random.uniform(0, original_h - new_h)
bottom = top + new_h
crop = torch.FloatTensor([int(left), int(top), int(right), int(bottom)])
# Calculate IoU between the crop and the bounding boxes
overlap = find_IoU(crop.unsqueeze(0), boxes) #(1, #objects)
overlap = overlap.squeeze(0)
# If not a single bounding box has a IoU of greater than the minimum, try again
if overlap.shape[0] == 0:
continue
if overlap.max().item() < mode:
continue
#Crop
new_image = image[:, int(top):int(bottom), int(left):int(right)] #(3, new_h, new_w)
#Center of bounding boxes
center_bb = (boxes[:, :2] + boxes[:, 2:])/2.0
#Find bounding box has been had center in crop
center_in_crop = (center_bb[:, 0] >left) * (center_bb[:, 0] < right
) *(center_bb[:, 1] > top) * (center_bb[:, 1] < bottom) #( #objects)
if not center_in_crop.any():
continue
#take matching bounding box
new_boxes = boxes[center_in_crop, :]
#take matching labels
new_labels = labels[center_in_crop]
#take matching difficulities
if difficulties is not None:
new_difficulties = difficulties[center_in_crop]
else:
new_difficulties = None
#Use the box left and top corner or the crop's
new_boxes[:, :2] = torch.max(new_boxes[:, :2], crop[:2])
#adjust to crop
new_boxes[:, :2] -= crop[:2]
new_boxes[:, 2:] = torch.min(new_boxes[:, 2:],crop[2:])
#adjust to crop
new_boxes[:, 2:] -= crop[:2]
return F.to_pil_image(new_image), new_boxes, new_labels, new_difficulties |