Spaces:
Sleeping
Sleeping
File size: 7,594 Bytes
281df87 |
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 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 |
import torch
from PIL import Image, ImageDraw
import torchvision.transforms as transforms
import torchvision
from zipfile import ZipFile
import os
import multiprocessing
import math
import numpy as np
import random
from io import BytesIO
VALID_IMAGE_TYPES = ['.jpg', '.jpeg', '.tiff', '.bmp', '.png']
def check_filenames_in_zipdata(filenames, ziproot):
samples = []
for fst in ZipFile(ziproot).infolist():
fname = fst.filename
if fname.endswith('/') or fname.startswith('.') or fst.file_size == 0:
continue
if os.path.splitext(fname)[1].lower() in VALID_IMAGE_TYPES:
samples.append((fname))
filenames = set(filenames)
samples = set(samples)
assert filenames.issubset(samples), 'Something wrong with your zip data'
def draw_box(img, boxes):
colors = ["red", "olive", "blue", "green", "orange", "brown", "cyan", "purple"]
draw = ImageDraw.Draw(img)
for bid, box in enumerate(boxes):
draw.rectangle([box[0], box[1], box[2], box[3]], outline =colors[bid % len(colors)], width=4)
# draw.rectangle([box[0], box[1], box[2], box[3]], outline ="red", width=2) # x0 y0 x1 y1
return img
def to_valid(x0, y0, x1, y1, image_size, min_box_size):
valid = True
if x0>image_size or y0>image_size or x1<0 or y1<0:
valid = False # no way to make this box vide, it is completely cropped out
return valid, (None, None, None, None)
x0 = max(x0, 0)
y0 = max(y0, 0)
x1 = min(x1, image_size)
y1 = min(y1, image_size)
if (x1-x0)*(y1-y0) / (image_size*image_size) < min_box_size:
valid = False
return valid, (None, None, None, None)
return valid, (x0, y0, x1, y1)
def recalculate_box_and_verify_if_valid(x, y, w, h, trans_info, image_size, min_box_size):
"""
x,y,w,h: the original annotation corresponding to the raw image size.
trans_info: what resizing and cropping have been applied to the raw image
image_size: what is the final image size
"""
x0 = x * trans_info["performed_scale"] - trans_info['crop_x']
y0 = y * trans_info["performed_scale"] - trans_info['crop_y']
x1 = (x + w) * trans_info["performed_scale"] - trans_info['crop_x']
y1 = (y + h) * trans_info["performed_scale"] - trans_info['crop_y']
# at this point, box annotation has been recalculated based on scaling and cropping
# but some point may fall off the image_size region (e.g., negative value), thus we
# need to clamp them into 0-image_size. But if all points falling outsize of image
# region, then we will consider this is an invalid box.
valid, (x0, y0, x1, y1) = to_valid(x0, y0, x1, y1, image_size, min_box_size)
if valid:
# we also perform random flip.
# Here boxes are valid, and are based on image_size
if trans_info["performed_flip"]:
x0, x1 = image_size-x1, image_size-x0
return valid, (x0, y0, x1, y1)
class BaseDataset(torch.utils.data.Dataset):
def __init__(self, image_root, random_crop, random_flip, image_size):
super().__init__()
self.image_root = image_root
self.random_crop = random_crop
self.random_flip = random_flip
self.image_size = image_size
self.use_zip = False
if image_root[-4::] == 'zip':
self.use_zip = True
self.zip_dict = {}
if self.random_crop:
assert False, 'NOT IMPLEMENTED'
def fetch_zipfile(self, ziproot):
pid = multiprocessing.current_process().pid # get pid of this process.
if pid not in self.zip_dict:
self.zip_dict[pid] = ZipFile(ziproot)
zip_file = self.zip_dict[pid]
return zip_file
def fetch_image(self, filename):
if self.use_zip:
zip_file = self.fetch_zipfile(self.image_root)
image = Image.open( BytesIO(zip_file.read(filename)) ).convert('RGB')
return image
else:
image = Image.open( os.path.join(self.image_root,filename) ).convert('RGB')
return image
def vis_getitem_data(self, index=None, out=None, return_tensor=False, name="res.jpg", print_caption=True):
if out is None:
out = self[index]
img = torchvision.transforms.functional.to_pil_image( out["image"]*0.5+0.5 )
canvas = torchvision.transforms.functional.to_pil_image( torch.ones_like(out["image"]) )
W, H = img.size
if print_caption:
caption = out["caption"]
print(caption)
print(" ")
boxes = []
for box in out["boxes"]:
x0,y0,x1,y1 = box
boxes.append( [float(x0*W), float(y0*H), float(x1*W), float(y1*H)] )
img = draw_box(img, boxes)
if return_tensor:
return torchvision.transforms.functional.to_tensor(img)
else:
img.save(name)
def transform_image(self, pil_image):
if self.random_crop:
assert False
arr = random_crop_arr(pil_image, self.image_size)
else:
arr, info = center_crop_arr(pil_image, self.image_size)
info["performed_flip"] = False
if self.random_flip and random.random()<0.5:
arr = arr[:, ::-1]
info["performed_flip"] = True
arr = arr.astype(np.float32) / 127.5 - 1
arr = np.transpose(arr, [2,0,1])
return torch.tensor(arr), info
def center_crop_arr(pil_image, image_size):
# We are not on a new enough PIL to support the `reducing_gap`
# argument, which uses BOX downsampling at powers of two first.
# Thus, we do it by hand to improve downsample quality.
WW, HH = pil_image.size
while min(*pil_image.size) >= 2 * image_size:
pil_image = pil_image.resize(
tuple(x // 2 for x in pil_image.size), resample=Image.BOX
)
scale = image_size / min(*pil_image.size)
pil_image = pil_image.resize(
tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC
)
# at this point, the min of pil_image side is desired image_size
performed_scale = image_size / min(WW, HH)
arr = np.array(pil_image)
crop_y = (arr.shape[0] - image_size) // 2
crop_x = (arr.shape[1] - image_size) // 2
info = {"performed_scale":performed_scale, 'crop_y':crop_y, 'crop_x':crop_x, "WW":WW, 'HH':HH}
return arr[crop_y : crop_y + image_size, crop_x : crop_x + image_size], info
def random_crop_arr(pil_image, image_size, min_crop_frac=0.8, max_crop_frac=1.0):
min_smaller_dim_size = math.ceil(image_size / max_crop_frac)
max_smaller_dim_size = math.ceil(image_size / min_crop_frac)
smaller_dim_size = random.randrange(min_smaller_dim_size, max_smaller_dim_size + 1)
# We are not on a new enough PIL to support the `reducing_gap`
# argument, which uses BOX downsampling at powers of two first.
# Thus, we do it by hand to improve downsample quality.
while min(*pil_image.size) >= 2 * smaller_dim_size:
pil_image = pil_image.resize(
tuple(x // 2 for x in pil_image.size), resample=Image.BOX
)
scale = smaller_dim_size / min(*pil_image.size)
pil_image = pil_image.resize(
tuple(round(x * scale) for x in pil_image.size), resample=Image.BICUBIC
)
arr = np.array(pil_image)
crop_y = random.randrange(arr.shape[0] - image_size + 1)
crop_x = random.randrange(arr.shape[1] - image_size + 1)
return arr[crop_y : crop_y + image_size, crop_x : crop_x + image_size]
|