#!/usr/bin/env python3

__version__ = "1.0"

import logging
import mimetypes
import multiprocessing as mp
import os
import traceback
import warnings
import time
from contextlib import redirect_stderr, contextmanager
from multiprocessing import Pool, Semaphore

import click
import cv2
import numpy as np
from peewee import *
from more_itertools import peekable

#os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"

#with redirect_stderr(open(os.devnull, "w")):
from keras.models import load_model
from keras.backend import set_session
import keras

import tensorflow as tf

#tf.get_logger().setLevel("WARNING")
tf.config.run_functions_eagerly(False)
warnings.filterwarnings("ignore")


database = SqliteDatabase(None, autocommit=False)  # Defer initialization


class Result(Model):
    image_file = CharField(primary_key=True)
    columns = IntegerField()

    class Meta:
        database = database

semaphore = Semaphore(32)


@contextmanager
def timing(description: str, logger=None) -> None:
    start = time.time()
    yield
    elapsed_time = time.time() - start

    msg = f"{description}: {elapsed_time:.3f}s"
    if logger:
        logger.debug(msg)
    else:
        print(msg)


def _imread_and_prepare(image_file: str, model_input_shape):
    """Read a single image from a file and prepare it for page prediction

    Must be defined at the top-level so it can be pickled for multiprocessing.
    """

    # XXX Butt-ugly to use a global here
    global semaphore
    with semaphore:
        img_in = cv2.imread(image_file)
        if not np.any(img_in):
            return None, None, image_file

        # img = self.otsu_copy(image)
        BLUR_TIMES = 1
        for _ in range(BLUR_TIMES):
            img = cv2.GaussianBlur(img_in, (5, 5), 0)

        # n_classes = model.layers[len(model.layers) - 1].output_shape[3]

        dim = model_input_shape[1:]
        img = img / float(255.0)
        img = _resize_image(img, dim[0], dim[1])

        assert img.shape == dim

        return img, img_in, image_file


# XXX HACK HACK HACK
def _imread_and_prepare_HACK(image_file: str):
    model_input_shape = (None, 448, 448, 3)  # XXX hardcoded shape
    return _imread_and_prepare(image_file, model_input_shape)


def _resize_image(img_in, input_height, input_width):
    """
    Resize image.

    Wrap cv2.resize so that the order of the height and width parameters match .shape's order.

    Must be defined at the top-level so it can be pickled for multiprocessing.
    """
    return cv2.resize(img_in, (input_width, input_height), interpolation=cv2.INTER_NEAREST)


