Computer Vision: Code Overview
Now that I've implemented a convolutional neural network in this tutorial. Now, I'm gonna start looking into the code.
What I'm going to do today is to go over all the code lines, and try and pickpocket and "figure out" what each of the lines of each of the classes do.
The first file that I'm going to go over is utils.py and document the pseudocode for this.
import numpy as np
import cv2
import matplotlib.pyplot as plt
def convert_exr_and_write(exr, dest_path, img_format, default_height = 400, default_width = 900):
im=cv2.imread(exr,-1)
height, width = im.shape[:2]
# convert to jpg
if img_format == 'jpg':
tonemap = cv2.createTonemap(gamma=1)
im = tonemap.process(im)
im = im[0:height-600,:]
im = cv2.normalize(im, None, alpha=0, beta=20000, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_32F)
im=np.uint16(im)
im = cv2.resize(im, (default_width, default_height))
cv2.imwrite(dest_path, im)
elif img_format == 'exr': # convert to exr
tonemap = cv2.createTonemap(gamma=1)
im = tonemap.process(im)
im = im[0:height-600,:]
im = cv2.resize(im, (default_width, default_height))
cv2.imwrite(dest_path, im)
else:
raise TypeError('Supported format jpg and exr')
def imshow(img):
cv2.imshow('image', img)
cv2.waitKey(0)
The first line loads an image.
The second line is named im.shape. This line returns the dimensions of this image.
createToneMap() creates a simple linear map with gamma correction. Whatever that means.
The process call merges the image and then the image is called from index 0 to the 600th last index.
After which the image is normalized where the minimum distance is 0 and the maximum distance is 20,000. I convert everything to integer values and subsequently resize the image and write the image to a certain destination path. This normalization is required for .img file, but not for the .exr file, as this is already compressed.
Now for the next file. It's called generate-feature-matrix.py. Here is the following code:
import os
import cv2
import json
import argparse
import numpy as np
from dataset.preprocess import DataGenerator
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Generate Feature Matrix')
parser.add_argument('-f', '--imgformat', type=str,
help='generate annotated images of a specific type: jpg or exr', default='jpg')
parser.add_argument('-sp', '--sourcepath', type=str, help='data source path', default=os.path.join('data', 'jpg_sample'))
parser.add_argument('-dp', '--destpath', type=str, help='data destination path', default='data')
# format: jpg or exr
args = parser.parse_args()
img_format = args.imgformat
img_format = str.lower(img_format)
if img_format == 'jpg' or img_format == 'jpeg' or img_format == 'exr':
pass
else:
raise TypeError("Supported format: jpg and exr")
# paths
source_path = args.sourcepath
dest_path = args.destpath
# generate feature matrix
matrix_generator = DataGenerator()
matrix_generator.generate_feature_and_label(os.path.join('data', 'jpg_sample'), 'data', img_format)
The first line is just to generate the description feature matrix.
3 arguments take strings, specifying the image format, the source path, and the destination path.
We want to make sure the image format is .exr and .jpg. The source path and destination path is subsequently generated and the feature and label is finally generated. How does this happen? Well, the labels of the image are generated
This is what happens when the feature and label are generated.
The next file I want to discuss is convert_exr.py. The file is as follows:
import os
import argparse
from utils import convert_exr_and_write
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Convert exr to exr or jpg both of size (400, 900)')
parser.add_argument('-f', '--imgformat', type=str,
help='The type images would be converted to', default='jpg')
parser.add_argument('-dp', '--datapath', type=str, help='data source path', default='/home/adrian/Downloads/laval-dataset')
parser.add_argument('-s', '--size', type=int,
help='numebr of images converted, if is -1, convert all the image', default=30)
args = parser.parse_args()
data_path = args.datapath
img_format = args.imgformat
size = args.size
# format check
if str.lower(img_format) != 'jpg' and str.lower(img_format) != 'jpeg' and str.lower(img_format) != 'exr':
raise TypeError("Format supported: jpg and exr")
if str.lower(img_format) == 'jpeg':
img_format= 'jpg'
img_format = str.lower(img_format)
# store images in different dirs based on format
dest_path = os.path.join('data', img_format)
if not os.path.isdir(dest_path):
os.mkdir(dest_path)
file_ind = 1
exr_lists = os.listdir(data_path)
# number of images is size
if size != -1 and size < len(exr_lists):
exr_lists = exr_lists[0:size]
# convert and write
for exr_file_name in exr_lists:
converted_file = os.path.join(dest_path, str(file_ind) + '.' + img_format)
convert_exr_and_write(os.path.join(data_path, exr_file_name), converted_file, img_format)
file_ind = file_ind + 1
This file converts .exr file to .jpg, since .jpg is much more efficient than .exr. If the format of the image isn't .jpg or .exr, then nothing will be able to be supported. We convert the format of the image into lowercase, and then we store images in different directories based on format into the said vectors.
The next file before the major Train and the neural network classes is the generated_annotated_examples.py. The following is the code for this class:
import os
import cv2
import argparse
import numpy as np
from dataset.preprocess import JPGLabeler, EXRLabeler
if __name__ == '__main__':
# parse arguments
parser = argparse.ArgumentParser(description='Generate example images')
parser.add_argument('-f', '--imgformat', type=str,
help='generate annotated images of a specific type: jpg or exr', default='exr')
parser.add_argument('-dp', '--datapath', type=str, help='data source path', default='exr_sample')
args = parser.parse_args()
img_format = args.imgformat
# decide a labeler
if str.lower(img_format) == 'jpg' or str.lower(img_format) == 'jpeg':
labler = JPGLabeler()
elif str.lower(img_format) == 'exr':
labler = EXRLabeler()
else:
raise TypeError("Format supported: jpg and exr")
data_path = os.path.join('data', args.datapath)
example_path = os.path.join('annotated_examples', str.lower(img_format))
if not os.path.isdir(example_path):
os.mkdir(example_path)
# get image files
img_names = os.listdir(os.path.join(data_path))
discarded_data_set = set()
img_names.sort()
# generate images
for img_name in img_names:
img = cv2.imread(os.path.join(data_path, img_name), cv2.IMREAD_UNCHANGED).astype(np.float32)
img = labler.generate_annotated_img(img)
# any img with fewer than 3 lights is removed.
if img is None:
discarded_data_set.add(img_name)
continue
# write images to the specific location
cv2.imwrite(os.path.join(example_path, img_name),img)
This class, I believe is just a training data class, generating annotated images of a particular dataset. First, we generate a labeler based on the format of the image. JPG_labeler takes the Gaussian form of a particular image and then blur, erodes and dilates an image in order to make it easier to process. This is stored and any image under 3 light sources is removed.
The next part of the code I want to go through is the program, called train.py. Here is the code:
import time
import os
import copy
import data
import torch
import torch.optim as optim
import numpy as np
import matplotlib.pyplot as plt
import argparse
import dataset.transformer as transformer
from datetime import datetime
from torch.utils.data import DataLoader
from torch.optim import lr_scheduler
from torchvision import transforms
from model.loss import average_difference_loss, location_success_count
from model.network import IlluminationPredictionNet
from dataset.dataset import EnvironmentJPGDataset
def train_model(model, criterion, optimizer, scheduler, num_epochs=25):
global train_dataloader
global test_dataloader
since = time.time()
best_model_wts = None
best_acc = 0.0
device = torch.device('cuda' if torch.cuda.is_available() else ('cpu'))
model.to(device)
# to draw figures
train_loss_epoch = np.zeros((num_epochs, 2))
train_acc_epoch = np.zeros((num_epochs, 2))
val_loss_epoch = np.zeros((num_epochs, 2))
val_acc_epoch = np.zeros((num_epochs, 2))
for epoch in range(num_epochs):
print('Epoch {}/{}'.format(epoch, num_epochs - 1))
print('-' * 10)
# Each epoch has a training and validation phase
for phase in ['train', 'validation']:
if phase == 'train':
model.train() # Set model to training mode
dataloader = train_dataloader
else:
model.eval() # Set model to evaluate mode
dataloader = test_dataloader
running_loss = 0.0
running_corrects = 0
# Iterate over data.
for (i, data) in enumerate(dataloader):
inputs = data[0].to(device)
labels = data[1].to(device)
# zero the parameter gradients
optimizer.zero_grad()
# forward
# track history if only in train
with torch.set_grad_enabled(phase == 'train'):
outputs = model(inputs)
outputs = torch.reshape(outputs, (-1, 3, 9))
loss = criterion(outputs, labels)
running_corrects += location_success_count(outputs, labels)
# backward + optimize only if in training phase
if phase == 'train':
loss.backward()
optimizer.step()
# statistics
running_loss += loss.item() * inputs.size(0)
if phase == 'train' and scheduler is not None:
scheduler.step()
epoch_loss = running_loss / len(dataloader)/dataloader.batch_size
epoch_acc = running_corrects / len(dataloader)/3/dataloader.batch_size
# record loss & acc during training
if phase == 'train':
train_loss_epoch[epoch][1] = epoch_loss
train_acc_epoch[epoch][1] = epoch_acc
train_loss_epoch[epoch][0] = epoch
train_acc_epoch[epoch][0] = epoch
else:
val_loss_epoch[epoch][1] = epoch_loss
val_acc_epoch[epoch][1] = epoch_acc
val_loss_epoch[epoch][0] = epoch
val_acc_epoch[epoch][0] = epoch
print('{} Loss: {:.4f} Acc: {:.4f}'.format(
phase, epoch_loss, epoch_acc))
if epoch_acc > best_acc and phase == 'validation':
best_model_wts = copy.deepcopy(model)
best_acc = epoch_acc
print()
time_elapsed = time.time() - since
print('Training complete in {:.0f}m {:.0f}s'.format(
time_elapsed // 60, time_elapsed % 60))
if best_model_wts is None and num_epochs >= 25:
raise TypeError("Accuracy Metric is Invalid")
if num_epochs >= 25:
torch.save(best_model_wts, os.path.join('checkpoint','naive_model_with_activation' + datetime.now().strftime("_%H_%M_%S_%d-%m-%Y")))
return best_model_wts, train_loss_epoch, train_acc_epoch, val_loss_epoch, val_acc_epoch
def plot_loss_acc(train_loss_epoch, train_acc_epoch, val_loss_epoch, val_acc_epoch):
epoch_num = len(train_loss_epoch)
# loss plot
l1, = plt.plot(train_loss_epoch[:,0], train_loss_epoch[:, 1], color='blue')
l2, = plt.plot(val_loss_epoch[:,0], val_loss_epoch[:, 1], color ='red')
plt.legend(handles=[l1,l2],labels=['train','validation'],loc='best')
plt.title('Training and Validation Loss')
plt.xlabel('epoch')
plt.ylabel('loss')
plt.savefig(os.path.join("figures", "epoch-vs-squared_loss" + datetime.now().strftime("_%H-%M-%S_%d-%m-%Y") + ".png"))
plt.close('all')
# acc plot
l1, = plt.plot(train_acc_epoch[:,0], train_acc_epoch[:,1], color='blue')
l2, = plt.plot(val_acc_epoch[:,0], val_acc_epoch[:,1], color ='red')
plt.legend(handles=[l1,l2],labels=['train','validation'],loc='best')
plt.title('Training and Validation Accuracy')
plt.xlabel('epoch')
plt.ylabel('accuracy')
plt.savefig(os.path.join("figures", "epoch-vs-accuracy"+ datetime.now().strftime("_%H-%M-%S_%d-%m-%Y") + ".png"))
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='train model')
parser.add_argument('-opt', '--optimizer', type=str,
help='choose SGD or Adam', default='sgd')
parser.add_argument('-lr', '--learningrate', type=float, help='learning rate', default=0.0001)
parser.add_argument('-mm', '--momentum', type=float, help='momentum for sgd', default=0.9)
parser.add_argument('-b1', '--beta1', type=float, help='beta1 parameter for Adam', default=0.9)
parser.add_argument('-b2', '--beta2', type=float, help='beta2 parameter for Adam', default=0.999)
parser.add_argument('-e', '--epsilon', type=float, help='eps parameter for Adam', default=1e-8)
parser.add_argument('-he', '--height', type=int, help = 'height of the input image', default=400)
parser.add_argument('-w', '--width', type=int, help = 'width of the input image', default=900)
parser.add_argument('-bs', '--batchsize', type=int, help='batch size', default=16)
parser.add_argument('-epoch', '--epoch', type=int, help='training epoch', default = 50)
args = parser.parse_args()
# args
choice_of_optimizer = str.lower(str(args.optimizer))
if choice_of_optimizer != 'sgd' and choice_of_optimizer != 'adam':
raise TypeError('Optimizer Must be SGD or Adam')
batch_size = args.batchsize
learning_rate = args.learningrate
sgd_momentum = args.momentum
beta1 = args.beta1
beta2 = args.beta2
adam_beta = (beta1, beta2)
adam_epsilon = args.epsilon
height = args.height
width = args.width
#transformer.CustomNormalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
# dataset
train_ds = EnvironmentJPGDataset(os.path.join('data', 'train_feature_matrix.npy'), os.path.join('data', 'train_label.npy'),\
transform= transforms.Compose([
transformer.ToTensor(),
transformer.CustomNormalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
, augmentation=(1, 2))
test_ds = EnvironmentJPGDataset(os.path.join('data', 'test_feature_matrix.npy'), os.path.join('data', 'test_label.npy'),\
transform= transforms.Compose([
transformer.ToTensor(),
transformer.CustomNormalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
, augmentation=(1, 2))
train_dataloader = DataLoader(train_ds, batch_size)
test_dataloader = DataLoader(test_ds, batch_size)
# model
model = IlluminationPredictionNet()
model.double()
# optimizer
if optim == 'sgd':
optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=sgd_momentum)
else:
optimizer = optim.Adam(model.parameters(), lr=learning_rate, betas=adam_beta, eps=adam_epsilon)
scheduler = lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.01)
# train
model, train_loss_epoch, train_acc_epoch, val_loss_epoch, val_acc_epoch = train_model(model, average_difference_loss, optimizer, scheduler, 1)
plot_loss_acc(train_loss_epoch, train_acc_epoch, val_loss_epoch, val_acc_epoch)
There is a line of code that keeps track of the time, and dedicates the code to the CPU or GPU depending on what part of the model the code is located in. There is a training set and a value set that is made, with num_epoch rows and 2 columns.
Each epoch has both a training and validation phase. If the epoch is in the training phase, then set the model into training mode. Else, set the model into evaluation/testing mode. The loss and corrects are set to 0 and then we start enumerating all of the data. First, we zero the parameter, and track the parameter history if only in training mode.
Torch.reshape reshapes the vectors to a 3x9 with 3 lights and 9 features for each light, including the RGB value. Get everything from the dataloader, and get the inputs and labels, and model everything with a convolutional neural network. We then compute the loss and implement the number of correct labels after training the network. Loss.backward() computes the optimizer loss and stores the gradient into tensors(vectors). Then, all the gradients are updated with the function optimizer.step(); We increment the total loss and then have an average epoch loss and epoch accuracy. There is a training loss epoch and value loss epoch. An unknown command is copy.deepcopy(). The best model and accuracy is copied and tricked, and a graph is ploted in case of epoch and accuracy, and saved based on the date.
In the main function, there's a lot of arguments, including, beta, and momentum, to include a few of them. We initialize a prediction network and run the model-update method. However, before this update is run the optimizer is initialized for either stochastic gradient descent or Adam optimization. The SGD is the gradient descent with momentum. Adam optimization is a stochastic gradient descent method that is based on adaptive estimation of first-order and second-order moments.
The next file I want to discuss here is the dataset.py. There are 2 methods in this class, EnvironmentJPGDataset, and EnvironmentEXRDataset. There are 2 files here, img_npy_file and label_npy_file. We shape the environment and the labels. After which comes the augmentation part. Augmentation is a tuple of 2 dimensions, first dimension is the number of rows, and the second dimension is the number of columns, as follows:
[num_rows, num_cols]. First we assume an environmental frame, and then a label frame. The for loops separates things into width/column number and height / row number, respectively. I'm not entirely sure what transform means, but I am not entirely sure, probably things to set up the api. Transform converts everything into a tensor vector which needs to be this way for pytorch to operate on it.
The next class is the transformer class and it is as follows:
import torch
import numpy as np
from skimage import transform
class Normalize(object):
'''Normalize image matrix'''
def __init__(self):
pass
def __call__(self, input):
image, labels = input[0], input[1]
h, w = image.shape[:2]
if isinstance(self.output_size, int):
if h > w:
new_h, new_w = self.output_size * h / w, self.output_size
else:
new_h, new_w = self.output_size, self.output_size * w / h
else:
new_h, new_w = self.output_size
new_h, new_w = int(new_h), int(new_w)
img = transform.resize(image, (new_h, new_w))
# h and w are swapped for landmarks because for images,
# x and y axes are axis 1 and 0 respectively
# landmarks = landmarks * [new_w / w, new_h / h]
return img, labels
''' source code from pytorch '''
class Rescale(object):
"""Rescale the image in a sample to a given size.
Args:
output_size (tuple or int): Desired output size. If tuple, output is
matched to output_size. If int, smaller of image edges is matched
to output_size keeping aspect ratio the same.
"""
def __init__(self, output_size):
assert isinstance(output_size, (int, tuple))
self.output_size = output_size
def __call__(self, input):
image, labels = input[0], input[1]
h, w = image.shape[:2]
if isinstance(self.output_size, int):
if h > w:
new_h, new_w = self.output_size * h / w, self.output_size
else:
new_h, new_w = self.output_size, self.output_size * w / h
else:
new_h, new_w = self.output_size
new_h, new_w = int(new_h), int(new_w)
img = transform.resize(image, (new_h, new_w))
# h and w are swapped for landmarks because for images,
# x and y axes are axis 1 and 0 respectively
# landmarks = landmarks * [int(new_w / w), int(new_h / h)]
return img, labels
class ToTensor(object):
"""Convert ndarrays in sample to Tensors."""
def __call__(self, input):
image, labels = input[0], input[1]
# swap color axis because
# numpy image: H x W x C
# torch image: C X H X
image = image.transpose((2, 0, 1))
return torch.from_numpy(image), torch.from_numpy(labels)
''' Normalize Images in torch.tensor'''
class CustomNormalize(object):
def __init__(self, mean, std):
self.mean = torch.from_numpy(np.resize(mean, (3, 1, 1)))
self.std = torch.from_numpy(np.resize(std, (3, 1, 1)))
def __call__(self, input):
image, labels = input[0], input[1]
output = (image - self.mean) / self.std
return output, labels
this class rearranges the size of the image, and normalizes this image to make it easier to process. There are also classes called CustomNormalize and ToTensor, and these classes. A numpy image is height x width x color, while torch image is color x height x width, which is what the transpose line of code does. Normalizing images resizes them and result in the difference from the mean. The loss is represented as a basic squared error difference approach, and the location accuracy count is based on the distance between the output and the labels. Network.py sets a simple neural network, and has a dense net network where each layer is interconnected with each other.
The last class I want to discuss is preprocessing, which performs everything that needs to be done for a neural network to run.
import os
import cv2
import math
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import imutils
from scipy.interpolate import interp2d
from functools import cmp_to_key
from imutils import contours
from skimage import measure
import utils
# extract max
# grow the seed
# find illuminated part
# repeat
class JPGLabeler():
def _get_threshold_img(self, img):
gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
blurred_img = cv2.GaussianBlur(gray_img, (11, 11), 0)
extreme = blurred_img.max()
img_mean = np.mean(img)
if img_mean < 85:
thresh_img = cv2.threshold(blurred_img, extreme*5/12, extreme, cv2.THRESH_BINARY)[1]
else:
thresh_img = cv2.threshold(blurred_img, extreme*2/3, extreme, cv2.THRESH_BINARY)[1]
thresh_img = cv2.erode(thresh_img, None, iterations=2)
thresh_img = cv2.dilate(thresh_img, None, iterations=4)
return thresh_img
def _find_connected_components(self, img, thresh_img, N):
labels = measure.label(thresh_img, connectivity=1, background=0)
masks = []
# loop over the unique components
for label in np.unique(labels):
# if this is the background label, ignore it
if label == 0:
continue
# otherwise, construct the label mask and count the
# number of pixels
labelMask = np.zeros(thresh_img.shape, dtype="uint8")
labelMask[labels == label] = 255
numPixels = cv2.countNonZero(labelMask)
# if the number of pixels in the component is sufficiently
# large, then add it to our mask of "large blobs"
if numPixels > 300:
masks.append([np.resize(labelMask, thresh_img.shape), np.sum(img[labelMask == 255])/np.count_nonzero(img[labelMask == 255]), numPixels])
# If not enough number of light, discard the image
if len(masks) < N:
return None
def compare_connected_components(x, y):
if x[1] > y[1]:
return 1
elif x[1] < y[1]:
return -1
elif x[2] > y[2]:
return 1
elif x[2] < y[2]:
return -1
else:
return 0
# otherwise, return the first N largest light
masks = sorted(masks, key =cmp_to_key(compare_connected_components), reverse=True)
return masks[0:N]
def _label_img(self, masks, img):
# for l in labels:
# l = [[(x1, y1, 1), (ax1, ax2), rotation1, ('R', 'G', 'B')],
# [(x2, y2, 1), (ax1, ax2), rotation2, ('R', 'G', 'B')],
# ...,
# [(xN, yN, 1), (ax1, ax2), rotationN, ('R', 'G', 'B')]]
labels = []
# for each illumination (denoted by mask)
for mask, _, _ in masks:
# find the contour
cnts = cv2.findContours(mask.copy(), cv2.RETR_EXTERNAL,
cv2.CHAIN_APPROX_SIMPLE)
cnts = imutils.grab_contours(cnts)
cnts = contours.sort_contours(cnts)[0]
# loop over the contours
for (i, c) in enumerate(cnts):
# draw the bright spot on the image
rect = cv2.minAreaRect(c)
x, y = rect[0]
r, g, b =np.divide(np.sum(img[mask == 255], axis=(0)), \
np.count_nonzero(img[mask==255], axis=(0)))
labels.append([x, y, 1 , rect[1][0], rect[1][1], rect[2], r, g, b])
return labels
def _annotate_img(self, masks, img):
for mask, _, _ in masks:
cnts = cv2.findContours(mask.copy(), cv2.RETR_EXTERNAL,
cv2.CHAIN_APPROX_SIMPLE)
cnts = imutils.grab_contours(cnts)
cnts = contours.sort_contours(cnts)[0]
# loop over the contours
for (i, c) in enumerate(cnts):
# draw the bright spot on the image
rect = cv2.minAreaRect(c)
elp = cv2.ellipse(img, rect, color = (0, 0, 255), thickness=2)
return img
def generate_annotated_img(self, img, N=3):
# preprocessing
thresh_img = self._get_threshold_img(img)
masks = self._find_connected_components(img, thresh_img, N)
if masks is None:
return None
img = self._annotate_img(masks, img)
return img
def generate_labels(self, img, N=3):
# preprocessing
thresh_img = self._get_threshold_img(img)
masks = self._find_connected_components(img, thresh_img, N)
if masks is None:
return None
label = self._label_img(masks, img)
return label
class EXRLabeler():
def _get_threshold_img(self, img):
# gray_img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
# blurred_img = cv2.GaussianBlur(gray_img, (11, 11), 0)
img_one_channel = np.sum(img, axis=-1)
extreme = img_one_channel.max()
# img_mean = np.mean(img)
# if img_mean < 85:
# thresh_img = cv2.threshold(blurred_img, extreme*5/12, extreme, cv2.THRESH_BINARY)[1]
# else:
# thresh_img = cv2.threshold(blurred_img, extreme*2/3, extreme, cv2.THRESH_BINARY)[1]
# thresh_img = cv2.erode(thresh_img, None, iterations=2)
# thresh_img = cv2.dilate(thresh_img, None, iterations=4)
thresh_img = np.zeros(shape=img_one_channel.shape, dtype='uint8')
thresh_img[img_one_channel > extreme*1/100] = 255
return thresh_img
def _find_connected_components(self, img, thresh_img, N):
labels = measure.label(thresh_img, connectivity=1, background=0)
masks = []
# loop over the unique components
for label in np.unique(labels):
# if this is the background label, ignore it
if label == 0:
continue
# otherwise, construct the label mask and count the
# number of pixels
labelMask = np.zeros(thresh_img.shape, dtype="uint8")
labelMask[labels == label] = 255
numPixels = cv2.countNonZero(labelMask)
# if the number of pixels in the component is sufficiently
# large, then add it to our mask of "large blobs"
if numPixels > 300:
masks.append([np.resize(labelMask, thresh_img.shape), np.sum(img[labelMask == 255])/np.count_nonzero(img[labelMask == 255]), numPixels])
# If not enough number of light, discard the image
if len(masks) < N:
return None
def compare_connected_components(x, y):
if x[1] > y[1]:
return 1
elif x[1] < y[1]:
return -1
elif x[2] > y[2]:
return 1
elif x[2] < y[2]:
return -1
else:
return 0
# otherwise, return the first N largest light
masks = sorted(masks, key =cmp_to_key(compare_connected_components), reverse=True)
return masks[0:N]
def _label_img(self, masks, img):
# for l in labels:
# l = [[(x1, y1, 1), (ax1, ax2), rotation1, ('R', 'G', 'B')],
# [(x2, y2, 1), (ax1, ax2), rotation2, ('R', 'G', 'B')],
# ...,
# [(xN, yN, 1), (ax1, ax2), rotationN, ('R', 'G', 'B')]]
labels = []
# for each illumination (denoted by mask)
for mask, _, _ in masks:
# find the contour
cnts = cv2.findContours(mask.copy(), cv2.RETR_EXTERNAL,
cv2.CHAIN_APPROX_SIMPLE)
cnts = imutils.grab_contours(cnts)
cnts = contours.sort_contours(cnts)[0]
# loop over the contours
for (i, c) in enumerate(cnts):
# find a rectangle (ellipse)
rect = cv2.minAreaRect(c)
x, y = rect[0]
r, g, b =np.divide(np.sum(img[mask == 255], axis=(0)), \
np.count_nonzero(img[mask==255], axis=(0)))
labels.append([x, y, 1 , rect[1][0], rect[1][1], rect[2], r, g, b])
return labels
def _annotate_img(self, masks, img):
max_intensity = np.max(img)
for mask, _, _ in masks:
cnts = cv2.findContours(mask.copy(), cv2.RETR_EXTERNAL,
cv2.CHAIN_APPROX_SIMPLE)
cnts = imutils.grab_contours(cnts)
cnts = contours.sort_contours(cnts)[0]
# loop over the contours
for (i, c) in enumerate(cnts):
# draw the bright spot on the image
rect = cv2.minAreaRect(c)
elp = cv2.ellipse(img, rect, color = (0, 0, 255), thickness=2)
return img
def generate_annotated_img(self, img, N=3):
# preprocessing
thresh_img = self._get_threshold_img(img)
masks = self._find_connected_components(img, thresh_img, N)
if masks is None:
return None
img = self._annotate_img(masks, img)
return img
def generate_labels(self, img, N=3):
# preprocessing
thresh_img = self._get_threshold_img(img)
masks = self._find_connected_components(img, thresh_img, N)
if masks is None:
return None
label = self._label_img(masks, img)
return label
class DataGenerator():
def __init__(self):
pass
def generate_feature_and_label(self, source_path, dest_path, img_format='jpg'):
img_format = str.lower(img_format)
if img_format == 'jpg' or img_format == 'jpeg':
labeler = JPGLabeler()
elif img_format == 'exr':
labeler = EXRLabeler()
img_names = os.listdir(os.path.join(source_path))
discarded_data_set = set()
img_names.sort()
label_dict = dict()
img_index = 0
labeled_images = []
labels = []
# generate labels
for img_name in img_names:
img = cv2.imread(os.path.join(source_path, img_name))
label = labeler.generate_labels(img)
if label is None:
discarded_data_set.add(img_name)
continue
# label_dict[img_index] = label
labels.append(label)
labeled_images.append(img)
img_index += 1
# store labeled images
# with open(os.path.join(label_path, 'labels.txt'), 'w') as f:
# f.write(json.dumps(label_dict))
labeled_images = np.array(labeled_images)
labels = np.array(labels, dtype='float64')
# num of X == num of Y
assert labels.shape[0] == labeled_images.shape[0]
# ratio of test sample
ratio = 0.2
num_of_samples = labeled_images.shape[0]
num_of_test_samples = int(ratio * num_of_samples)
num_of_train_samples = int(num_of_samples - num_of_test_samples)
train_matrix = labeled_images[0:num_of_train_samples]
train_labels = labels[0:num_of_train_samples]
test_matrix = labeled_images[num_of_train_samples:-1]
test_labels = labels[num_of_train_samples:-1]
np.save(os.path.join(dest_path, 'train_feature_matrix.npy'), train_matrix, allow_pickle=True)
np.save(os.path.join(dest_path, 'train_label.npy'), train_labels, allow_pickle=True)
np.save(os.path.join(dest_path, 'test_feature_matrix.npy'), test_matrix, allow_pickle=True)
np.save(os.path.join(dest_path, 'test_label.npy'), test_labels, allow_pickle=True)
class NaiveCropper():
def __init__(self, environment_map):
self.map = environment_map
self.height = environment_map.shape[0]
self.width = environment_map.shape[1]
self.type = environment_map.dtype
self.channel = environment_map.shape[2]
def generate_image(self, view_direction_theta, view_direction_phi, FOV_horizontal, image_width, image_height):
"""
Return a cropped image
Parameters
----------
view_direction_theta : float.
The theta of the view vector, which defines in spherical coordinate
view_direction_phi : float.
The phi of the view vector, which defines in spherical coordinate
FOV_horizontal : float.
The field of view of the cropped image horizontally
image_height, image_width : int.
Returns
-------
out : 2darray
A 2darray with size = image_heigh x image_width
Examples
--------
image = handle.generate_image(0, 0, math.pi/3, 1080, 720)
"""
FOV_vertical = FOV_horizontal/image_width * image_height
theta_start = view_direction_theta - FOV_horizontal/2
phi_start = view_direction_phi - FOV_vertical/2
image_center = np.array([1 * np.cos(view_direction_theta) * np.sin(view_direction_phi), \
1 * np.sin(view_direction_theta) * np.sin(view_direction_phi),1 * np.cos(view_direction_phi)])
horizontal_size_half = np.tan(FOV_horizontal/2)
vertical_size_half = np.tan(FOV_vertical/2)
horizontal_displacement = np.array([-horizontal_size_half * np.sin(view_direction_theta), horizontal_size_half * np.cos(view_direction_theta), 0])
vertical_displacement = np.array([-vertical_size_half*np.cos(view_direction_phi)*np.cos(view_direction_theta), -vertical_size_half * np.cos(view_direction_phi) * np.sin(view_direction_theta), \
vertical_size_half*np.sin(view_direction_phi)])
top_left = image_center + vertical_displacement - horizontal_displacement
top_right = image_center + vertical_displacement + horizontal_displacement
bot_left = image_center - vertical_displacement - horizontal_displacement
bot_right = image_center - vertical_displacement + horizontal_displacement
# form the image plane
img_rect = np.zeros(shape=(image_height,image_width, 3), dtype=np.float32)
left = top_left
right = top_right
height_array = np.linspace(top_left ,bot_left, image_height)
for i in range(image_height):
left = height_array[i]
right = left + horizontal_displacement*2
img_rect[i,:,:] = np.linspace(left,right,image_width)
# transfer the image plane back to spherical
theta = np.arctan2(img_rect[:,:,1], img_rect[:,:,0])
theta[theta < 0] = theta[theta < 0] + math.pi*2
phi = np.arccos(img_rect[:,:,2]/np.linalg.norm(img_rect, axis=2))
# Read from the original image
# u = np.floor(theta * self.width / math.pi)
# v = np.floor(phi * self.height / (math.pi/2))
img = np.zeros(shape=(image_height,image_width, self.channel), dtype=self.type)
for i in range(self.channel):
f = interp2d(np.linspace(0, math.pi * 2,self.width), np.linspace(0,math.pi, self.height), self.map[:,:,i], kind='linear')
for j in range(theta.shape[0]):
for k in range(theta.shape[1]):
img[j,k,i] = (f(theta[j,k], phi[j,k]))[0]
return img
def generate_image_detail(self, starting_theta, starting_phi, theta_size, phi_size, height_resolution, width_resolution):
# create four points on the sphere
top_left_sph = np.array([1, starting_theta, starting_phi], dtype = np.float32)
top_right_sph = np.array([1, starting_theta + theta_size, starting_phi], dtype = np.float32)
bot_left_sph = np.array([1, starting_theta, starting_phi+phi_size], dtype = np.float32)
bot_right_sph = np.array([1, starting_theta + theta_size, starting_phi+phi_size], dtype = np.float32)
# transfer them to rect system
top_left = np.array([top_left_sph[0] * np.cos(top_left_sph[1]) * np.sin(top_left_sph[2]), \
top_left_sph[0] * np.sin(top_left_sph[1]) * np.sin(top_left_sph[2]), top_left_sph[0] * np.cos(top_left_sph[2])])
top_right = np.array([top_right_sph[0] * np.cos(top_right_sph[1]) * np.sin(top_right_sph[2]), \
top_right_sph[0] * np.sin(top_right_sph[1]) * np.sin(top_right_sph[2]), top_right_sph[0] * np.cos(top_right_sph[2])])
bot_left = np.array([bot_left_sph[0] * np.cos(bot_left_sph[1]) * np.sin(bot_left_sph[2]), \
bot_left_sph[0] * np.sin(bot_left_sph[1]) * np.sin(bot_left_sph[2]), bot_left_sph[0] * np.cos(bot_left_sph[2])])
bot_right = np.array([bot_right_sph[0] * np.cos(bot_right_sph[1]) * np.sin(bot_right_sph[2]), \
bot_right_sph[0] * np.sin(bot_right_sph[1]) * np.sin(bot_right_sph[2]), bot_right_sph[0] * np.cos(bot_right_sph[2])])
# form the image plane
img_width_difference = top_right - top_left
img_height_difference = bot_left - top_left
img_rect = np.zeros(shape=(height_resolution,width_resolution, 3), dtype=np.float32)
left = top_left
right = top_right
height_array = np.linspace(top_left ,bot_left, height_resolution)
for i in range(height_resolution):
left = height_array[i]
right = left + img_width_difference
img_rect[i,:,:] = np.linspace(left,right,width_resolution)
# transfer the image plane back to spherical
theta = np.arctan2(img_rect[:,:,1], img_rect[:,:,0])
theta[theta < 0] = theta[theta < 0] + math.pi/2
phi = np.arccos(img_rect[:,:,2]/np.linalg.norm(img_rect, axis=2))
# print(theta)
# print(phi)
# Read from the original image
u = np.floor(theta * self.width / math.pi)
v = np.floor(phi * self.height / (math.pi/2))
img = np.zeros(shape=(height_resolution,width_resolution, self.channel), dtype=self.type)
# x = np.linspace(0, math.pi,self.width)
# y = np.linspace(0, math.pi/2,self.height)
# points_x, points_y = np.meshgrid(x, y)
# points_x_lin = np.reshape(points_x, -1)
# points_y_lin = np.reshape(points_y, -1)
# points = np.stack((points_x_lin, points_y_lin), axis = 1)
# print(points.shape)
# grid_x, grid_y = np.meshgrid(u,v)
# print(u.shape)
for i in range(self.channel):
f = interp2d(np.linspace(0, math.pi * 2,self.width), np.linspace(0,math.pi, self.height), self.map[:,:,i], kind='linear')
for j in range(theta.shape[0]):
img[j,:,i] = (f(theta[j,:], phi[j,:]))[0,:]
# print(img[:,:,i])
# for i in range(self.channel):
# img[:,:,i] = self.map[v,u,i]
# print(img[:,:,i])
return img
# test
if __name__ == '__main__':
# preprocessing
img = cv2.imread('./test_conversion/5.jpg')
l = ImageLabeler()
thresh_img = l._get_threshold_img(img)
masks = l._find_connected_components(img, thresh_img, 3)
img = l._annotate_image(masks, img)
for mask in masks:
plt.imshow(mask)
plt.show()
# exr => label, jpeg => train
# ambient light
# chopped image
get_threshold_image erodes, create, and dilates an image. Find_connected components assigns a full value to the array of anything that is connected. The label img labels stuff for each illumination.
Comments
Post a Comment