import base64
from types import new_class
import ipywidgets as widgets
from pathlib import PurePath
from jupyter_bbox_widget import BBoxWidget
from ipywidgets import Layout
import numpy as np
[docs]
class Labeler:
def __init__(self, dataset=None):
self.dataset = dataset
[docs]
def StartPyLaber(self, new_classes=None, image=None, yolo_model=None):
"""Display the bbox widget loaded with images and annotations from this dataset."""
if "google.colab" in str(get_ipython()):
from google.colab import output
output.enable_custom_widget_manager()
dataset = self.dataset
widget_output = None
files = dataset.df.img_filename.unique()
files = files.tolist()
global file_index
if image == None:
file_index = 0
image = files[0]
else:
file_index = files.index(image)
def GetBBOXs(image):
# Make a dataframe with the annotations for a single image
img_df = dataset.df.loc[dataset.df["img_filename"] == image]
img_df_subset = img_df[
[
"cat_name",
"ann_bbox_height",
"ann_bbox_width",
"ann_bbox_xmin",
"ann_bbox_ymin",
]
]
# Rename the columns to match the format used by jupyter_bbox_widget
img_df_subset.columns = ["label", "height", "width", "x", "y"]
# Drop rows that have NaN, invalid bounding boxes
img_df_subset = img_df_subset.dropna()
bboxes_dict = img_df_subset.to_dict(orient="records")
return bboxes_dict
def GetImageLabel(image):
"""Returns the imagename or imagename + (not annotated)"""
# Make a dataframe with the annotations for a single image
img_df = dataset.df.loc[dataset.df["img_filename"] == image]
if img_df.iloc[0].annotated == 1:
return image
else:
return f"{image} (not annotated)"
return
bboxes_dict = GetBBOXs(image)
img_folder = dataset.df.loc[dataset.df["img_filename"] == image].iloc[0][
"img_folder"
]
file_paths = [
str(PurePath(dataset.path_to_annotations, img_folder, file))
for file in files
]
def encode_image(filepath):
with open(filepath, "rb") as f:
image_bytes = f.read()
encoded = str(base64.b64encode(image_bytes), "utf-8")
return "data:image/jpg;base64," + encoded
def UpdateCategoryList(cat_dict, new_categories):
from math import isnan
# Remove invalid entries
cat_dict.pop("", None)
# cat_dict = {k: v for k, v in cat_dict.items() if not isnan(v)}
for cat in new_categories:
if len(cat_dict) == 0:
cat_dict[cat] = "0"
elif cat in list(cat_dict.keys()):
continue
else:
# Create a new cat id that 1+ the highest cat id value
new_cat_id = max([int(v) for v in cat_dict.values()]) + 1
cat_dict[cat] = str(new_cat_id)
return cat_dict
def on_submit(b):
# save annotations for current image
import pandas as pd
global widget_output
global file_index
img_filename = files[file_index]
widget_output = pd.DataFrame.from_dict(w_bbox.bboxes)
# Check if there are any bounding boxes for the current image:
if not widget_output.empty:
# If there are bounding boxes then add them to the dataset
widget_output = widget_output.rename(
columns={
"label": "cat_name",
"height": "ann_bbox_height",
"width": "ann_bbox_width",
"x": "ann_bbox_xmin",
"y": "ann_bbox_ymin",
}
)
widget_output["ann_area"] = (
widget_output["ann_bbox_height"] * widget_output["ann_bbox_width"]
)
widget_output["cat_name"] = widget_output["cat_name"].astype("string")
categories = dict(zip(dataset.df.cat_name, dataset.df.cat_id))
categories = UpdateCategoryList(
categories, list(widget_output.cat_name)
)
widget_output["cat_id"] = widget_output["cat_name"].map(categories)
else:
# Build a entry in the table with empty bounding box columns
widget_output = pd.DataFrame(
[],
columns=[
"cat_name",
"cat_id",
"ann_bbox_height",
"ann_bbox_width",
"ann_bbox_xmin",
"ann_bbox_ymin",
"ann_area",
],
)
widget_output.loc[0] = [
"",
np.NaN,
np.NaN,
np.NaN,
np.NaN,
np.NaN,
np.NaN,
]
widget_output["img_filename"] = str(img_filename)
widget_output["img_filename"] = widget_output["img_filename"].astype(
"string"
)
widget_output.index.name = "id"
img_df = dataset.df.loc[dataset.df["img_filename"] == img_filename]
# Get the metadata associated with this image from the dataset
metadata = img_df.iloc[0].to_frame().T
metadata["img_filename"] = metadata["img_filename"].astype("string")
# Drop the fields that are in the widget_output dataframe
metadata.drop(
[
"cat_name",
"cat_id",
"ann_area",
"ann_bbox_height",
"ann_bbox_width",
"ann_bbox_xmin",
"ann_bbox_ymin",
],
axis=1,
inplace=True,
)
widget_output = widget_output.merge(
metadata, left_on="img_filename", right_on="img_filename"
)
widget_output = widget_output[dataset.df.columns]
# Set annotated = 1, which means the annotates have been reviewed and accepted
widget_output["annotated"] = 1
# Now we have a dataframe with output of the bbox widget
# Drop the current annotations for the image and add the the new ones
dataset.df.drop(
dataset.df[dataset.df["img_filename"] == img_filename].index,
inplace=True,
)
dataset.df.reset_index(drop=True, inplace=True)
dataset.df = dataset.df.append(widget_output).reset_index(drop=True)
# move on to the next file
on_next(b)
def on_next(b):
global file_index
file_index += 1
# open new image in the widget
image_file = file_paths[file_index]
w_bbox.image = encode_image(image_file)
w_bbox.bboxes = GetBBOXs(files[file_index])
progress_label.value = f"{file_index+1} / {len(files)}"
current_img_name_label.value = GetImageLabel(files[file_index])
def on_back(b):
global file_index
file_index -= 1
# open new image in the widget
image_file = file_paths[file_index]
w_bbox.image = encode_image(image_file)
w_bbox.bboxes = GetBBOXs(files[file_index])
progress_label.value = f"{file_index+1} / {len(files)}"
current_img_name_label.value = GetImageLabel(files[file_index])
def on_add_class(b):
if new_class_text.value.strip() != "":
class_list = list(w_bbox.classes)
class_list.append(new_class_text.value)
w_bbox.classes = list(set(class_list))
new_class_text.value = ""
def on_predict(b):
global file_index
image_file = file_paths[file_index]
result = yolo_model(image_file)
result = result.pandas().xyxy[0]
result["width"] = result.xmax - result.xmin
result["height"] = result.ymax - result.ymin
result.drop(["class", "confidence", "xmax", "ymax"], axis=1, inplace=True)
result.columns = ["x", "y", "label", "width", "height"]
result = result[["label", "height", "width", "x", "y"]]
bboxes_dict = result.to_dict(orient="records")
w_bbox.bboxes = bboxes_dict
# Add the predicted classes to the widget so they can be selected by the user
w_bbox.classes = list(set(w_bbox.classes + list(result["label"])))
if new_classes:
classes = dataset.analyze.classes + new_classes
else:
classes = dataset.analyze.classes
# remove empty labels and duplicate labels
classes = list(
set([str(c).strip() for c in classes if len(str(c).strip()) > 0])
)
# Load BBoxWidget for first load on page
w_bbox = BBoxWidget(
image=encode_image(file_paths[file_index]),
classes=classes,
bboxes=bboxes_dict,
hide_buttons=True,
)
progress_txt = f"{file_index+1} / {len(files)}"
left_arrow_btn = widgets.Button(
icon="fa-arrow-left", layout=Layout(width="35px")
)
progress_label = widgets.Label(value=progress_txt)
current_img_name_label = widgets.Label(value=GetImageLabel(image))
right_arrow_btn = widgets.Button(
icon="fa-arrow-right", layout=Layout(width="35px")
)
save_btn = widgets.Button(
icon="fa-check", description="Save", layout=Layout(width="70px")
)
predict_btn = widgets.Button(
icon="fa-eye", description="Predict", layout=Layout(width="100px")
)
train_btn = widgets.Button(
icon="fa-refresh", description="Train", layout=Layout(width="100px")
)
add_class_label = widgets.Label(value="Add class:")
new_class_text = widgets.Text(layout=Layout(width="200px"))
plus_btn = widgets.Button(icon="fa-plus", layout=Layout(width="35px"))
button_row_list = [left_arrow_btn, progress_label, right_arrow_btn, save_btn]
# If model arg is empty hide predict and train buttons
if yolo_model != None:
button_row_list = button_row_list + [predict_btn] # train_btn
button_row = widgets.HBox(button_row_list)
current_img_details_row = widgets.HBox([current_img_name_label])
bottom_row = widgets.HBox([add_class_label, new_class_text, plus_btn])
pylabler = widgets.VBox(
[current_img_details_row, button_row, w_bbox, bottom_row]
)
save_btn.on_click(on_submit)
left_arrow_btn.on_click(on_back)
right_arrow_btn.on_click(on_next)
predict_btn.on_click(on_predict)
plus_btn.on_click(on_add_class)
new_class_text.on_submit(on_add_class)
# Returning the container will show the widgets
return pylabler