class sbb_column_classifier:
    def __init__(self, dir_models):

        # Setup logging
        logging.basicConfig(
            level=logging.DEBUG,
            format="%(asctime)s %(name)-12s %(levelname)-8s %(message)s",
            datefmt="%Y-%m-%d %H:%M:%S",
        )
        self.logger = logging.getLogger("sbb_column_classifier")
        self.logger.setLevel(logging.DEBUG)
        logging.getLogger("peewee").setLevel(logging.INFO)

        self.kernel = np.ones((5, 5), np.uint8)

        config = tf.compat.v1.ConfigProto()
        config.gpu_options.allow_growth = True
        session = tf.compat.v1.Session(config=config)
        set_session(session)

        self.model_page_file = os.path.join(dir_models, "model_page_mixed_best.h5")
        self.model_page = self.our_load_model(self.model_page_file)
        self.model_classifier_file = os.path.join(dir_models, "model_scale_classifier.h5")
        self.model_classifier = self.our_load_model(self.model_classifier_file)

        session.graph.finalize()

        self.time_last_batch = time.time()

    def our_load_model(self, model_file):
        self.logger.debug("Loading model {}...".format(os.path.basename(model_file)))

        model = load_model(model_file, compile=False)
        self.logger.debug(f"  input_shape:  {model.layers[0].input_shape}")
        self.logger.debug(f"  output_shape: {model.layers[len(model.layers) - 1].output_shape}")

        self.logger.debug("Loading model done.")

        return model

    @tf.function
    def model_page_predict(self, X):
        return self.model_page(X, training=False)
    @tf.function
    def model_classifier_predict(self, X):
        return self.model_classifier(X, training=False)



    @staticmethod
    def _crop_image_inside_box(box, img_org_copy):
        # input: coordinates of corner + dimensions
        #        box (x, y, w, h) or (y, x, h, w)
        # returns corner coordinates
        image_box = img_org_copy[box[1] : box[1] + box[3], box[0] : box[0] + box[2]]
        return image_box, [box[1], box[1] + box[3], box[0], box[0] + box[2]]

    @staticmethod
    def find_bounding_rect_of_largest_blob(img):
        # Find the largest blob in the prediction - it is our page (and its countour is the border)
        contours, _ = cv2.findContours(img, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

        if len(contours) > 0:
            cnt_size = np.array([cv2.contourArea(contours[j]) for j in range(len(contours))])
            cnt = contours[np.argmax(cnt_size)]

            # bounding rectangle of the largest blob
            x, y, w, h = cv2.boundingRect(cnt)

            # TODO Put 30 into a constant
            # If the prediction is near the img borders, go fully to the border
            if x <= 30:
                w = w + x
                x = 0
            if (img.shape[1] - (x + w)) <= 30:
                w = w + (img.shape[1] - (x + w))

            if y <= 30:
                h = h + y
                y = 0
            if (img.shape[0] - (y + h)) <= 30:
                h = h + (img.shape[0] - (y + h))

            box = [x, y, w, h]
        else:
            # If no contour is found, assume the full image as the page
            box = [0, 0, img.shape[1], img.shape[0]]

        return box

    N_WORKERS = min(3, mp.cpu_count())
    BATCH_SIZE = 16

    def _crop_page_from_pred(self, pred, img_in):
        seg = np.argmax(pred, axis=2)
        seg_color = np.repeat(seg[:, :, np.newaxis], 3, axis=2)

        # XXX Here we seem to resize it (just to resize it smaller again later)
        prediction_true = _resize_image(seg_color, img_in.shape[0], img_in.shape[1])
        prediction_true = prediction_true.astype(np.uint8)

        # XXX Why make an RGB image above and then convert it back to grayscale again?
        img_page_prediction = prediction_true
        img_gray = cv2.cvtColor(img_page_prediction, cv2.COLOR_BGR2GRAY)
        _, thresh = cv2.threshold(img_gray, 0, 255, 0)
        thresh = cv2.dilate(thresh, self.kernel, iterations=3)

        box = self.find_bounding_rect_of_largest_blob(thresh)
        cropped_page, page_coord = self._crop_image_inside_box(box, img_in)

        return cropped_page


    def crop_pages_batchwise(self, image_files):
        """Crop pages, batch for batch"""
        batch = []
        # X = np.empty((self.batch_size, *dim))
        with Pool(self.N_WORKERS) as pool:
            prepared_images = peekable(pool.imap(_imread_and_prepare_HACK, image_files))
            for x, img_in, image_file in prepared_images:
                if not np.any(x):
                    self.logger.error(f"Error reading {image_file}")
                    continue

                batch.append([x, img_in, image_file])

                # We have either a full batch or the last batch (= peekable iterator is exhausted):
                if len(batch) >= self.BATCH_SIZE or not prepared_images:
                    X = np.stack((x for x, _, _ in batch), axis=0)
                    X = tf.convert_to_tensor(X, dtype=tf.float64)
                    pred_batch = self.model_page_predict(X)

                    # TODO This doesn't run parallelized
                    with timing("Cropping images", logger=self.logger):
                        cropped_pages = []
                        for label_p_pred, (img_in, image_file2) in zip(pred_batch, ((img_in, image_file2) for _, img_in, image_file2 in batch)):
                            cropped_page = self._crop_page_from_pred(label_p_pred, img_in)
                            cropped_pages.append((cropped_page, image_file2))
                        batch = []

                    yield cropped_pages

    def number_of_columns(self, image_files):
        for cropped_pages_batch in self.crop_pages_batchwise(image_files):
            batch = []
            # TODO This doesn't run parallelized
            with timing("Resizing down again for the 2nd nnet", logger=self.logger):
                for cropped_page, image_file in cropped_pages_batch:
                    # TODO ... just to resize it down again
                    cropped_page = cropped_page / 255.0
                    cropped_page = cv2.resize(cropped_page, (448, 448), interpolation=cv2.INTER_NEAREST)  # XXX hardcoded shape
                    batch.append((cropped_page, image_file))

            X = np.stack((x for x, _ in batch), axis=0)
            X = tf.convert_to_tensor(X, dtype=tf.float64)
            label_p_pred = self.model_classifier_predict(X)
            num_col_batch = np.argmax(label_p_pred, axis=1) + 1

            duration_this_batch = time.time() - self.time_last_batch
            self.time_last_batch = time.time()

            self.logger.debug(f"Batch of {len(batch)} images done ({duration_this_batch:.3f}s, {duration_this_batch/len(batch):.3f}s/image).")
            yield from zip(num_col_batch, (image_file for _, image_file in batch))


@click.command()
@click.option(
    "--model",
    "-m",
    help="Directory of models (page extractor and classifier)",
    required=True,
    type=click.Path(exists=True, file_okay=False),
)
@click.option("--db-out", help="Write output as SQLite3 db", type=click.Path(exists=False, dir_okay=False))
@click.argument("images", required=True, type=click.Path(exists=True, dir_okay=True), nargs=-1)
def main(model, db_out, images):
    """
    Determine the number of columns in the document image IMAGES.

    Input document images should be in RGB. If a directory is given as IMAGES,
    we will process any image in the directory and its subdirectories.
    """

    #if mp.get_start_method(allow_none=True) != 'spawn':
    #     mp.set_start_method('spawn', force=True)

    cl = sbb_column_classifier(model)

    def is_image(fn):
        guessed_mimetype = mimetypes.guess_type(fn)[0]
        return guessed_mimetype is not None and guessed_mimetype.startswith("image/")

    def already_done(fn):
        return db_out and Result.get_or_none(Result.image_file == fn)

    def process_walk(i, explicitly_given=False):
        if os.path.isdir(i):
            root = i
            # Using os.scandir() here for better performance over os.walk()
            with os.scandir(root) as it:
                for entry in it:
                    yield from process_walk(os.path.join(root, entry.name))
        elif os.path.isfile(i):
            if explicitly_given or is_image(i):
                if already_done(i):
                    cl.logger.debug("Skipping {!r}, it is already done.".format(i))
                else:
                    yield i

    def process_walk_outer(images):
        for i in images:
            yield from process_walk(i, explicitly_given=True)

    if db_out:
        database.init(db_out, pragmas={'journal_mode': 'wal'})
        database.create_tables([Result])

    if db_out:
        database.begin()

    for i, (number_of_columns, image_file) in enumerate(cl.number_of_columns(process_walk_outer(images))):
        if db_out:
            r = Result.create(image_file=image_file, columns=number_of_columns)
            if i % 4*32 == 0:
                database.commit()
                database.begin()
        else:
            print("{!r},{}".format(image_file, number_of_columns))

        global semaphore
        semaphore.release()
    if db_out:
        database.commit()

    # TODO
    # try:
    #    cl.run(image_file)
    # except Exception:
    #    print(traceback.format_exc())


if __name__ == "__main__":
    main()

#!/usr/bin/env python3

__version__ = "1.0"

import logging
import mimetypes
import multiprocessing as mp
import os
import traceback
import warnings
import time
from contextlib import redirect_stderr, contextmanager
from multiprocessing import Pool, Semaphore

import click
import cv2
import numpy as np
from peewee import *
from more_itertools import peekable

#os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"

#with redirect_stderr(open(os.devnull, "w")):
from keras.models import load_model
from keras.backend import set_session
import keras

import tensorflow as tf

#tf.get_logger().setLevel("WARNING")
tf.config.run_functions_eagerly(False)
warnings.filterwarnings("ignore")


database = SqliteDatabase(None, autocommit=False)  # Defer initialization


class Result(Model):
    image_file = CharField(primary_key=True)
    columns = IntegerField()

    class Meta:
        database = database

semaphore = Semaphore(32)


@contextmanager
def timing(description: str, logger=None) -> None:
    start = time.time()
    yield
    elapsed_time = time.time() - start

    msg = f"{description}: {elapsed_time:.3f}s"
    if logger:
        logger.debug(msg)
    else:
        print(msg)


def _imread_and_prepare(image_file: str, model_input_shape):
    """Read a single image from a file and prepare it for page prediction

    Must be defined at the top-level so it can be pickled for multiprocessing.
    """

    # XXX Butt-ugly to use a global here
    global semaphore
    with semaphore:
        img_in = cv2.imread(image_file)
        if not np.any(img_in):
            return None, None, image_file

        # img = self.otsu_copy(image)
        BLUR_TIMES = 1
        for _ in range(BLUR_TIMES):
            img = cv2.GaussianBlur(img_in, (5, 5), 0)

        # n_classes = model.layers[len(model.layers) - 1].output_shape[3]

        dim = model_input_shape[1:]
        img = img / float(255.0)
        img = _resize_image(img, dim[0], dim[1])

        assert img.shape == dim

        return img, img_in, image_file


# XXX HACK HACK HACK
def _imread_and_prepare_HACK(image_file: str):
    model_input_shape = (None, 448, 448, 3)  # XXX hardcoded shape
    return _imread_and_prepare(image_file, model_input_shape)


def _resize_image(img_in, input_height, input_width):
    """
    Resize image.

    Wrap cv2.resize so that the order of the height and width parameters match .shape's order.

    Must be defined at the top-level so it can be pickled for multiprocessing.
    """
    return cv2.resize(img_in, (input_width, input_height), interpolation=cv2.INTER_NEAREST)


class sbb_column_classifier:
    def __init__(self, dir_models):

        # Setup logging
        logging.basicConfig(
            level=logging.DEBUG,
            format="%(asctime)s %(name)-12s %(levelname)-8s %(message)s",
            datefmt="%Y-%m-%d %H:%M:%S",
        )
        self.logger = logging.getLogger("sbb_column_classifier")
        self.logger.setLevel(logging.DEBUG)
        logging.getLogger("peewee").setLevel(logging.INFO)

        self.kernel = np.ones((5, 5), np.uint8)

        config = tf.compat.v1.ConfigProto()
        config.gpu_options.allow_growth = True
        session = tf.compat.v1.Session(config=config)
        set_session(session)

        self.model_page_file = os.path.join(dir_models, "model_page_mixed_best.h5")
        self.model_page = self.our_load_model(self.model_page_file)
        self.model_classifier_file = os.path.join(dir_models, "model_scale_classifier.h5")
        self.model_classifier = self.our_load_model(self.model_classifier_file)

        session.graph.finalize()

        self.time_last_batch = time.time()

    def our_load_model(self, model_file):
        self.logger.debug("Loading model {}...".format(os.path.basename(model_file)))

        model = load_model(model_file, compile=False)
        self.logger.debug(f"  input_shape:  {model.layers[0].input_shape}")
        self.logger.debug(f"  output_shape: {model.layers[len(model.layers) - 1].output_shape}")

        self.logger.debug("Loading model done.")

        return model

    @tf.function
    def model_page_predict(self, X):
        return self.model_page(X, training=False)
    @tf.function
    def model_classifier_predict(self, X):
        return self.model_classifier(X, training=False)



    @staticmethod
    def _crop_image_inside_box(box, img_org_copy):
        # input: coordinates of corner + dimensions
        #        box (x, y, w, h) or (y, x, h, w)
        # returns corner coordinates
        image_box = img_org_copy[box[1] : box[1] + box[3], box[0] : box[0] + box[2]]
        return image_box, [box[1], box[1] + box[3], box[0], box[0] + box[2]]

    @staticmethod
    def find_bounding_rect_of_largest_blob(img):
        # Find the largest blob in the prediction - it is our page (and its countour is the border)
        contours, _ = cv2.findContours(img, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

        if len(contours) > 0:
            cnt_size = np.array([cv2.contourArea(contours[j]) for j in range(len(contours))])
            cnt = contours[np.argmax(cnt_size)]

            # bounding rectangle of the largest blob
            x, y, w, h = cv2.boundingRect(cnt)

            # TODO Put 30 into a constant
            # If the prediction is near the img borders, go fully to the border
            if x <= 30:
                w = w + x
                x = 0
            if (img.shape[1] - (x + w)) <= 30:
                w = w + (img.shape[1] - (x + w))

            if y <= 30:
                h = h + y
                y = 0
            if (img.shape[0] - (y + h)) <= 30:
                h = h + (img.shape[0] - (y + h))

            box = [x, y, w, h]
        else:
            # If no contour is found, assume the full image as the page
            box = [0, 0, img.shape[1], img.shape[0]]

        return box

    N_WORKERS = min(3, mp.cpu_count())
    BATCH_SIZE = 16

    def _crop_page_from_pred(self, pred, img_in):
        seg = np.argmax(pred, axis=2)
        seg_color = np.repeat(seg[:, :, np.newaxis], 3, axis=2)

        # XXX Here we seem to resize it (just to resize it smaller again later)
        prediction_true = _resize_image(seg_color, img_in.shape[0], img_in.shape[1])
        prediction_true = prediction_true.astype(np.uint8)

        # XXX Why make an RGB image above and then convert it back to grayscale again?
        img_page_prediction = prediction_true
        img_gray = cv2.cvtColor(img_page_prediction, cv2.COLOR_BGR2GRAY)
        _, thresh = cv2.threshold(img_gray, 0, 255, 0)
        thresh = cv2.dilate(thresh, self.kernel, iterations=3)

        box = self.find_bounding_rect_of_largest_blob(thresh)
        cropped_page, page_coord = self._crop_image_inside_box(box, img_in)

        return cropped_page


    def crop_pages_batchwise(self, image_files):
        """Crop pages, batch for batch"""
        batch = []
        # X = np.empty((self.batch_size, *dim))
        with Pool(self.N_WORKERS) as pool:
            prepared_images = peekable(pool.imap(_imread_and_prepare_HACK, image_files))
            for x, img_in, image_file in prepared_images:
                if not np.any(x):
                    self.logger.error(f"Error reading {image_file}")
                    continue

                batch.append([x, img_in, image_file])

                # We have either a full batch or the last batch (= peekable iterator is exhausted):
                if len(batch) >= self.BATCH_SIZE or not prepared_images:
                    X = np.stack((x for x, _, _ in batch), axis=0)
                    X = tf.convert_to_tensor(X, dtype=tf.float64)
                    pred_batch = self.model_page_predict(X)

                    # TODO This doesn't run parallelized
                    with timing("Cropping images", logger=self.logger):
                        cropped_pages = []
                        for label_p_pred, (img_in, image_file2) in zip(pred_batch, ((img_in, image_file2) for _, img_in, image_file2 in batch)):
                            cropped_page = self._crop_page_from_pred(label_p_pred, img_in)
                            cropped_pages.append((cropped_page, image_file2))
                        batch = []

                    yield cropped_pages

    def number_of_columns(self, image_files):
        for cropped_pages_batch in self.crop_pages_batchwise(image_files):
            batch = []
            # TODO This doesn't run parallelized
            with timing("Resizing down again for the 2nd nnet", logger=self.logger):
                for cropped_page, image_file in cropped_pages_batch:
                    # TODO ... just to resize it down again
                    cropped_page = cropped_page / 255.0
                    cropped_page = cv2.resize(cropped_page, (448, 448), interpolation=cv2.INTER_NEAREST)  # XXX hardcoded shape
                    batch.append((cropped_page, image_file))

            X = np.stack((x for x, _ in batch), axis=0)
            X = tf.convert_to_tensor(X, dtype=tf.float64)
            label_p_pred = self.model_classifier_predict(X)
            num_col_batch = np.argmax(label_p_pred, axis=1) + 1

            duration_this_batch = time.time() - self.time_last_batch
            self.time_last_batch = time.time()

            self.logger.debug(f"Batch of {len(batch)} images done ({duration_this_batch:.3f}s, {duration_this_batch/len(batch):.3f}s/image).")
            yield from zip(num_col_batch, (image_file for _, image_file in batch))


@click.command()
@click.option(
    "--model",
    "-m",
    help="Directory of models (page extractor and classifier)",
    required=True,
    type=click.Path(exists=True, file_okay=False),
)
@click.option("--db-out", help="Write output as SQLite3 db", type=click.Path(exists=False, dir_okay=False))
@click.argument("images", required=True, type=click.Path(exists=True, dir_okay=True), nargs=-1)
def main(model, db_out, images):
    """
    Determine the number of columns in the document image IMAGES.

    Input document images should be in RGB. If a directory is given as IMAGES,
    we will process any image in the directory and its subdirectories.
    """

    #if mp.get_start_method(allow_none=True) != 'spawn':
    #     mp.set_start_method('spawn', force=True)

    cl = sbb_column_classifier(model)

    def is_image(fn):
        guessed_mimetype = mimetypes.guess_type(fn)[0]
        return guessed_mimetype is not None and guessed_mimetype.startswith("image/")

    def already_done(fn):
        return db_out and Result.get_or_none(Result.image_file == fn)

    def process_walk(i, explicitly_given=False):
        if os.path.isdir(i):
            root = i
            # Using os.scandir() here for better performance over os.walk()
            with os.scandir(root) as it:
                for entry in it:
                    yield from process_walk(os.path.join(root, entry.name))
        elif os.path.isfile(i):
            if explicitly_given or is_image(i):
                if already_done(i):
                    cl.logger.debug("Skipping {!r}, it is already done.".format(i))
                else:
                    yield i

    def process_walk_outer(images):
        for i in images:
            yield from process_walk(i, explicitly_given=True)

    if db_out:
        database.init(db_out, pragmas={'journal_mode': 'wal'})
        database.create_tables([Result])

    if db_out:
        database.begin()

    for i, (number_of_columns, image_file) in enumerate(cl.number_of_columns(process_walk_outer(images))):
        if db_out:
            r = Result.create(image_file=image_file, columns=number_of_columns)
            if i % 4*32 == 0:
                database.commit()
                database.begin()
        else:
            print("{!r},{}".format(image_file, number_of_columns))

        global semaphore
        semaphore.release()
    if db_out:
        database.commit()

    # TODO
    # try:
    #    cl.run(image_file)
    # except Exception:
    #    print(traceback.format_exc())


if __name__ == "__main__":
    main()

#!/usr/bin/env python3 #os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2" #with redirect_stderr(open(os.devnull, "w")): #tf.get_logger().setLevel("WARNING") # Defer initialization Read a single image from a file and prepare it for page prediction      Must be defined at the top-level so it can be pickled for multiprocessing. # XXX Butt-ugly to use a global here # img = self.otsu_copy(image) # n_classes = model.layers[len(model.layers) - 1].output_shape[3] # XXX HACK HACK HACK # XXX hardcoded shape Resize image.      Wrap cv2.resize so that the order of the height and width parameters match .shape's order.      Must be defined at the top-level so it can be pickled for multiprocessing. # Setup logging # input: coordinates of corner + dimensions #        box (x, y, w, h) or (y, x, h, w) # returns corner coordinates # Find the largest blob in the prediction - it is our page (and its countour is the border) # bounding rectangle of the largest blob # TODO Put 30 into a constant # If the prediction is near the img borders, go fully to the border # If no contour is found, assume the full image as the page # XXX Here we seem to resize it (just to resize it smaller again later) # XXX Why make an RGB image above and then convert it back to grayscale again? Crop pages, batch for batch # X = np.empty((self.batch_size, *dim)) # We have either a full batch or the last batch (= peekable iterator is exhausted): # TODO This doesn't run parallelized # TODO This doesn't run parallelized # TODO ... just to resize it down again # XXX hardcoded shape Determine the number of columns in the document image IMAGES.      Input document images should be in RGB. If a directory is given as IMAGES,     we will process any image in the directory and its subdirectories. #if mp.get_start_method(allow_none=True) != 'spawn': #     mp.set_start_method('spawn', force=True) # Using os.scandir() here for better performance over os.walk() # TODO # try: #    cl.run(image_file) # except Exception: #    print(traceback.format_exc())