Spaces:
Moustached
/
Runtime error

Vengery / dataset /base_dataset.py
liuhaotian's picture
Upload app
f9a674e
raw
history blame
7.59 kB
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]