ddrescue_visualisation/ddresc_visu/draw.py

from .mapfile import MapFile

import pygame

colors = {
	'?' : (192, 192, 192), # Non-tried
	'*' : (0  , 0  , 255), # Non-trimmed
	'/' : (255, 255, 0)  , # Non-scraped
	'-' : (255, 0  , 0)  , # Bad
	'+' : (0  , 128, 0)  , # Recovered
	None: (0  , 0  , 0)  , # Not on disk
}

BORDERCOLOR = (64,64,64)

W = 1920 * 2 // 3
H = 1060
CELLSIZE = 8
CELLBORDER = 0
CELLMARGIN = 3

CELLOFFSET = CELLSIZE + 2*CELLBORDER + CELLMARGIN
COLUMNS = W // CELLOFFSET
ROWS = H // CELLOFFSET
CELLS = COLUMNS * ROWS

def get_symb(hist):
	## Majority:
	#symb, _count = max(hist.items(), key=lambda x: x[1])
	#return symb
	
	# The worst:
	order = ['-', '/', '*', '?', '+']
	for ch in order:
		if ch in hist:
			return ch

class Visualisation:
	def __init__(self, fn, **kwargs):
		pygame.init()
		self.disp = pygame.display.set_mode((W, H))

		self.mapfile = MapFile(fn, **kwargs)
		self.mapfile.load()

	def draw(self):
		self.disp.fill((0,0,0)) # Dark theme :-)

		sz = self.mapfile.size
		sqsz = sz / CELLS
		for cid in range(CELLS):
			start = cid * sqsz
			end = (cid + 1) * sqsz
			hist = self.mapfile.get_hist(start, end)
			symb = get_symb(hist)
			color = colors[symb]

			line = cid // COLUMNS
			col = cid % COLUMNS
			pos = (line, col)

			rect = pygame.Rect(col * CELLOFFSET + CELLBORDER, line * CELLOFFSET + CELLBORDER, CELLSIZE+2*CELLBORDER, CELLSIZE+2*CELLBORDER)
			pygame.draw.rect(self.disp, color, rect)
			if CELLBORDER != 0: pygame.draw.rect(self.disp, BORDERCOLOR, rect, width=CELLBORDER)

		pygame.display.flip()

	def reload(self):
		self.mapfile.load()
		self.draw()