import streamlit as st
import numpy as np
import plotly.express as px
import cv2
from src.error_analysis import ErrorAnalysis, transform_gt_bbox_format
import yaml
import os
from src.confusion_matrix import ConfusionMatrix
from plotly.subplots import make_subplots
import plotly.graph_objects as go
import pandas as pd
def amend_cm_df(cm_df, labels_dict):
"""Helper function to amend the index and column name for readability
Example - index currently is 0, 1 ... -> GT - person
Likewise in Column - 0, 1 ... -> Pred - person etc
Args:
cm_df (_type_): confusion matrix dataframe.
labels_dict (_type_): dictionary of the class labels
Returns:
cm_df: confusion matrix dataframe with index and column names filled
"""
index_list = list(labels_dict.values())
index_list.append("background")
cm_df = cm_df.set_axis([f"GT - {elem}" for elem in index_list])
cm_df = cm_df.set_axis([f"Pred - {elem}" for elem in index_list], axis=1)
cm_df = cm_df.astype(int)
return cm_df
def find_top_left_pos(mask):
"""gets the top left position of the mask
Args:
mask (_type_): _description_
Returns:
_type_: _description_
"""
return np.unravel_index(np.argmax(mask, axis=None), mask.shape)
class ImageTool:
def __init__(self, cfg_path="cfg/cfg.yml"):
# getting the config object
cfg_file = open(cfg_path)
self.cfg_obj = yaml.load(cfg_file, Loader=yaml.FullLoader)
# initialising the model and getting the annotations
self.ea_obj = ErrorAnalysis(cfg_path)
self.inference_folder = self.ea_obj.inference_folder
self.ea_obj.get_annots()
self.gt_annots = self.ea_obj.gt_dict
self.all_img = os.listdir(self.inference_folder)
self.ea_obj.model.score_threshold = self.cfg_obj["visual_tool"]["conf_threshold"]
self.ea_obj.model.iou_threshold = self.cfg_obj["visual_tool"]["iou_threshold"]
# for labels
self.labels_dict = self.cfg_obj["error_analysis"]["labels_dict"]
self.labels_dict = {v: k for k, v in self.labels_dict.items()}
self.inference_labels_dict = self.cfg_obj["error_analysis"]["inference_labels_dict"]
self.inference_labels_dict = {v: k for k, v in self.inference_labels_dict.items()}
self.idx_base = self.cfg_obj["error_analysis"]["idx_base"]
# for visualisation
self.bbox_thickness = self.cfg_obj["visual_tool"]["bbox_thickness"]
self.font_scale = self.cfg_obj["visual_tool"]["font_scale"]
self.font_thickness = self.cfg_obj["visual_tool"]["font_thickness"]
self.pred_colour = tuple(self.cfg_obj["visual_tool"]["pred_colour"])
self.gt_colour = tuple(self.cfg_obj["visual_tool"]["gt_colour"])
def show_img(self, img_fname="000000011149.jpg", show_preds=False, show_gt=False):
"""generate img with option to overlay with GT and/or preds
Args:
img_fname (str, optional): Filename of the image. Defaults to "000000011149.jpg".
show_preds (bool, optional): Toggle True to run model to get the preds. Defaults to False.
show_gt (bool, optional): Toggle True to get the GT labels/boxes. Defaults to False.
Returns:
fig (Plotly Figure): image with overlays if toggled True
cm_df (pd.DataFrame): confusion matrix of the pred versus GT
cm_tpfpfn_dict (Dict): confusion matrix dictionary of tp/fp/fn
"""
# get the image's file path. Concatenates with the folder in question
img = cv2.imread(f"{self.inference_folder}{img_fname}")
labels = {"x": "X", "y": "Y", "color": "Colour"}
if show_preds:
preds = self.get_preds(img_fname)
if self.ea_obj.task == "det":
img = self.draw_pred_bboxes(img, preds)
elif self.ea_obj.task == "seg":
img = self.draw_pred_masks(img, preds)
if show_gt:
gt_annots = self.get_gt_annot(img_fname)
if self.ea_obj.task == "det":
img = self.draw_gt_bboxes(img, preds)
elif self.ea_obj.task == "seg":
img = self.draw_gt_masks(img, gt_annots)
fig = px.imshow(img[..., ::-1], aspect="equal", labels=labels)
if show_gt and show_preds:
cm_df, cm_tpfpfn_dict = self.generate_cm_one_image(preds, gt_annots)
return [fig, cm_df, cm_tpfpfn_dict]
return fig
def show_img_sbs(self, img_fname="000000011149.jpg"):
"""generate two imageso with confusion matrix and tp/fp/fn. fig1 is image with GT overlay, while fig2 is the image witih pred overlay.
Args:
img_fname (str, optional): Filename of the image. Defaults to "000000011149.jpg".
Returns:
list: fig1 - imshow of image with GT overlay
fig2 - imshow of image with pred overlay
cm_df - confusion matrix dataframe
cm_tpfpfn_df - confusion matrix dictionary of tp/fp/fn
"""
# shows the image side by side
img = cv2.imread(f"{self.inference_folder}{img_fname}")
labels = {"x": "X", "y": "Y", "color": "Colour"}
img_pred = img.copy()
img_gt = img.copy()
preds = self.get_preds(img_fname)
gt_annots = self.get_gt_annot(img_fname)
if self.ea_obj.task == 'det':
img_pred = self.draw_pred_bboxes(img_pred, preds)
img_gt = self.draw_gt_bboxes(img_gt, gt_annots)
elif self.ea_obj.task == 'seg':
img_pred = self.draw_pred_masks(img_pred, preds)
img_gt = self.draw_gt_masks(img_gt, gt_annots)
fig1 = px.imshow(img_gt[..., ::-1], aspect="equal", labels=labels)
fig2 = px.imshow(img_pred[..., ::-1], aspect="equal", labels=labels)
fig2.update_yaxes(visible=False)
cm_df, cm_tpfpfn_df = self.generate_cm_one_image(preds, gt_annots)
return [fig1, fig2, cm_df, cm_tpfpfn_df]
def generate_cm_one_image(self, preds, gt_annots):
"""Generates confusion matrix between the inference and the Ground Truth of an image
Args:
preds (array): inference output of the model on the image
gt_annots (array): Ground Truth labels of the image
Returns:
cm_df (DataFrame): Confusion matrix dataframe.
cm_tpfpfn_df (DataFrame): TP/FP/FN dataframe
"""
num_classes = len(list(self.cfg_obj["error_analysis"]["labels_dict"].keys()))
idx_base = self.cfg_obj["error_analysis"]["idx_base"]
conf_threshold, iou_threshold = (
self.ea_obj.model.score_threshold,
self.ea_obj.model.iou_threshold,
)
cm = ConfusionMatrix(
num_classes=num_classes,
CONF_THRESHOLD=conf_threshold,
IOU_THRESHOLD=iou_threshold,
)
if self.ea_obj.task == 'det':
gt_annots[:, 0] -= idx_base
preds[:, -1] -= idx_base
elif self.ea_obj.task == 'seg':
gt_annots = [[gt[0] - idx_base, gt[1]] for gt in gt_annots]
cm.process_batch(preds, gt_annots, task = self.ea_obj.task)
confusion_matrix_df = cm.return_as_df()
cm.get_tpfpfn()
cm_tpfpfn_dict = {
"True Positive": cm.tp,
"False Positive": cm.fp,
"False Negative": cm.fn,
}
cm_tpfpfn_df = pd.DataFrame(cm_tpfpfn_dict, index=[0])
cm_tpfpfn_df = cm_tpfpfn_df.set_axis(["Values"], axis=0)
cm_tpfpfn_df = cm_tpfpfn_df.astype(int)
# amend df
confusion_matrix_df = amend_cm_df(confusion_matrix_df, self.labels_dict)
# print (cm.matrix)
return confusion_matrix_df, cm_tpfpfn_df
def get_preds(self, img_fname="000000011149.jpg"):
"""Using the model in the Error Analysis object, run inference to get outputs
Args:
img_fname (str): Image filename. Defaults to "000000011149.jpg".
Returns:
outputs (array): Inference output of the model on the image
"""
# run inference using the error analysis object per image
outputs, img_shape = self.ea_obj.generate_inference(img_fname)
if self.ea_obj.task == 'det':
# converts image coordinates from normalised to integer values
# image shape is [Y, X, C] (because Rows are Y)
# So don't get confused!
outputs[:, 0] *= img_shape[1]
outputs[:, 1] *= img_shape[0]
outputs[:, 2] *= img_shape[1]
outputs[:, 3] *= img_shape[0]
return outputs
def get_gt_annot(self, img_fname):
"""Retrieve the Ground Truth annotations of the image.
Args:
img_fname (_type_): Image filename
Returns:
grount_truth (array): GT labels of the image
"""
ground_truth = self.gt_annots[img_fname].copy()
img = cv2.imread(f"{self.inference_folder}{img_fname}")
# converts image coordinates from normalised to integer values
# image shape is [Y, X, C] (because Rows are Y)
# So don't get confused!
if self.ea_obj.task == 'det':
img_shape = img.shape
ground_truth = transform_gt_bbox_format(ground_truth, img_shape, format="coco")
ground_truth[:, 1] *= img_shape[1]
ground_truth[:, 2] *= img_shape[0]
ground_truth[:, 3] *= img_shape[1]
ground_truth[:, 4] *= img_shape[0]
return ground_truth
def draw_pred_masks(self, img_pred, inference_outputs):
"""Overlay mask onto img_pred
Args:
img_pred (_type_): _description_
preds (_type_): _description_
"""
pred_mask = sum([output[0] for output in inference_outputs])
pred_mask = np.where(pred_mask > 1, 1, pred_mask)
# mask_3d = np.stack((mask,mask,mask),axis=0)
# mask_3d = mask_3d.reshape(mask.shape[0], mask.shape[1], 3)
colour = np.array(self.pred_colour, dtype='uint8')
masked_img = np.where(pred_mask[...,None], colour, img_pred)
masked_img = masked_img.astype(np.uint8)
img_pred = cv2.addWeighted(img_pred, 0.7, masked_img, 0.3, 0)
def put_text_ina_mask(output, img):
coords = find_top_left_pos(output[0])
img = cv2.putText(img, self.inference_labels_dict[output[2]], (coords[1], coords[0] + 5), fontFace = cv2.FONT_HERSHEY_SIMPLEX, fontScale = self.font_scale,
color = self.pred_colour, thickness = self.font_thickness)
return img
for output in inference_outputs:
img_pred = put_text_ina_mask(output, img_pred)
return img_pred
def draw_gt_masks(self, img_gt, gt_outputs):
"""Overlay mask onto img_pred
Args:
img_pred (_type_): _description_
preds (_type_): _description_
"""
gt_mask = sum([output[1] for output in gt_outputs])
gt_mask = np.where(gt_mask > 1, 1, gt_mask)
# mask_3d = np.stack((mask,mask,mask),axis=0)
# mask_3d = mask_3d.reshape(mask.shape[0], mask.shape[1], 3)
colour = np.array(self.gt_colour, dtype='uint8')
masked_img = np.where(gt_mask[...,None], colour, img_gt)
def put_text_ina_mask(output, img):
coords = find_top_left_pos(output[1])
img = cv2.putText(img, self.labels_dict[output[0]], (coords[1], coords[0] + 5), fontFace = cv2.FONT_HERSHEY_SIMPLEX, fontScale = self.font_scale,
color = self.gt_colour, thickness = self.font_thickness)
return img
img_gt = cv2.addWeighted(img_gt, 0.7, masked_img, 0.3,0)
for output in gt_outputs:
img_gt = put_text_ina_mask(output, img_gt)
return img_gt
def draw_pred_bboxes(self, img_pred, preds):
"""Draws the preds onto the image
Args:
img_pred (array): image
preds (array): model inference outputs
Returns:
img_pred (array): image with outputs on overlay
"""
for pred in preds:
pred = pred.astype(int)
img_pred = cv2.rectangle(
img_pred,
(pred[0], pred[1]),
(pred[2], pred[3]),
color=self.pred_colour,
thickness=self.bbox_thickness,
)
img_pred = cv2.putText(
img_pred,
self.labels_dict[pred[5]],
(pred[0] + 5, pred[1] + 25),
color=self.pred_colour,
fontFace=cv2.FONT_HERSHEY_SIMPLEX,
fontScale=self.font_scale,
thickness=self.font_thickness,
)
return img_pred
def draw_gt_bboxes(self, img_gt, gt_annots, **kwargs):
"""Draws the GT onto the image
Args:
img_gt (array): image
gt_annots (array): GT labels
Returns:
img_gt (array): image with GT overlay
"""
for annot in gt_annots:
annot = annot.astype(int)
# print (annot)
img_gt = cv2.rectangle(
img_gt,
(annot[1], annot[2]),
(annot[3], annot[4]),
color=self.gt_colour,
thickness=self.bbox_thickness,
)
img_gt = cv2.putText(
img_gt,
self.labels_dict[annot[0]],
(annot[1] + 5, annot[2] + 25),
color=(0, 255, 0),
fontFace=cv2.FONT_HERSHEY_SIMPLEX,
fontScale=self.font_scale,
thickness=self.font_thickness,
)
return img_gt
def plot_with_preds_gt(self, option, side_by_side=False, plot_type=None):
"""Rules on what plot to generate
Args:
option (_string_): image filename. Toggled on the app itself. See app.py
side_by_side (bool, optional): Whether to have two plots side by side.
Defaults to False.
plot_type (_type_, optional): "all" - both GT and pred will be plotted,
"pred" - only preds,
"GT" - only ground truth
None - only image generated
Will be overridden if side_by_side = True
Defaults to None.
"""
if plot_type == "all":
plot, df, cm_tpfpfn_df = self.show_img(
option, show_preds=True, show_gt=True
)
st.plotly_chart(plot, use_container_width=True)
st.caption("Blue: Model BBox, Green: GT BBox")
st.table(df)
st.table(cm_tpfpfn_df)
elif plot_type == "pred":
st.plotly_chart(
self.show_img(option, show_preds=True), use_container_width=True
)
elif plot_type == "gt":
st.plotly_chart(
self.show_img(option, show_gt=True), use_container_width=True
)
elif side_by_side:
plot1, plot2, df, cm_tpfpfn_df = self.show_img_sbs(option)
col1, col2 = st.columns(2)
with col1:
col1.subheader("Ground Truth")
st.plotly_chart(plot1, use_container_width=True)
with col2:
col2.subheader("Prediction")
st.plotly_chart(plot2, use_container_width=True)
st.table(df)
st.table(cm_tpfpfn_df)
else:
st.plotly_chart(self.show_img(option), use_container_width=True)