Created
October 28, 2019 20:28
-
-
Save JACKHAHA363/56746e6c8d163d3614840dda4addfec7 to your computer and use it in GitHub Desktop.
Plot probability with original image for demo
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import cv2 | |
from PIL import Image | |
import matplotlib.pyplot as plt | |
import numpy as np | |
def _rounding(ratio, max_val): | |
""" Return a rounded valud from 0 to max_val - 1 """ | |
val = int(ratio * max_val) | |
return min(max_val - 1, val) | |
def _matplotfig2data(fig): | |
""" matplotlib fig to data. Return a cv2 image """ | |
# draw the renderer | |
fig.canvas.draw() | |
# Get the RGBA buffer from the figure | |
w, h = fig.canvas.get_width_height() | |
buf = np.frombuffer(fig.canvas.tostring_argb(), dtype=np.uint8) | |
buf.shape = (w, h, 4) | |
buf = np.roll(buf, 3, axis=2) | |
img_pil = Image.frombytes("RGBA", (w, h), buf.tostring()) | |
np_img = np.array(img_pil.convert('RGB')) | |
return np_img[:, :, ::-1] | |
def bbox_on_img(img, bbox, text, thickness=3, | |
text_size=16, color=(0, 0, 0)): | |
""" | |
:param img: A Pillow Image | |
:param bbox: A tuple of left upper corner x, y and right down corner x, y | |
from 0 | |
:param text: A str | |
:param color: Bounding box/text color | |
:return A rendered img | |
""" | |
h, w = img.shape[:2] | |
# Draw bbox | |
ul_x = _rounding(bbox[0], w) | |
ul_y = _rounding(bbox[1], h) | |
br_x = _rounding(bbox[2], w) | |
br_y = _rounding(bbox[3], h) | |
img = cv2.rectangle(img, (ul_x, ul_y), (br_x, br_y), color, | |
thickness=thickness) | |
img = cv2.putText(img, text, | |
(ul_x, ul_y-5), | |
fontFace=cv2.FONT_HERSHEY_SIMPLEX, | |
fontScale=0.5, | |
color=(0, 0, 0), | |
thickness=thickness) | |
return img | |
def add_prob_bar(img, labels, probs): | |
""" add probablistic bar """ | |
height, width = img.shape[:2] | |
# Generate histogram | |
fig, axs = plt.subplots() | |
axs.bar(labels, probs) | |
axs.set_ylim([0, 1]) | |
axs.set_ylabel('probablity') | |
data_img = _matplotfig2data(fig) | |
data_img = cv2.resize(data_img, (int(width / 1.5), height)) | |
full_img = np.concatenate([img, data_img], axis=1) | |
return full_img | |
img = cv2.imread('sample.jpg') | |
img = add_prob_bar(img, ['dog', 'cat', 'butterfly'], [0.4, 0.5, 0.1]) | |
img_pil = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB)) | |
img_pil.show() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
sample.jpg
Result