def fit(self, completions, workdir=None, **kwargs):
annotations = []
for completion in completions:
if is_skipped(completion):
continue
image_path = get_image_local_path(completion['data'][self.value])
image_label = get_choice(completion)
annotations.append((image_path, image_label))
# Create dataset
ds = tf.data.Dataset.from_tensor_slices(annotations)
def prepare_item(item):
label = tf.argmax(item[1] == self.labels)
img = tf.io.read_file(item[0])
img = tf.image.decode_jpeg(img, channels=3)
img = tf.image.resize(img, [self.image_height, self.image_width])
return img, label
ds = ds.map(prepare_item, num_parallel_calls=tf.data.AUTOTUNE)
ds = ds.cache().shuffle(buffer_size=1000).batch(
self.batch_size).prefetch(buffer_size=tf.data.AUTOTUNE)
self.model.compile(optimizer=tf.keras.optimizers.Adam(),
loss=tf.keras.losses.SparseCategoricalCrossentropy(
from_logits=True),
metrics=['acc'])
self.model.fit(ds, epochs=self.epochs)
model_file = os.path.join(workdir, 'checkpoint')
self.model.save_weights(model_file)
return {'model_file': model_file}
def predict(self, tasks, **kwargs):
assert len(tasks) == 1
task = tasks[0]
image_path = get_image_local_path(task['data'][self.value],
project_dir=self.project_path)
model_results = inference_detector(self.model, image_path)
results = []
all_scores = []
img_width, img_height = get_image_size(image_path)
for bboxes, label in zip(model_results, self.model.CLASSES):
output_label = self.label_map.get(label, label)
if output_label not in self.labels_in_config:
print(output_label + ' label not found in project config.')
continue
for bbox in bboxes:
bbox = list(bbox)
if not bbox:
continue
score = float(bbox[-1])
if score < self.score_thresh:
continue
x, y, xmax, ymax = bbox[:4]
results.append({
'from_name': self.from_name,
'to_name': self.to_name,
'type': 'rectanglelabels',
'value': {
'rectanglelabels': [output_label],
'x': x / img_width * 100,
'y': y / img_height * 100,
'width': (xmax - x) / img_width * 100,
'height': (ymax - y) / img_height * 100
},
'score': score
})
all_scores.append(score)
avg_score = sum(all_scores) / len(all_scores)
return [{'result': results, 'score': avg_score}]
def predict(self, tasks, **kwargs):
image_path = get_image_local_path(tasks[0]['data'][self.value])
image = Image.open(image_path).resize(
(self.image_width, self.image_height))
image = np.array(image) / 255.0
result = self.model.predict(image[np.newaxis, ...])
predicted_label_idx = np.argmax(result[0], axis=-1)
predicted_label_score = result[0][predicted_label_idx]
predicted_label = self.labels[predicted_label_idx]
return [{
'result': [{
'from_name': self.from_name,
'to_name': self.to_name,
'type': 'choices',
'value': {
'choices': [str(predicted_label.numpy(), 'utf-8')]
}
}],
'score':
float(predicted_label_score)
}]