license-plate-detect/ultralytics/solutions/analytics.py

313 lines
11 KiB
Python
Raw Normal View History

2024-08-26 20:19:30 +08:00
# Ultralytics YOLO 🚀, AGPL-3.0 license
import warnings
from itertools import cycle
import cv2
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
from matplotlib.figure import Figure
class Analytics:
"""A class to create and update various types of charts (line, bar, pie, area) for visual analytics."""
def __init__(
self,
type,
writer,
im0_shape,
title="ultralytics",
x_label="x",
y_label="y",
bg_color="white",
fg_color="black",
line_color="yellow",
line_width=2,
points_width=10,
fontsize=13,
view_img=False,
save_img=True,
max_points=50,
):
"""
Initialize the Analytics class with various chart types.
Args:
type (str): Type of chart to initialize ('line', 'bar', 'pie', or 'area').
writer (object): Video writer object to save the frames.
im0_shape (tuple): Shape of the input image (width, height).
title (str): Title of the chart.
x_label (str): Label for the x-axis.
y_label (str): Label for the y-axis.
bg_color (str): Background color of the chart.
fg_color (str): Foreground (text) color of the chart.
line_color (str): Line color for line charts.
line_width (int): Width of the lines in line charts.
points_width (int): Width of line points highlighter
fontsize (int): Font size for chart text.
view_img (bool): Whether to display the image.
save_img (bool): Whether to save the image.
max_points (int): Specifies when to remove the oldest points in a graph for multiple lines.
"""
self.bg_color = bg_color
self.fg_color = fg_color
self.view_img = view_img
self.save_img = save_img
self.title = title
self.writer = writer
self.max_points = max_points
self.line_color = line_color
self.x_label = x_label
self.y_label = y_label
self.points_width = points_width
self.line_width = line_width
self.fontsize = fontsize
# Set figure size based on image shape
figsize = (im0_shape[0] / 100, im0_shape[1] / 100)
if type in {"line", "area"}:
# Initialize line or area plot
self.lines = {}
self.fig = Figure(facecolor=self.bg_color, figsize=figsize)
self.canvas = FigureCanvas(self.fig)
self.ax = self.fig.add_subplot(111, facecolor=self.bg_color)
if type == "line":
(self.line,) = self.ax.plot([], [], color=self.line_color, linewidth=self.line_width)
elif type in {"bar", "pie"}:
# Initialize bar or pie plot
self.fig, self.ax = plt.subplots(figsize=figsize, facecolor=self.bg_color)
self.ax.set_facecolor(self.bg_color)
color_palette = [
(31, 119, 180),
(255, 127, 14),
(44, 160, 44),
(214, 39, 40),
(148, 103, 189),
(140, 86, 75),
(227, 119, 194),
(127, 127, 127),
(188, 189, 34),
(23, 190, 207),
]
self.color_palette = [(r / 255, g / 255, b / 255, 1) for r, g, b in color_palette]
self.color_cycle = cycle(self.color_palette)
self.color_mapping = {}
# Ensure pie chart is circular
self.ax.axis("equal") if type == "pie" else None
# Set common axis properties
self.ax.set_title(self.title, color=self.fg_color, fontsize=self.fontsize)
self.ax.set_xlabel(x_label, color=self.fg_color, fontsize=self.fontsize - 3)
self.ax.set_ylabel(y_label, color=self.fg_color, fontsize=self.fontsize - 3)
self.ax.tick_params(axis="both", colors=self.fg_color)
def update_area(self, frame_number, counts_dict):
"""
Update the area graph with new data for multiple classes.
Args:
frame_number (int): The current frame number.
counts_dict (dict): Dictionary with class names as keys and counts as values.
"""
x_data = np.array([])
y_data_dict = {key: np.array([]) for key in counts_dict.keys()}
if self.ax.lines:
x_data = self.ax.lines[0].get_xdata()
for line, key in zip(self.ax.lines, counts_dict.keys()):
y_data_dict[key] = line.get_ydata()
x_data = np.append(x_data, float(frame_number))
max_length = len(x_data)
for key in counts_dict.keys():
y_data_dict[key] = np.append(y_data_dict[key], float(counts_dict[key]))
if len(y_data_dict[key]) < max_length:
y_data_dict[key] = np.pad(y_data_dict[key], (0, max_length - len(y_data_dict[key])), "constant")
# Remove the oldest points if the number of points exceeds max_points
if len(x_data) > self.max_points:
x_data = x_data[1:]
for key in counts_dict.keys():
y_data_dict[key] = y_data_dict[key][1:]
self.ax.clear()
colors = ["#E1FF25", "#0BDBEB", "#FF64DA", "#111F68", "#042AFF"]
color_cycle = cycle(colors)
for key, y_data in y_data_dict.items():
color = next(color_cycle)
self.ax.fill_between(x_data, y_data, color=color, alpha=0.6)
self.ax.plot(
x_data,
y_data,
color=color,
linewidth=self.line_width,
marker="o",
markersize=self.points_width,
label=f"{key} Data Points",
)
self.ax.set_title(self.title, color=self.fg_color, fontsize=self.fontsize)
self.ax.set_xlabel(self.x_label, color=self.fg_color, fontsize=self.fontsize - 3)
self.ax.set_ylabel(self.y_label, color=self.fg_color, fontsize=self.fontsize - 3)
legend = self.ax.legend(loc="upper left", fontsize=13, facecolor=self.bg_color, edgecolor=self.fg_color)
# Set legend text color
for text in legend.get_texts():
text.set_color(self.fg_color)
self.canvas.draw()
im0 = np.array(self.canvas.renderer.buffer_rgba())
self.write_and_display(im0)
def update_line(self, frame_number, total_counts):
"""
Update the line graph with new data.
Args:
frame_number (int): The current frame number.
total_counts (int): The total counts to plot.
"""
# Update line graph data
x_data = self.line.get_xdata()
y_data = self.line.get_ydata()
x_data = np.append(x_data, float(frame_number))
y_data = np.append(y_data, float(total_counts))
self.line.set_data(x_data, y_data)
self.ax.relim()
self.ax.autoscale_view()
self.canvas.draw()
im0 = np.array(self.canvas.renderer.buffer_rgba())
self.write_and_display(im0)
def update_multiple_lines(self, counts_dict, labels_list, frame_number):
"""
Update the line graph with multiple classes.
Args:
counts_dict (int): Dictionary include each class counts.
labels_list (int): list include each classes names.
frame_number (int): The current frame number.
"""
warnings.warn("Display is not supported for multiple lines, output will be stored normally!")
for obj in labels_list:
if obj not in self.lines:
(line,) = self.ax.plot([], [], label=obj, marker="o", markersize=self.points_width)
self.lines[obj] = line
x_data = self.lines[obj].get_xdata()
y_data = self.lines[obj].get_ydata()
# Remove the initial point if the number of points exceeds max_points
if len(x_data) >= self.max_points:
x_data = np.delete(x_data, 0)
y_data = np.delete(y_data, 0)
x_data = np.append(x_data, float(frame_number)) # Ensure frame_number is converted to float
y_data = np.append(y_data, float(counts_dict.get(obj, 0))) # Ensure total_count is converted to float
self.lines[obj].set_data(x_data, y_data)
self.ax.relim()
self.ax.autoscale_view()
self.ax.legend()
self.canvas.draw()
im0 = np.array(self.canvas.renderer.buffer_rgba())
self.view_img = False # for multiple line view_img not supported yet, coming soon!
self.write_and_display(im0)
def write_and_display(self, im0):
"""
Write and display the line graph
Args:
im0 (ndarray): Image for processing
"""
im0 = cv2.cvtColor(im0[:, :, :3], cv2.COLOR_RGBA2BGR)
cv2.imshow(self.title, im0) if self.view_img else None
self.writer.write(im0) if self.save_img else None
def update_bar(self, count_dict):
"""
Update the bar graph with new data.
Args:
count_dict (dict): Dictionary containing the count data to plot.
"""
# Update bar graph data
self.ax.clear()
self.ax.set_facecolor(self.bg_color)
labels = list(count_dict.keys())
counts = list(count_dict.values())
# Map labels to colors
for label in labels:
if label not in self.color_mapping:
self.color_mapping[label] = next(self.color_cycle)
colors = [self.color_mapping[label] for label in labels]
bars = self.ax.bar(labels, counts, color=colors)
for bar, count in zip(bars, counts):
self.ax.text(
bar.get_x() + bar.get_width() / 2,
bar.get_height(),
str(count),
ha="center",
va="bottom",
color=self.fg_color,
)
# Display and save the updated graph
canvas = FigureCanvas(self.fig)
canvas.draw()
buf = canvas.buffer_rgba()
im0 = np.asarray(buf)
self.write_and_display(im0)
def update_pie(self, classes_dict):
"""
Update the pie chart with new data.
Args:
classes_dict (dict): Dictionary containing the class data to plot.
"""
# Update pie chart data
labels = list(classes_dict.keys())
sizes = list(classes_dict.values())
total = sum(sizes)
percentages = [size / total * 100 for size in sizes]
start_angle = 90
self.ax.clear()
# Create pie chart without labels inside the slices
wedges, autotexts = self.ax.pie(sizes, autopct=None, startangle=start_angle, textprops={"color": self.fg_color})
# Construct legend labels with percentages
legend_labels = [f"{label} ({percentage:.1f}%)" for label, percentage in zip(labels, percentages)]
self.ax.legend(wedges, legend_labels, title="Classes", loc="center left", bbox_to_anchor=(1, 0, 0.5, 1))
# Adjust layout to fit the legend
self.fig.tight_layout()
self.fig.subplots_adjust(left=0.1, right=0.75)
# Display and save the updated chart
im0 = self.fig.canvas.draw()
im0 = np.array(self.fig.canvas.renderer.buffer_rgba())
self.write_and_display(im0)
if __name__ == "__main__":
Analytics("line", writer=None, im0_shape=None)