|
- """Class for logging."""
- import math
- import numpy as np
- from visdom import Visdom
- from PIL import ImageDraw
-
-
- class Logger():
- """Logger for training."""
-
- def __init__(self, enable_visdom=False, curve_names=None):
- self.curve_names = curve_names
- if enable_visdom:
- self.vis = Visdom()
- assert self.vis.check_connection()
- self.curve_x = np.array([0])
- else:
- self.curve_names = None
-
- 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)
|
