1900zyh
3 years ago
parent
4180346273
commit
000b2223c4
@ -0,0 +1,326 @@
|
|||||||
|
import os
|
||||||
|
import io
|
||||||
|
import scipy
|
||||||
|
import torch
|
||||||
|
import random
|
||||||
|
import time
|
||||||
|
import lmdb
|
||||||
|
import math
|
||||||
|
import argparse
|
||||||
|
import pickle
|
||||||
|
import numpy as np
|
||||||
|
from PIL import Image, ImageDraw
|
||||||
|
from skimage.feature import canny
|
||||||
|
from xml.etree import ElementTree as ET
|
||||||
|
|
||||||
|
import torchvision.transforms as transforms
|
||||||
|
import torchvision.transforms.functional as F
|
||||||
|
from torch.utils.data import DataLoader
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
def random_bbox(img_width, img_height):
|
||||||
|
"""Generate a random tlhw regular mask
|
||||||
|
"""
|
||||||
|
vertical_margin = horizontal_margin = 0
|
||||||
|
mask_height = img_height//2
|
||||||
|
mask_width = img_width//2
|
||||||
|
max_delta_height = img_height // 8
|
||||||
|
max_delta_width = img_width // 8
|
||||||
|
maxt = img_height - vertical_margin - mask_height
|
||||||
|
maxl = img_width - horizontal_margin - mask_width
|
||||||
|
mask = np.zeros((img_height, img_width), np.uint8)
|
||||||
|
|
||||||
|
t = np.random.randint(vertical_margin, maxt)
|
||||||
|
l = np.random.randint(horizontal_margin, maxl)
|
||||||
|
h = np.random.randint(max_delta_height//2+1)
|
||||||
|
w = np.random.randint(max_delta_width//2+1)
|
||||||
|
mask[t+h:t+mask_height-h,
|
||||||
|
l+w:l+mask_width-w] = 1
|
||||||
|
return mask
|
||||||
|
|
||||||
|
|
||||||
|
def center_bbox(img_width, img_height):
|
||||||
|
"""Generate a center square mask
|
||||||
|
"""
|
||||||
|
mask = np.zeros((img_height, img_width), np.uint8)
|
||||||
|
mask[:, img_height//4:img_height//4*3,
|
||||||
|
img_width//4:img_width//4*3] = 1
|
||||||
|
return mask
|
||||||
|
|
||||||
|
|
||||||
|
def random_stroke(img_width, img_height):
|
||||||
|
min_num_vertex = 4
|
||||||
|
max_num_vertex = 12
|
||||||
|
mean_angle = 2*math.pi / 5
|
||||||
|
angle_range = 2*math.pi / 15
|
||||||
|
min_width = 12
|
||||||
|
max_width = 40
|
||||||
|
average_radius = math.sqrt(img_height*img_height+img_width*img_width) / 8
|
||||||
|
mask = Image.new('L', (img_width, img_height), 0)
|
||||||
|
|
||||||
|
steps = 6
|
||||||
|
for _ in range(np.random.randint(1, steps)):
|
||||||
|
num_vertex = np.random.randint(min_num_vertex, max_num_vertex)
|
||||||
|
angle_min = mean_angle - np.random.uniform(0, angle_range)
|
||||||
|
angle_max = mean_angle + np.random.uniform(0, angle_range)
|
||||||
|
angles = []
|
||||||
|
vertex = []
|
||||||
|
for i in range(num_vertex):
|
||||||
|
if i % 2 == 0:
|
||||||
|
angles.append(
|
||||||
|
2*math.pi - np.random.uniform(angle_min, angle_max))
|
||||||
|
else:
|
||||||
|
angles.append(np.random.uniform(angle_min, angle_max))
|
||||||
|
|
||||||
|
h, w = mask.size
|
||||||
|
vertex.append((int(np.random.randint(0, w)),
|
||||||
|
int(np.random.randint(0, h))))
|
||||||
|
for i in range(num_vertex):
|
||||||
|
r = np.clip(
|
||||||
|
np.random.normal(loc=average_radius,
|
||||||
|
scale=average_radius//2),
|
||||||
|
0, 2*average_radius)
|
||||||
|
new_x = np.clip(vertex[-1][0] + r * math.cos(angles[i]), 0, w)
|
||||||
|
new_y = np.clip(vertex[-1][1] + r * math.sin(angles[i]), 0, h)
|
||||||
|
vertex.append((int(new_x), int(new_y)))
|
||||||
|
|
||||||
|
draw = ImageDraw.Draw(mask)
|
||||||
|
width = int(np.random.uniform(min_width, max_width))
|
||||||
|
draw.line(vertex, fill=1, width=width)
|
||||||
|
for v in vertex:
|
||||||
|
draw.ellipse((v[0] - width//2,
|
||||||
|
v[1] - width//2,
|
||||||
|
v[0] + width//2,
|
||||||
|
v[1] + width//2),
|
||||||
|
fill=1)
|
||||||
|
|
||||||
|
if np.random.normal() > 0:
|
||||||
|
mask.transpose(Image.FLIP_LEFT_RIGHT)
|
||||||
|
if np.random.normal() > 0:
|
||||||
|
mask.transpose(Image.FLIP_TOP_BOTTOM)
|
||||||
|
mask = np.asarray(mask, np.uint8)
|
||||||
|
return mask
|
||||||
|
|
||||||
|
|
||||||
|
def bbox2np(img_width, img_height, bbox, pad_rate=0):
|
||||||
|
mask = np.zeros((img_height, img_width), np.uint8)
|
||||||
|
for (xmin, ymin, xmax, ymax) in bbox:
|
||||||
|
pad = int(pad_rate*min(xmax-xmin, ymax-ymin))
|
||||||
|
mask[max(0,ymin-pad):min(ymax+pad, img_height),
|
||||||
|
max(0,xmin-pad):min(xmax+pad, img_width)] = 1
|
||||||
|
return Image.fromarray(mask*255)
|
||||||
|
|
||||||
|
|
||||||
|
def outside_xml(oriw, orih, bbox):
|
||||||
|
if bbox is None:
|
||||||
|
return random_bbox(oriw, orih)
|
||||||
|
|
||||||
|
mask = np.zeros((orih, oriw))
|
||||||
|
logo_mask = np.array(bbox2np(oriw, orih, bbox))//255
|
||||||
|
random_mask = small_block(oriw, orih)
|
||||||
|
|
||||||
|
if np.sum(logo_mask) < 0.6*oriw*orih:
|
||||||
|
mask = ((mask + random_mask)>0) * (1-logo_mask)
|
||||||
|
iters = 5
|
||||||
|
while np.sum(mask) < 0.1*oriw*orih and iters>0:
|
||||||
|
random_mask = small_block(oriw, orih)
|
||||||
|
mask = ((mask+random_mask)>0) * (1-logo_mask)
|
||||||
|
iters -= 1
|
||||||
|
else:
|
||||||
|
mask = small_block(oriw, orih)
|
||||||
|
return np.array(mask).astype(np.uint8)
|
||||||
|
|
||||||
|
|
||||||
|
def small_block(width, height, nums=3):
|
||||||
|
mask = np.zeros((height, width)).astype(np.uint8)
|
||||||
|
margin_width = width//8
|
||||||
|
margin_height = height//8
|
||||||
|
for i in range(nums):
|
||||||
|
x = random.randint(margin_width, width - margin_width)
|
||||||
|
y = random.randint(margin_height, height - margin_height)
|
||||||
|
w = random.randint(margin_width, width//2)
|
||||||
|
h = random.randint(margin_height, height//2)
|
||||||
|
mask[y:min(y+h, height), x:min(x+w, width)] = 1
|
||||||
|
return mask
|
||||||
|
|
||||||
|
|
||||||
|
|
||||||
|
class LmdbReader(object):
|
||||||
|
lmdb_env = None
|
||||||
|
|
||||||
|
def __init__(self):
|
||||||
|
super(LmdbReader, self).__init__()
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def build_lmdb_env(lmdb_path):
|
||||||
|
if LmdbReader.lmdb_env is None:
|
||||||
|
LmdbReader.lmdb_env = lmdb.open(lmdb_path, max_readers=64, readonly=True,
|
||||||
|
lock=False, readahead=False, meminit=False,)
|
||||||
|
return LmdbReader.lmdb_env
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def read(path, key, val_type='img'):
|
||||||
|
env = LmdbReader.build_lmdb_env(path)
|
||||||
|
with env.begin(write=False) as txn:
|
||||||
|
try:
|
||||||
|
if val_type == 'int':
|
||||||
|
val = txn.get(key).decode('utf-8')
|
||||||
|
val = int(val)
|
||||||
|
elif val_type == 'img':
|
||||||
|
val = io.BytesIO(txn.get(key))
|
||||||
|
val = Image.open(val)
|
||||||
|
elif val_type == 'list':
|
||||||
|
val = io.BytesIO(txn.get(key))
|
||||||
|
val = pickle.load(val)
|
||||||
|
val = list(val)
|
||||||
|
except:
|
||||||
|
val = None
|
||||||
|
return val
|
||||||
|
|
||||||
|
|
||||||
|
class Dataset(torch.utils.data.Dataset):
|
||||||
|
def __init__(self, path, size, mask_type='bbox', split='train'):
|
||||||
|
super(Dataset, self).__init__()
|
||||||
|
self.lmdb_path = path
|
||||||
|
self.mask_type = mask_type
|
||||||
|
self.split = split
|
||||||
|
self.size = size
|
||||||
|
self.train_total = LmdbReader.read(path, 'train-total'.encode('utf-8'), val_type='int')
|
||||||
|
self.test_total = LmdbReader.read(path, 'test-total'.encode('utf-8'), val_type='int')
|
||||||
|
|
||||||
|
self._train_transform = transforms.Compose([
|
||||||
|
transforms.ToTensor(),
|
||||||
|
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ]
|
||||||
|
)
|
||||||
|
|
||||||
|
self._test_transform = transforms.Compose([
|
||||||
|
transforms.ToTensor(),
|
||||||
|
transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)), ]
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
def set_split(self, split='train'):
|
||||||
|
if split == 'train':
|
||||||
|
self.split = split
|
||||||
|
else:
|
||||||
|
self.split = 'test'
|
||||||
|
|
||||||
|
|
||||||
|
def __len__(self):
|
||||||
|
if self.split == 'train':
|
||||||
|
return self.train_total
|
||||||
|
else:
|
||||||
|
return self.test_total
|
||||||
|
|
||||||
|
|
||||||
|
def set_subset(self, start, end):
|
||||||
|
self.data = self.data[start:end]
|
||||||
|
|
||||||
|
|
||||||
|
def __getitem__(self, index):
|
||||||
|
key = f'{self.split}-{str(index).zfill(7)}-image'
|
||||||
|
orig_img = LmdbReader.read(self.lmdb_path, key.encode('utf-8'), val_type='img')
|
||||||
|
if self.split == 'train':
|
||||||
|
# obtain mask
|
||||||
|
if self.mask_type != 'xml':
|
||||||
|
img = transforms.Compose([transforms.RandomHorizontalFlip(),
|
||||||
|
transforms.Resize(self.size),
|
||||||
|
transforms.RandomCrop(self.size),])(orig_img)
|
||||||
|
if self.mask_type == 'bbox':
|
||||||
|
mask = random_bbox(self.size, self.size)
|
||||||
|
elif self.mask_type == 'stroke':
|
||||||
|
mask = random_stroke(self.size, self.size)
|
||||||
|
elif self.mask_type == 'comp':
|
||||||
|
mask = np.array(np.logical_or(random_bbox(self.size, self.size),
|
||||||
|
random_stroke(self.size, self.size))).astype(np.uint8)
|
||||||
|
mask = Image.fromarray(mask*255)
|
||||||
|
else: # self.mask_type == 'xml':
|
||||||
|
oriw, orih = orig_img.size
|
||||||
|
label_name = f'{self.split}-{str(index).zfill(7)}-label'
|
||||||
|
bbox = LmdbReader.read(self.lmdb_path, label_name.encode('utf-8'), val_type='list')
|
||||||
|
mask = Image.fromarray(outside_xml(oriw, orih, bbox)*255)
|
||||||
|
|
||||||
|
# resize
|
||||||
|
rate = 512.0 / min(orih, oriw)
|
||||||
|
neww, newh = int(oriw*rate), int(orih*rate)
|
||||||
|
img = orig_img.resize((neww, newh), Image.BILINEAR)
|
||||||
|
mask = mask.resize((neww, newh), Image.NEAREST)
|
||||||
|
|
||||||
|
x = random.randint(0, np.maximum(0, img.size[0] - self.size))
|
||||||
|
y = random.randint(0, np.maximum(0, img.size[1] - self.size))
|
||||||
|
img = img.crop((x, y, x+self.size, y+self.size))
|
||||||
|
mask = mask.crop((x, y, x+self.size, y+self.size))
|
||||||
|
|
||||||
|
if np.random.normal() > 0:
|
||||||
|
mask = mask.transpose(Image.FLIP_LEFT_RIGHT)
|
||||||
|
img = img.transpose(Image.FLIP_LEFT_RIGHT)
|
||||||
|
|
||||||
|
if np.sum(np.array(mask)) == 0:
|
||||||
|
mask = Image.fromarray(small_block(self.size, self.size)*255)
|
||||||
|
|
||||||
|
img = self._train_transform(img)
|
||||||
|
return key, img, transforms.ToTensor()(mask)
|
||||||
|
else:
|
||||||
|
oriw, orih = orig_img.size
|
||||||
|
rate = (self.size + 0.0) / max(orih, oriw)
|
||||||
|
neww, newh = int(oriw*rate), int(orih*rate)
|
||||||
|
img = orig_img.resize((neww, newh), Image.BILINEAR)
|
||||||
|
|
||||||
|
label_name = f'{self.split}-{str(index).zfill(7)}-label'
|
||||||
|
bbox = LmdbReader.read(self.lmdb_path, label_name.encode('utf-8'), val_type='list')
|
||||||
|
if bbox is not None:
|
||||||
|
orig_mask = bbox2np(oriw, orih, bbox)
|
||||||
|
mask = orig_mask.resize((neww, newh), Image.NEAREST)
|
||||||
|
else:
|
||||||
|
mask = random_bbox(neww, newh)
|
||||||
|
mask = Image.fromarray(mask*255)
|
||||||
|
|
||||||
|
img = F.pad(img, (0, 0, max(512 - neww, 0), max(512 - newh, 0)), fill=0, padding_mode='reflect')
|
||||||
|
mask = F.pad(mask, (0, 0, max(512-neww, 0), max(512 - newh, 0)), fill=0, padding_mode='reflect')
|
||||||
|
|
||||||
|
orig_img = self._test_transform(orig_img)
|
||||||
|
img = self._test_transform(img)
|
||||||
|
mask = transforms.ToTensor()(mask)
|
||||||
|
return key, orig_img, img, mask
|
||||||
|
|
||||||
|
def create_iterator(self, batch_size):
|
||||||
|
while True:
|
||||||
|
sample_loader = DataLoader(
|
||||||
|
dataset=self, batch_size=batch_size, drop_last=True)
|
||||||
|
for item in sample_loader:
|
||||||
|
yield item
|
||||||
|
|
||||||
|
|
||||||
|
if __name__ == '__main__':
|
||||||
|
|
||||||
|
parser = argparse.ArgumentParser(description="ca")
|
||||||
|
parser.add_argument('--lmdb_path', type=str, default='/data07/t-yazen/lsun_data/logos')
|
||||||
|
parser.add_argument('--size', type=int, default=512)
|
||||||
|
args = parser.parse_args()
|
||||||
|
|
||||||
|
|
||||||
|
d = Dataset(args.lmdb_path, args.size)
|
||||||
|
d.set_split('test')
|
||||||
|
print(len(d), ' for testing')
|
||||||
|
for i in range(5):
|
||||||
|
key, orig_img, img, mask = d[i]
|
||||||
|
print(orig_img.size(), np.unique(mask.numpy()))
|
||||||
|
orig_img = (orig_img.permute(1,2,0).numpy()+1)/2*255
|
||||||
|
orig_img = Image.fromarray(orig_img.astype(np.uint8))
|
||||||
|
|
||||||
|
oriw, orih = orig_img.size
|
||||||
|
rate = (args.size+0.0) / max(orih, oriw)
|
||||||
|
neww, newh = int(oriw*rate), int(orih*rate)
|
||||||
|
|
||||||
|
img = img*(1.-mask) + mask
|
||||||
|
img = (img.permute(1,2,0).numpy()+1)/2*255
|
||||||
|
img = Image.fromarray(img.astype(np.uint8))
|
||||||
|
mask = mask.squeeze().numpy()
|
||||||
|
mask = Image.fromarray(mask.astype(np.uint8))
|
||||||
|
img = F.crop(img, 0, 0, newh, neww)
|
||||||
|
mask = F.crop(mask, 0, 0, newh, neww)
|
||||||
|
print(oriw, orih, neww, newh)
|
||||||
|
mask = np.expand_dims(mask, axis=-1)
|
||||||
|
img = np.array(np.array(img)*(1.-mask) + mask*255).astype(np.uint8)
|
||||||
|
Image.fromarray(img).save(f'comp-{key}.png')
|
||||||
Loading…
Reference in new issue