import numpy as np
import os
from skimage import exposure, filters
import chainer
from chainercv import utils
from chainercv import transforms
import warnings
root = './Data/'
root2 = './Data/'
class DetectionDataset(chainer.dataset.DatasetMixin):
"""Base class for defect defection dataset
"""
def __init__(self, data_dir='auto', split='', img_size=1024, resize=False):
if data_dir == 'auto':
data_dir = root
self.data_dir = data_dir
self.img_size = img_size
self.resize = resize
images_file = os.path.join(self.data_dir, '{}images.txt'.format(split))
self.images = [line.strip() for line in open(images_file)]
def obtain_image_name(self,i):
#print(i)
return self.images[i]
def __len__(self):
return len(self.images)
def get_example(self, i):
"""Returns the i-th example.
Args:
i (int): The index of the example.
Returns:
tuple of an image and its label.
The image is in CHW format and its color channel is ordered in
RGB.
a bounding box is appended to the returned value.
"""
img = utils.read_image(
os.path.join(self.data_dir, 'images', self.images[i]),
color=True)
# Add processing to the other two channels
with warnings.catch_warnings():
warnings.simplefilter("ignore")
img[1, :, :] = exposure.rescale_intensity(exposure.equalize_adapthist(
exposure.rescale_intensity(img[1, :, :])), out_range=(0, 255))
img[2, :, :] = exposure.rescale_intensity(filters.gaussian(
exposure.rescale_intensity(img[2, :, :])), out_range=(0, 255))
# bbs should be a matrix (m by 4). m is the number of bounding
# boxes in the image
# labels should be an integer array (m by 1). m is the same as the bbs
bbs_file = os.path.join(self.data_dir, 'bounding_boxes', self.images[i][0:-4]+'.txt')
bbs = np.stack([line.strip().split() for line in open(bbs_file)]).astype(np.float32)
label = np.stack([0]*bbs.shape[0]).astype(np.int32)
_, H, W = img.shape
if self.resize and (H != self.img_size or W != self.img_size):
img = transforms.resize(img, (self.img_size, self.img_size))
bbs = transforms.resize_bbox(bbs, (H, W), (self.img_size, self.img_size))
return img, bbs, label
class MultiDetectionDataset(chainer.dataset.DatasetMixin):
"""Base class for multi defect defection dataset
"""
def __init__(self, data_dir='auto', split='', img_size=1024, resize=False):
if data_dir == 'auto':
data_dir = root2
self.data_dir = data_dir
self.img_size = img_size
self.resize = resize
images_file = os.path.join(self.data_dir, '{}images.txt'.format(split))
self.images = [
line.strip() for line in open(images_file)]
def __len__(self):
return len(self.images)
def get_example(self, i):
"""Returns the i-th example.
Args:
i (int): The index of the example.
Returns:
tuple of an image and its label.
The image is in CHW format and its color channel is ordered in
RGB.
a bounding box is appended to the returned value.
"""
img = utils.read_image(
os.path.join(self.data_dir, 'images', self.images[i]),
color=True)
# Add processing to the other two channels
with warnings.catch_warnings():
warnings.simplefilter("ignore")
img[1, :, :] = exposure.rescale_intensity(exposure.equalize_adapthist(
exposure.rescale_intensity(img[1, :, :])), out_range=(0, 255))
img[2, :, :] = exposure.rescale_intensity(filters.gaussian(
exposure.rescale_intensity(img[2, :, :])), out_range=(0, 255))
# bbs should be a matrix (m by 4). m is the number of bounding
# boxes in the image
# labels should be an integer array (m by 1). m is the same as the bbs
bbs_file = os.path.join(self.data_dir, 'bounding_boxes', self.images[i][0:-4]+'.txt')
print(bbs_file)
#bbs_file = "./Data/bounding_boxes/a.txt"
label_bbs = np.loadtxt(bbs_file, dtype=np.float32)
print(label_bbs)
# only 1D for medical case
#label = label_bbs[:,0].astype(np.int32)
#bbs = label_bbs[:,1:5]
label = label_bbs[0].astype(np.int32)
bbs = label_bbs[1:5]
_, H, W = img.shape
if self.resize and (H != self.img_size or W != self.img_size):
img = transforms.resize(img, (self.img_size, self.img_size))
bbs = transforms.resize_bbox(bbs, (H, W), (self.img_size, self.img_size))
return img, bbs, label
Datasets
Models
