61 lines
2.3 KiB
Python
61 lines
2.3 KiB
Python
# -*- coding: utf-8 -*-
|
|
import math
|
|
import numpy as np
|
|
from visdom import Visdom
|
|
from PIL import ImageDraw
|
|
|
|
|
|
class Logger():
|
|
"""Logger for training."""
|
|
|
|
def __init__(self, curve_names=None):
|
|
self.curve_names = curve_names
|
|
if curve_names:
|
|
self.vis = Visdom()
|
|
assert self.vis.check_connection()
|
|
self.curve_x = np.array([0])
|
|
|
|
def log(self, xval=None, win_name='loss', **kwargs):
|
|
"""Log and print the information."""
|
|
print("##############################################################")
|
|
for key, value in kwargs.items():
|
|
print(key, value, sep='\t')
|
|
|
|
if self.curve_names:
|
|
if not xval:
|
|
xval = self.curve_x
|
|
for i in range(len(self.curve_names)):
|
|
name = self.curve_names[i]
|
|
if name not in kwargs:
|
|
continue
|
|
yval = np.array([kwargs[name]])
|
|
self.vis.line(Y=yval, X=xval, win=win_name, update='append',
|
|
name=name, opts=dict(showlegend=True))
|
|
self.curve_x += 1
|
|
|
|
def plot_curve(self, yvals, xvals, win_name='pr_curves'):
|
|
"""Plot curve."""
|
|
self.vis.line(Y=np.array(yvals), X=np.array(xvals), win=win_name)
|
|
|
|
def plot_marking_points(self, image, marking_points, win_name='mk_points'):
|
|
"""Plot marking points on visdom."""
|
|
width, height = image.size
|
|
draw = ImageDraw.Draw(image)
|
|
for point in marking_points:
|
|
p0_x = width * point.x
|
|
p0_y = height * point.y
|
|
p1_x = p0_x + 50*math.cos(point.direction)
|
|
p1_y = p0_y + 50*math.sin(point.direction)
|
|
draw.line((p0_x, p0_y, p1_x, p1_y), fill=(255, 0, 0))
|
|
p2_x = p0_x - 50*math.sin(point.direction)
|
|
p2_y = p0_y + 50*math.cos(point.direction)
|
|
if point.shape > 0.5:
|
|
draw.line((p2_x, p2_y, p0_x, p0_y), fill=(255, 0, 0))
|
|
else:
|
|
p3_x = p0_x + 50*math.sin(point.direction)
|
|
p3_y = p0_y - 50*math.cos(point.direction)
|
|
draw.line((p2_x, p2_y, p3_x, p3_y), fill=(255, 0, 0))
|
|
image = np.asarray(image, dtype="uint8")
|
|
image = np.transpose(image, (2, 0, 1))
|
|
self.vis.image(image, win=win_name)
|