def main(args):
probability_threshold = 50.0
if (not args.input_tsv_file):
raise ValueError(
'You must supply input TSV file with --input_tsv_file.')
if (not args.output_tsv_file):
raise ValueError(
'You must supply output TSV file with --output_tsv_file.')
if (not os.path.isfile(args.input_tsv_file)):
return (False)
model_root_dir = NetworkFactory.model_deploy_dir()
last_network = 'ONet'
face_detector = FaceDetector(last_network, model_root_dir)
classifier_object = Classifier()
if (not classifier_object.load_dataset(args.dataset_dir)):
return (False)
if (not classifier_object.load_model(args.checkpoint_path, args.model_name,
args.gpu_memory_fraction)):
return (False)
network_size = classifier_object.network_image_size()
number_of_images = 0
good_images = 0
input_tsv_file = open(args.input_tsv_file, 'r')
output_tsv_file = open(args.output_tsv_file, 'w')
while (True):
input_data = input_tsv_file.readline().strip()
if (not input_data):
break
number_of_images += 1
fields = input_data.split('\t')
line_number = str(fields[0])
image_string = fields[1]
decoded_image_string = base64.b64decode(image_string)
image_data = np.fromstring(decoded_image_string, dtype=np.uint8)
input_image = cv2.imdecode(image_data, cv2.IMREAD_COLOR)
height, width, channels = input_image.shape
cv2.imwrite('image.png', input_image)
#misc.imsave('image.png', input_image)
input_image = misc.imread('image.png')
input_clone = np.copy(input_image)
boxes_c, landmarks = face_detector.detect(input_clone)
face_probability = 0.0
found = False
crop_box = []
for index in range(boxes_c.shape[0]):
if (boxes_c[index, 4] > face_probability):
found = True
face_probability = boxes_c[index, 4]
bounding_box = boxes_c[index, :4]
crop_box = [
int(max(bounding_box[0], 0)),
int(max(bounding_box[1], 0)),
int(min(bounding_box[2], width)),
int(min(bounding_box[3], height))
]
if (found):
cropped_image = input_image[crop_box[1]:crop_box[3],
crop_box[0]:crop_box[2], :]
else:
cropped_image = input_image
#resized_image = cv2.resize(cropped_image, (network_size, network_size), interpolation=cv2.INTER_LINEAR)
resized_image = misc.imresize(cropped_image,
(network_size, network_size),
interp='bilinear')
class_names_probabilities = classifier_object.classify(
resized_image, print_results=False)
predicted_name = ''
probability = 0.0
if (len(class_names_probabilities) > 0):
names = map(operator.itemgetter(0), class_names_probabilities)
probabilities = map(operator.itemgetter(1),
class_names_probabilities)
predicted_name = str(names[0])
probability = probabilities[0]
if ((probability > probability_threshold)
or (probability >
(probabilities[1] + probabilities[2] / 2.0))):
good_images += 1
print(number_of_images, ', predicted_name -', predicted_name,
', probability -', probability)
print('Accuracy = ', (good_images * 100.0) / number_of_images, ' for ',
number_of_images, ' images.')
#cv2.imshow('image', cropped_image)
#cv2.waitKey();
output_tsv_file.write(line_number + '\t' + str(predicted_name) + '\t' +
str(probability) + '\n')
print('Accuracy = ', (good_images * 100.0) / number_of_images, ' for ',
number_of_images, ' images.')
return (True)
def main(args):
output_dir = os.path.expanduser(args.output_dir)
if(not os.path.exists(output_dir)):
os.mkdir(output_dir)
is_processed = {}
probability_threshold = [ 0.95, 0.90, 0.85, 0.80 ]
if(not args.input_tsv_file):
raise ValueError('You must supply input TSV file with --input_tsv_file.')
if(not os.path.isfile(args.input_tsv_file)):
return(False)
model_root_dir = NetworkFactory.model_deploy_dir()
last_network='ONet'
face_detector = FaceDetector(last_network, model_root_dir)
classifier_object = Classifier()
if(not classifier_object.load_dataset(args.dataset_dir)):
return(False)
if(not classifier_object.load_model(args.checkpoint_path, args.model_name, args.gpu_memory_fraction)):
return(False)
network_size = classifier_object.network_image_size()
celebrity_count = 0
for current_threshold in probability_threshold:
input_tsv_file = open(args.input_tsv_file, 'r')
while( True ):
input_data = input_tsv_file.readline().strip()
if( not input_data ):
break
fields = input_data.split('\t')
class_name = str(fields[2])
if class_name in is_processed.keys():
continue
image_string = fields[1]
image_search_rank = fields[3]
decoded_image_string = base64.b64decode(image_string)
image_data = np.fromstring(decoded_image_string, dtype=np.uint8)
input_image = cv2.imdecode(image_data, cv2.IMREAD_COLOR)
height, width, channels = input_image.shape
class_dir = fields[2]
img_name = class_dir + '.png'
cv2.imwrite('image.png', input_image)
#misc.imsave('image.png', input_image)
input_image = misc.imread('image.png')
input_clone = np.copy(input_image)
boxes_c, landmarks = face_detector.detect(input_clone)
face_probability = 0.0
found = False
crop_box = []
for index in range(boxes_c.shape[0]):
if(boxes_c[index, 4] > face_probability):
found = True
face_probability = boxes_c[index, 4]
bounding_box = boxes_c[index, :4]
crop_box = [int(max(bounding_box[0],0)), int(max(bounding_box[1],0)), int(min(bounding_box[2], width)), int(min(bounding_box[3], height))]
if(found):
cropped_image = input_image[crop_box[1]:crop_box[3],crop_box[0]:crop_box[2],:]
else:
cropped_image = input_image
#resized_image = cv2.resize(cropped_image, (network_size, network_size), interpolation=cv2.INTER_LINEAR)
resized_image = misc.imresize(cropped_image, (network_size, network_size), interp='bilinear')
class_names_probabilities = classifier_object.classify(resized_image, print_results=False)
predicted_name = ''
probability = 0.0
if(len(class_names_probabilities) > 0):
names = map(operator.itemgetter(0), class_names_probabilities)
probabilities = map(operator.itemgetter(1), class_names_probabilities)
predicted_name = str(names[0])
probability = probabilities[0]
if( class_name != predicted_name ):
continue
if(probability < current_threshold):
continue
is_processed[class_name] = True
full_class_dir = os.path.join(output_dir, class_dir)
if not os.path.exists(full_class_dir):
os.mkdir(full_class_dir)
celebrity_count = celebrity_count + 1
full_path = os.path.join(full_class_dir, img_name)
cv2.imwrite(full_path, resized_image)
#cv2.imshow('image', cropped_image)
#cv2.waitKey();
print('Processed ', celebrity_count, 'celebrities.')
return(True)