pbm layout

combine_pbm.py

import os.path, sys

memory_available = 4 * (10**9)
outfile = "out.pbm"
default_tsv_file="image_lookup_table.tsv"
DEBUG = False

def debug(*x):
    if DEBUG:
        print(*x)
def progress(*x):
    print(*x, file=sys.stderr)

class PBMImageRead():
    def __init__(self, path):
        self.path = path
        self.f = open(path, "rb")
        self._read_header()
    def _read_header(self):
        assert self.f.read(3) == b'P4\n'
        wh = self.f.readline().decode(encoding='ascii')
        width, height = wh.split()
        self.width, self.height = int(width), int(height)
    def read_bits(self, bits):
        assert bits > 0
        assert bits % 8 == 0
        r = self.f.read(bits//8)
        assert len(r) == bits//8, "{} != {}".format(len(r), bits//8)
        return r
    def close(self):
        self.f.close()

class PBMImageWrite():
    def __init__(self, path, width, height):
        self.path = path
        self.f = open(path, "wb")
        self.width = width
        self.height = height
        self._write_header()
    def _write_header(self):
        self.f.write(b'P4\n')
        self.f.write(bytes(str(self.width), encoding='ascii'))
        self.f.write(b' ')
        self.f.write(bytes(str(self.height), encoding='ascii'))
        self.f.write(b'\n')
    def write_bits(self, bits, data):
        assert bits % 8 == 0
        assert isinstance(data, bytes) or isinstance(data, bytearray), data # This could later become a bitarray to support not-divisible-by-8 rectangle widths or positions
        assert len(data) == (bits//8), "{} != {} ({})".format(len(data), bits//8, bits)
        return self.f.write(data)
    def close(self):
        self.f.close()

class Rectangle():
    def __init__(self, x1, x2, y1, y2, path, line):
        assert x1 < x2
        assert x1 % 8 == 0, "Rectangles must be byte-aligned (widths and positions divisible by 8)"
        assert (x2+1) % 8 == 0
        self.x1 = x1
        self.x2 = x2
        assert y1 < y2
        self.y1 = y1
        self.y2 = y2
        self.width = x2 - x1 + 1
        assert self.width % 8 == 0
        self.height = y2 - y1 + 1
        self.path = path
        self.line = line
    def __hash__(self):
        return hash((self.x1, self.x2, self.y1, self.y2, self.path))
    @staticmethod
    def has_overlap(r1, r2):
        if r1.x2 < r2.x1: # R1 is left of R2
            return False
        if r1.x1 > r2.x2: # R1 is right of R2
            return False
        if r1.y2 < r2.y1: # R1 is above R2
            return False
        if r1.y1 > r2.y2: # R1 is below R2
            return False
        return True
    def overlaps(self, r2):
        return Rectangle.has_overlap(self, r2)
    def open(self):
        self.pbm = PBMImageRead(self.path)
        assert self.pbm.width == (self.x2-self.x1+1) == self.width
        assert self.pbm.height == (self.y2-self.y1+1) == self.height
    def close(self):
        self.pbm.close()
        self.pbm = None
    def __repr__(self):
        return "<Rectangle(x1={} x2={} y1={} y2={} line={} path={})>".format(self.x1, self.x2, self.y1, self.y2, self.line, self.path)

def load_rectangles(path=default_tsv_file, check_files=True):
    rectangles = [] # TSV of x coordinates, y coordinates, path to monochrome pbm
    progress("Reading layout TSV: {}".format(path))
    with open(path, "r") as f:
        for i, line in enumerate(f):
            if (i==0):
                continue
            x1, x2, y1, y2, r_path = line[:-1].split("\t")
            x1, x2, y1, y2 = int(x1), int(x2), int(y1), int(y2)
            rect = Rectangle(x1, x2, y1, y2, r_path, i+1)
            rectangles.append(rect)
    assert len(rectangles) == len(set(rectangles)), "Rectangles are not all unique"

    progress("Checking rectangles for overlap...")
    # Brute force, this doesn't scale well. Comment it out if it becomes an issue
    for r1 in rectangles:
        for r2 in rectangles:
            if r1==r2:
                continue
            assert not r1.overlaps(r2), "Two rectangles overlap: {} and {}".format(r1, r2)

    if check_files:
        progress("Making sure rectangle .pbm files exist...")
        for r in rectangles:
            assert os.path.exists(r.path), "Path does not exist: {}".format(r.path)
    progress("Rectangles loaded.")
    return rectangles

rectangles = load_rectangles()
minY, maxY = min(r.y1 for r in rectangles), max(r.y2 for r in rectangles)
minX, maxX = min(r.x1 for r in rectangles), max(r.x2 for r in rectangles)
assert minX == 0 # actually needed
assert minY == 0 # should work anyway
width, height = (maxX-minX+1), (maxY-minY+1)
assert width % 8 == 0
row_batch_size = min(memory_available*8//(width*2), height)
debug("minX", minX)
debug("maxX", maxX)
debug("minY", minY)
debug("maxY", maxY)
debug("width", width)
debug("height", height)
debug("row_batch_size", row_batch_size)

progress("Outputting rectangles...")
output = PBMImageWrite(outfile, width, height)
for row in range(0, minY):
    output.write_bits(width, bytes(width//8))
for start_row in range(minY, maxY+1, row_batch_size):
    end_row = min(maxY, start_row + row_batch_size) # not inclusive
    moving_window_size = (end_row-start_row)*width
    assert moving_window_size > 0
    assert moving_window_size % 8 == 0
    debug("start_row", start_row)
    debug("end_row", end_row)
    debug("moving_window_size", moving_window_size)
    moving_window = bytearray(moving_window_size//8)
    moving_window_rect = Rectangle(minX, maxX, start_row, end_row-1, None, None)
    overlapping_rectangles = [r for r in rectangles if moving_window_rect.overlaps(r)]
    for r in overlapping_rectangles:
        assert r.width % 8 == 0
        if start_row <= r.y1:
            r.open()
        assert r.y1 <= end_row-1, "{} !<= {}".format(r.y1, end_row-1)
        start_rect_row = max(start_row, r.y1)
        assert start_row <= r.y2, "{} !<= {}".format(start_row, r.y2)
        end_rect_row = min(end_row-1, r.y2) # inclusive
        assert minY <= start_row <= start_rect_row <= end_rect_row <= end_row - 1 <= maxY
        overlap_height = end_rect_row - start_rect_row + 1
        r_overlap = r.pbm.read_bits(r.width*overlap_height)
        for y in range(start_rect_row, end_rect_row+1):
            assert len(moving_window) == moving_window_size//8, (start_row, r, y)
            r_offset = y-start_rect_row
            w_offset = y-start_row
            assert (w_offset*width + r.x1 + r.width)//8 <= len(moving_window), "{} {} {} {} {}".format(y, width, r.x1, r.width, len(moving_window))
            assert (r_offset+1)*r.width//8 <= len(r_overlap)
            
            moving_window[(w_offset*width + r.x1)//8 : (w_offset*width + r.x1 + r.width)//8] = r_overlap[r_offset*r.width//8:(r_offset+1)*r.width//8]
            try:
                assert len(moving_window) == moving_window_size//8, (start_row, r, y)
            except AssertionError:
                debug("r", r)
                raise
        if r.y2 <= end_row:
            r.close()
    output.write_bits(moving_window_size, moving_window)
    progress("{:.1f}% done...".format(end_row/maxY*100))
progress("Layout done. Available as {}".format(outfile))