import cv2
import numpy as np
import requests
from PIL import Image
# import pyocr
import pytesseract
import io
import re
import os

class image_process:

    def __init__(self):
        self.input_img = None
        self.result_img = None
        self.id_recog_result_img = None
        self.result_img_byte = None
        self.id_recog_result_img_byte = []
        self.leftscore = 0
        self.rightscore = 0
        self.message = ''
        self.player_id = ''
        self.clan_mems = []
        self.clanmem_id = []
        self.flag = False
        self.num_vs = 0
        self.match_result = ''
        self.m = ''

    def image_process(self, url):
        self.url = url
        resp = requests.get(self.url, stream=True).raw
        img = np.asarray(bytearray(resp.read()), dtype="uint8")
        img = cv2.imdecode(img, cv2.IMREAD_COLOR)
        self.input_img = cv2.resize(img, (1920, 1080))
        self.result_img = self.input_img.copy()
        self.judge_battle()
        self.player_recog()

    def image_process_dev(self, input_img_path):
        self.input_img = cv2.imread(input_img_path)
        self.input_img = cv2.resize(self.input_img, (1920, 1080))
        self.result_img = self.input_img.copy()
        self.judge_battle()
        self.player_recog()

    def judge_battle(self):
        img = cv2.cvtColor(self.input_img, cv2.COLOR_RGB2GRAY)

        ret, bin_target = cv2.threshold(img, 200, 255, cv2.THRESH_BINARY)

        # cv2.imwrite('bin_target.png', bin_target)

        label = []
        label_x_sum = 0

        for i in range(10):
            temp_name = f'./template_image_bin/large{i}.png'
            template_ori = cv2.imread(temp_name, 0)

            w, h = template_ori.shape[::-1]
            res = cv2.matchTemplate(bin_target, template_ori, cv2.TM_CCOEFF_NORMED)
            threshold = 0.9
            loc = np.where(res >= threshold)
            loc = list(zip(*loc[::-1]))
            loc.sort()
            # print(i, loc)
            cnt = 1

            if loc:
                x_pre, y_pre = x_sum, y_sum = loc[0][0], loc[0][1]
                cv2.rectangle(self.result_img, (x_pre, y_pre), (x_pre + w, y_pre + h), (0, 0, 255), 2)
                self.flag = True

            for x, y in loc[1:]:

                if abs(x_pre - x) > 20 or abs(y_pre - y) > 20:
                    cv2.rectangle(self.result_img, (x, y), (x + w, y + h), (0, 0, 255), 2)
                    label.append([x_sum / cnt, y_sum / cnt, i])
                    label_x_sum += x_sum / cnt
                    x_pre, y_pre = x_sum, y_sum = 0, 0
                    cnt = 0

                x_sum += x
                y_sum += y
                x_pre = x
                y_pre = y
                cnt += 1

                cv2.rectangle(self.result_img, (x, y), (x + w, y + h), (0, 0, 255), 2)

            if loc:
                cv2.rectangle(self.result_img, (x_pre, y_pre), (x_pre + w, y_pre + h), (0, 0, 255), 2)
                label.append([x_sum / cnt, y_sum / cnt, i])
                label_x_sum += x_sum / cnt

        label.sort()
        # print("\n", label)
        # cv2.imshow("self.result_img", self.result_img)
        # cv2.waitKey(0)
        # cv2.destroyAllWindows()
        try:
            label_x_ave = label_x_sum / len(label)

            left_score, right_score = '', ''

            for x, y, i in label:

                if x < label_x_ave:
                    left_score += str(i)

                else:
                    right_score += str(i)

            self.leftscore = int(left_score)
            self.rightscore = int(right_score)
            # print(f'{file} {left_score} vs {right_score}')

            # cv2.imwrite("self.result_img.png", self.result_img)
            # cv2.waitKey(0)
            # cv2.destroyAllWindows()
            if int(self.leftscore) > int( self.rightscore):
                judge = '勝ち～！'
                self.match_result = 'WIN'

            elif int(self.leftscore) < int( self.rightscore):
                judge = '負けちゃったね…'
                self.match_result = 'LOSE'

            else:
                judge = 'めずらしい！引き分けだね'
                self.match_result = 'DRAW'

            self.m = f'{self.leftscore} vs { self.rightscore}で{judge}\n'

        except Exception as e:
            print(e)
            self.m = '勝敗がわかりませんでした…\n'

    def detect_color(self, img, mode):
        # HSV色空間に変換
        hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)

        if mode == 'green':
            hsv_min = np.array([80, 50, 150])
            hsv_max = np.array([140, 255, 255])

        elif mode == 'brown':
            hsv_min = np.array([0, 20, 100])
            hsv_max = np.array([60, 180, 255])

        mask = cv2.inRange(hsv, hsv_min, hsv_max)

        # マスキング処理
        masked_img = cv2.bitwise_and(img, img, mask=mask)

        return mask, masked_img

    def player_recog(self):
        # OCRエンジンを取得
        # os.environ["TESSDATA_PREFIX"] = "/usr/share/tesseract-ocr/5/tessdata"

        LOC_X = 1525
        LOC_Y = 242
        target_img = self.input_img[LOC_Y:LOC_Y + 323, LOC_X:LOC_X + 300]
        Estimated_players = []
        Estimated_clan_mems = []

        for i in range(6):
            trim_img = target_img[i * 48:32 + i * 48, :]
            # trim_img = target_img
            participant_img = cv2.cvtColor(trim_img, cv2.COLOR_RGB2GRAY)

            ret, bin_participant_img = cv2.threshold(participant_img, 127, 255, cv2.THRESH_BINARY)
            txt_participants = pytesseract.image_to_string(Image.fromarray(bin_participant_img), lang="eng")
            green_mask, green_masked_img = self.detect_color(trim_img, 'green')
            brown_mask, brown_masked_img = self.detect_color(trim_img, 'brown')
            # green_masked_img = cv2.cvtColor(green_masked_img, cv2.COLOR_BGR2GRAY)
            green_masked_img = cv2.blur(green_masked_img, (1, 1))
            brown_masked_img = cv2.blur(brown_masked_img, (1, 1))
            txt_player = pytesseract.image_to_string(Image.fromarray(green_masked_img), lang="eng")
            txt_player = re.sub(r"[,.!?:;'\s\r\n]", "", txt_player)
            txt_clan = pytesseract.image_to_string(Image.fromarray(brown_masked_img), lang="eng")
            txt_clan = re.sub(r"[,.!?:;'\s\r\n]", "", txt_clan)
            image_size = green_mask.size
            whitePixels_green = cv2.countNonZero(green_mask)
            whitePixels_brown = cv2.countNonZero(brown_mask)

            if txt_player != '':
                Estimated_players.append([txt_player, round(whitePixels_green / image_size, 5), (LOC_X + 300, LOC_Y + 32 + i * 48), green_masked_img])
                self.num_vs += 1

            elif txt_clan != '':
                Estimated_clan_mems.append([txt_clan, round(whitePixels_brown / image_size, 5), (LOC_X + 300, LOC_Y + 32 + i * 48), brown_masked_img])
                self.num_vs += 1

            elif txt_participants != '':
                self.num_vs += 1

            # print(txt_participants)
            # cv2.imshow('bin', bin_participant_img)

            # cv2.imshow('trim', trim_img)

            # cv2.imshow('trim', green_mask)
            # print()
            # cv2.imshow('trim', brown_masked_img)
            # cv2.waitKey(0)
            # cv2.destroyAllWindows()

        if Estimated_players:
            if len(Estimated_players) == 1:
                # player_id = re.sub(r"[,.!?:;' ]", "", Estimated_players[0][0])
                # player_id = re.sub('[\r\n]+$', '', player_id, flags=re.MULTILINE)
                player_id = re.sub(r"[,.!?:;'\s\r\n]", "", Estimated_players[0][0])
                player_id_loc = Estimated_players[0][2]
                self.id_recog_result_img = Estimated_players[0][3]
                # print(file, player_id)

            else:
                es_players = sorted(Estimated_players, key=lambda x: x[1], reverse=True)
                player_id = re.sub(r"[,.!?:;'\s\r\n]", "", es_players[0][0])
                # player_id = re.sub('[\r\n]+$', '', player_id, flags=re.MULTILINE)
                player_id_loc = es_players[0][2]
                self.id_recog_result_img = es_players[0][3]
                # print(sorted(Estimated_players, key=lambda x: x[1]))
                # print(file, player_id)

            self.player_id = player_id
            # print(judge.judge(input_img), player_id)

        if self.player_id:
            self.m += f'プレイヤーのidは {self.player_id} です！\n'
            cv2.rectangle(self.result_img, (LOC_X, player_id_loc[1] - 32), player_id_loc, (0, 255, 0), 3)

        else:
            self.m += 'プレイヤーのidはわかりませんでした…\n'

        if Estimated_clan_mems:
            self.clan_mems = Estimated_clan_mems.copy()
            self.m += 'クランメンバーのidは…\n'
            for mem in self.clan_mems:
                self.m += f'{mem[0]}\n'
                self.clanmem_id.append(mem[0])
                cv2.rectangle(self.result_img, (LOC_X, mem[2][1] - 32), mem[2], (0, 0, 255), 3)

            self.m += "です!"

        ret, img = cv2.imencode('.png', self.result_img)
        self.result_img_byte = io.BytesIO(img)
        ret, img = cv2.imencode('.png', self.id_recog_result_img)
        self.id_recog_result_img_byte.append(io.BytesIO(img))
        if self.clan_mems:
            for mem in self.clan_mems:
                ret, img = cv2.imencode('.png', mem[3])
                self.id_recog_result_img_byte.append(io.BytesIO(img))


"""model accu
Construct model accu
"""
import math
from mindspore.train.callback import RunContext
from mindspore import context
from mindspore.context import ParallelMode
from mindspore import Model, connect_network_with_dataset
from mindspore.common.dtype import pytype_to_dtype
from mindspore._c_expression import init_exec_dataset
from mindspore.train.train_thor.dataset_helper import DatasetHelper
from mindspore.parallel._utils import _need_to_full, _to_full_tensor

def _convert_type(types):
    """
    Convert from numpy type to tensor type.

    Args:
        types (list): Numpy type list of element in dataset.

    Returns:
        list, list of element in dataset.
    """
    ms_types = []
    for np_type in types:
        ms_type = pytype_to_dtype(np_type)
        ms_types.append(ms_type)
    return ms_types

def _get_types_and_shapes(dataset):
    """Get dataset types and shapes."""
    dataset_types = _convert_type(dataset.output_types())
    dataset_shapes = dataset.output_shapes()
    return dataset_types, dataset_shapes

def _exec_datagraph(exec_dataset, dataset_size, phase='dataset'):
    """Initialize and execute the dataset graph."""
    batch_size = exec_dataset.get_batch_size()
    input_indexs = exec_dataset.input_indexs

    # transform data format
    dataset_types, dataset_shapes = _get_types_and_shapes(exec_dataset)
    init_exec_dataset(exec_dataset.__transfer_dataset__.queue_name,
                      dataset_size,
                      batch_size,
                      dataset_types,
                      dataset_shapes,
                      input_indexs,
                      phase=phase,
                      need_run=False)

class Model_ACCU(Model):
    """"Construct Model_ACCU"""
    def __init__(self, network, loss_fn=None, optimizer=None, metrics=None, eval_network=None,
                 eval_indexes=None, amp_level="O0", **kwargs):
        super(Model_ACCU, self).__init__(network, loss_fn, optimizer, metrics, eval_network,
                                         eval_indexes, amp_level, **kwargs)
        self._frequency = context.get_auto_parallel_context("grad_accumulation_step")
        # used to stop training for early stop, such as stopAtTIme or stopATStep
        self.should_stop = False
        self.switch_branch_one = True
        self.index_second_order = 0
        self.train_network_init_flag = True
        self.has_do_dataset_init = False
        self._train_network = self._build_train_network()

    def _exec_preprocess(self, network, is_train, phase, dataset, dataset_sink_mode, sink_size=-1,
                         epoch_num=1, iter_update_order=1):
        """Initializes dataset."""
        if dataset_sink_mode and not is_train:
            dataset.__loop_size__ = 1
        dataset_helper = DatasetHelper(dataset, dataset_sink_mode, sink_size, epoch_num, iter_update_order)

        if dataset_sink_mode and context.get_context("device_target") != "GPU":
            network = connect_network_with_dataset(network, dataset_helper)
        network.set_train(is_train)
        network.phase = phase

        if self._parallel_mode in (ParallelMode.SEMI_AUTO_PARALLEL, ParallelMode.AUTO_PARALLEL):
            network.set_auto_parallel()

        return dataset_helper, network

    def _train_gpu_sink_step(self, cb_params, inputs, list_callback, iter_accu_order, run_context):
        """train gpu sink step"""
        if self.switch_branch_one:
            cb_params.cur_step_num += 1
            if self.train_network_init_flag:
                self._train_network.add_flags_recursive(accumulation=True)
            self._train_network.phase = 'train0'
            outputs = self._train_network(*inputs)
            cb_params.net_outputs = outputs
            self.index_second_order += 1
            if self.index_second_order == iter_accu_order:
                self.index_second_order = 0
                self.switch_branch_one = not self.switch_branch_one
                list_callback.step_end(run_context)
        else:
            cb_params.cur_step_num += 1
            if self.train_network_init_flag:
                self._train_network.add_flags_recursive(accumulation=False)
                self.train_network_init_flag = False
            self._train_network.phase = 'train1'
            self.switch_branch_one = not self.switch_branch_one
            outputs = self._train_network(*inputs)
            cb_params.net_outputs = outputs
            list_callback.step_end(run_context)

    def _train_ascend_sink_step(self, cb_params, train_dataset, iter_accu_order, inputs, list_callback, run_context):
        """train ascend sink step"""
        if self.switch_branch_one:
            cb_params.cur_step_num += iter_accu_order
            if self.train_network_init_flag:
                self._train_network.add_flags_recursive(accumulation=True)
            self._train_network.phase = 'train0'
        else:
            cb_params.cur_step_num += 1
            if self.train_network_init_flag:
                self._train_network.add_flags_recursive(accumulation=False)
                self.train_network_init_flag = False
            self._train_network.phase = 'train1'
            if not self.has_do_dataset_init:
                _exec_datagraph(train_dataset, 1, phase='train1_dataset')
                self.has_do_dataset_init = True
        self.switch_branch_one = not self.switch_branch_one
        outputs = self._train_network(*inputs)
        cb_params.net_outputs = outputs
        list_callback.step_end(run_context)

    def _train_dataset_sink_process(self, epoch, train_dataset, list_callback=None, cb_params=None, sink_size=-1):
        """
        Training process. The data would be passed to network through dataset channel.

        Args:
            epoch (int): Total number of iterations on the data.
            train_dataset (Dataset): A training dataset iterator. If there is no
                                     loss_fn, a tuple with multiple data (data1, data2, data3, ...) should be
                                     returned and passed to the network. Otherwise, a tuple (data, label) should
                                     be returned. The data and label would be passed to the network and loss
                                     function respectively.
            list_callback (Callback): Executor of callback list. Default: None.
            cb_params (_InternalCallbackParam): Callback parameters. Default: None.
            sink_size (int): Control the amount of data in each sink. Default: -1.
        """
        if sink_size == -1:
            epoch_num = epoch
        else:
            epoch_num = math.ceil(epoch * sink_size / train_dataset.get_dataset_size())

        iter_update_order = 1
        iter_accu_order = self._frequency - 1
        if context.get_context("device_target") == "GPU":
            train_dataset.__loop_size__ = 1
        else:
            train_dataset.__loop_size__ = iter_accu_order
        dataset_helper, train_network = self._exec_preprocess(self._train_network,
                                                              is_train=True,
                                                              phase='train',
                                                              dataset=train_dataset,
                                                              dataset_sink_mode=True,
                                                              sink_size=sink_size,
                                                              epoch_num=epoch_num,
                                                              iter_update_order=iter_update_order)

        self._train_network = train_network
        cb_params.train_network = self._train_network
        cb_params.cur_step_num = 0

        run_context = RunContext(cb_params)
        list_callback.begin(run_context)

        for i in range(epoch):
            cb_params.cur_epoch_num = i + 1
            list_callback.epoch_begin(run_context)
            # for data sink dataset_helper only iter once, other wise iter epoch_size times.
            for inputs in dataset_helper:
                if _need_to_full() and context.get_context("device_target") == "GPU":
                    inputs = _to_full_tensor(inputs, self._device_number, self._global_rank)
                list_callback.step_begin(run_context)
                if context.get_context("device_target") == "GPU":
                    self._train_gpu_sink_step(cb_params, inputs, list_callback, iter_accu_order, run_context)
                else:
                    self._train_ascend_sink_step(cb_params, train_dataset, iter_accu_order, inputs, list_callback,
                                                 run_context)
            list_callback.epoch_end(run_context)
            self.should_stop = self.should_stop or run_context.get_stop_requested()
            if self.should_stop:
                break
        dataset_helper.stop_send()

        list_callback.end(run_context)


# python 3.6
"""Inverts given images to latent codes with In-Domain GAN Inversion.

Basically, for a particular image (real or synthesized), this script first
employs the domain-guided encoder to produce a initial point in the latent
space and then performs domain-regularized optimization to refine the latent
code.
"""

import os
import argparse
import pickle
from training.misc import progress as tqdm
import numpy as np
import tensorflow as tf
from dnnlib import tflib

from utils.logger import setup_logger
from utils.visualizer import adjust_pixel_range
from utils.visualizer import HtmlPageVisualizer
from utils.visualizer import save_image, load_image, resize_image

def parse_args():
  """Parses arguments."""
  parser = argparse.ArgumentParser()
  parser.add_argument('model_path', type=str,
                      help='Path to the pre-trained model.')
  parser.add_argument('image_list', type=str,
                      help='List of images to invert.')
  parser.add_argument('-o', '--output_dir', type=str, default='',
                      help='Directory to save the results. If not specified, '
                           '`./results/inversion/${IMAGE_LIST}` '
                           'will be used by default.')
  parser.add_argument('--batch_size', type=int, default=50,
                      help='Batch size. (default: 32)')
  parser.add_argument('--viz_size', type=int, default=256,
                      help='Image size for visualization. (default: 256)')
  parser.add_argument('--gpu_id', type=str, default='0',
                      help='Which GPU(s) to use. (default: `0`)')
  parser.add_argument('--radius', type=int, default=3,
                      help='radius of each latent w')
  parser.add_argument('--start', type=int, default=0,
                      help='decide start from which images')
  parser.add_argument('--end', type=int, default=1,
                      help='decide end with which images')
  parser.add_argument('--max_num_layers', type=int, default=8,
                      help='the maximum layer to replace')
  parser.add_argument('--total_nums', type=int, default=1e+5,
                      help='the maximum layer to replace')
  parser.add_argument('--save_raw', action='store_true',
                      help='Whether to save raw images')
  return parser.parse_args()

def main():
  """Main function."""
  args = parse_args()
  os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu_id
  assert os.path.exists(args.image_list)
  image_list_name = os.path.splitext(os.path.basename(args.image_list))[0]
  output_dir = args.output_dir or f'results/inversion/{image_list_name}'
  logger = setup_logger(output_dir, 'inversion.log', 'inversion_logger')

  logger.info(f'Loading model.')
  tflib.init_tf({'rnd.np_random_seed': 1000})
  with open(args.model_path, 'rb') as f:
    E, _, _, Gs = pickle.load(f)
  # Get input size.
  image_size = E.input_shape[2]
  assert image_size == E.input_shape[3]
  input_shape = E.input_shape
  num_layers, z_dim = Gs.components.synthesis.input_shape[1:3]

  # ATTRS = {'age': [1,2,3], 'pose': [0,1,2], 'male': [2,3,4,5], 'expression': [3,4,5], 'glass': [1,2,3]}
  # Build graph.
  logger.info(f'Building graph.')
  sess = tf.get_default_session()
  x = tf.placeholder(tf.float32, shape=input_shape, name='real_image')
  latent_w = tf.placeholder(tf.float32, shape=[None, num_layers, z_dim], name='latent_w')
  w_enc = E.get_output_for(x, is_training=False)
  w_enc = tf.reshape(w_enc, [-1, 14, 512])
  coef = args.radius / np.sqrt(z_dim / 3.0)

  x_rec = Gs.components.synthesis.get_output_for(latent_w, randomize_noise=False)
  # Load image list.
  logger.info(f'Loading image list.')
  image_list = []
  with open(args.image_list, 'r') as f:
    for line in f:
      image_list.append(line.strip())
  image_list = image_list[args.start:args.end]
  images = []
  names  = []
  for image_name in image_list:
    image = load_image(image_name)
    images.append(np.transpose(image, [2, 0, 1]))
    names.append(os.path.splitext(os.path.basename(image_name))[0])
  images = np.asarray(images)
  logger.info(f'images shape {images.shape}')
  images = images.astype(np.float32) / 255 * 2.0 - 1.0
  aug = np.zeros([args.batch_size, num_layers, z_dim])
  for idx in range(images.shape[0]):
    imgs = images[idx:idx+1]
    w_enc_ = sess.run(w_enc, {x: imgs})
    x_rec_ = sess.run(x_rec, {latent_w: w_enc_})
    imgs_ = adjust_pixel_range(x_rec_)
    if args.save_raw:
      save_image(f'{output_dir}/{names[idx]}_enc_init.png', imgs_[0])
    w_enc_ = np.tile(w_enc_, [args.batch_size, 1, 1])
    images_npy = []
    latent_w_s = []
    for it in tqdm(range(0, args.total_nums, args.batch_size)):
      num_layer_replace = np.random.randint(low=1, high=args.max_num_layers)
      start_layer_replace = np.random.randint(low=0, high=args.max_num_layers - num_layer_replace)
      end_layer_replace = start_layer_replace + num_layer_replace
      uniform = np.random.uniform(low=-1.0, high=1.0, size=[args.batch_size, num_layer_replace, z_dim]) * coef
      aug_square = np.concatenate([aug[:, :start_layer_replace], uniform, aug[:, end_layer_replace:]], axis=1)
      latent_w_square = w_enc_ + aug_square
      x_rec_ = sess.run(x_rec, {latent_w: latent_w_square})
      imgs_ = adjust_pixel_range(x_rec_)
      if args.save_raw:
        for ii in range(imgs_.shape[0]):
          save_image(f'{output_dir}/{names[idx]}_rand_{it +ii:04d}.png', imgs_[ii])
      latent_w_s.append(latent_w_square)
      images_npy.append(imgs_)
    latent_w_s = np.concatenate(latent_w_s, axis=0)
    images_npy = np.concatenate(images_npy, axis=0)
    logger.info(f'latent_w_s shape: {latent_w_s.shape}, Saving...')
    np.save(f'{output_dir}/{names[idx]}_latent.npy', latent_w_s)
    logger.info(f'images shape: {images_npy.shape}, Saving...')
    np.save(f'{output_dir}/{names[idx]}_image.npy', images_npy)

if __name__ == '__main__':
  main()


from tqdm import tqdm as tqdm
import numpy as np
import skimage
import lmdb
import os
from skimage.transform import resize
import deepracing.imutils
import ChannelOrder_pb2
import Image_pb2
import TimestampedImage_pb2
import cv2
import time
import google.protobuf.empty_pb2 as Empty_pb2
import yaml
import PIL, PIL.Image as PILImage
import torchvision, torchvision.transforms.functional as F
import torchvision.transforms as TF
import json
import warnings

def pbImageToNpImage(im_pb : Image_pb2.Image):
    im = None
    if im_pb.channel_order == ChannelOrder_pb2.BGR:
        im = np.reshape(np.frombuffer(im_pb.image_data,dtype=np.uint8),np.array((im_pb.rows, im_pb.cols, 3)))
        im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)
    elif im_pb.channel_order == ChannelOrder_pb2.RGB:
        im = np.reshape(np.frombuffer(im_pb.image_data,dtype=np.uint8),np.array((im_pb.rows, im_pb.cols, 3)))
    elif im_pb.channel_order == ChannelOrder_pb2.GRAYSCALE:
        im = np.reshape(np.frombuffer(im_pb.image_data,dtype=np.uint8),np.array((im_pb.rows, im_pb.cols)))
    elif im_pb.channel_order == ChannelOrder_pb2.RGBA:
        im = np.reshape(np.frombuffer(im_pb.image_data,dtype=np.uint8),np.array((im_pb.rows, im_pb.cols, 4)))
    elif im_pb.channel_order == ChannelOrder_pb2.BGRA:
        im = np.reshape(np.frombuffer(im_pb.image_data,dtype=np.uint8),np.array((im_pb.rows, im_pb.cols, 4)))
        im = cv2.cvtColor(im, cv2.COLOR_BGRA2RGBA)
    else:
        raise ValueError("Unknown channel order: " + im_pb.channel_order)
    return im#.copy()
class ImageArrayBackend():
    def __init__(self, keyfile ):
        with open(keyfile,"r") as keyfile_:
            self.keys = keyfile_.readlines()
        for i in range(len(self.keys)):
            self.keys[i] = self.keys[i].replace("\n","")

        self.keymap = {self.keys[i]:i for i in range(len(self.keys))}
        self.images = None
    def loadImages(self, image_folder, ROI = None, imsize = (66,200) ):
        print("Loading images in %s" %(image_folder))
        self.images = np.zeros((len(self.keys),imsize[0],imsize[1],3), dtype=np.uint8)
        if ROI is not None:
            x = ROI[0]
            y = ROI[1]
            w = ROI[2]
            h = ROI[3]
        for i, key in tqdm(enumerate(self.keys), total=len(self.keys)):
            fp = os.path.join(image_folder,key+".jpg")
           # print(fp)
            imin = deepracing.imutils.readImage(fp)
            if ROI is not None:
                imin = imin[y:y+h,x:x+w,:]
            self.images[i] = deepracing.imutils.resizeImage(imin, imsize)
    def getImage( self, key : str ):
        return self.images[self.keymap[key]]
class ImageFolderWrapper():
    def __init__(self, image_folder):
        self.image_folder = image_folder
    def getImage( self, key : str ):
        fp = os.path.join(self.image_folder,key+".jpg")
        return deepracing.imutils.readImage(fp)
class ImageLMDBWrapper():
    def __init__(self, encoding = "ascii"):
        self.env = None
        self.encoding = encoding
        self.spare_txns=1
        self.internal_cache = {}
    def readImages(self, image_files, keys, db_path, im_size, ROI=None, mapsize=int(1e10), tqdm=True):
        assert(len(image_files) > 0)
        assert(len(image_files) == len(keys))
        env = lmdb.open(db_path, map_size=mapsize)
        print("Loading image data", flush=True)
        cropped_images_dir = os.path.join(os.path.dirname(db_path),"cropped_images")
        os.makedirs(cropped_images_dir,exist_ok=True)
        if tqdm:
            t = tqdm(enumerate(keys), total=len(keys))
        else:
            t = enumerate(keys)
        for i, key in t:
            imgfile = image_files[i]
            base, ext = os.path.splitext(imgfile)
            system_time = None
            try:
                with open(base+".json","r") as f:
                    imdict = json.load(f)
                assert(imdict["imageFile"]==os.path.basename(imgfile))
                system_time = imdict["timestamp"]/1000.0
            except:
                try:
                    with open(base+".pb","rb") as f:
                        pb = TimestampedImage_pb2.TimestampedImage()
                        pb.ParseFromString(f.read())
                    assert(pb.image_file==os.path.basename(imgfile))
                    system_time = pb.timestamp/1000.0
                except:
                    warnings.warn("Could not get a timestamp for image file: %s" %(imgfile,), ResourceWarning)
            imgin = deepracing.imutils.readImage(imgfile)
            impil = F.to_pil_image(imgin)
            if bool(ROI):
                r = ROI[0]
                c = ROI[1]
                h = ROI[2]
                w = ROI[3]
                impilresize = F.resized_crop(impil,r,c,h,w,im_size,interpolation=TF.InterpolationMode.LANCZOS)
            else:
                impilresize = F.resize(impil,im_size, interpolation=TF.InterpolationMode.LANCZOS)
            impilresize.save(os.path.join(cropped_images_dir, "%s.jpg" % (key,)))
            im = np.asarray(impilresize)
            entry = Image_pb2.Image( rows=im.shape[0] , cols=im.shape[1] , channel_order=ChannelOrder_pb2.RGB , image_data=im.flatten().tobytes(), system_time=-1.0 )
            if system_time is not None:
                entry.system_time=system_time
            with env.begin(write=True) as write_txn:
                write_txn.put(key.encode(self.encoding), entry.SerializeToString())
        env.close()

    def writeImage(self, key, image):
        imarr = np.asarray(image)
        entry = Image_pb2.Image( rows=imarr.shape[0] , cols=imarr.shape[1] , channel_order=ChannelOrder_pb2.RGB , image_data=imarr.flatten().tobytes() )
        with self.env.begin(write=True) as write_txn:
            write_txn.put(key.encode(self.encoding), entry.SerializeToString())
        return entry
    def getDBPath(self):
        return self.env.path()
    def clearStaleReaders(self):
        self.env.reader_check()
    def resetEnv(self):
        if self.env is not None:
            path = self.env.path()
            mapsize = self.env.info()['map_size']
            self.env.close()
            del self.env
            time.sleep(1)
            self.readDatabase(path, mapsize=mapsize, max_spare_txns=self.spare_txns, readonly=True, lock=False)
    def readDatabase(self, db_path : str, mapsize=int(1e10), max_spare_txns=125, readonly=True, lock=False):
        if not os.path.isdir(db_path):
            raise IOError("Path " + db_path + " is not a directory")
        self.spare_txns = max_spare_txns
        self.env = lmdb.open(db_path, map_size=round(mapsize,None), max_spare_txns=max_spare_txns,\
            create=False, lock=lock, readonly=readonly)
    def getImagePB(self, key : str):
        im_pb = Image_pb2.Image()
        with self.env.begin(write=False) as txn:
            entry = txn.get( key.encode( self.encoding ) )
            if (entry is None):
                raise ValueError("Invalid key: %s on image database: %s" %(key, str(self.env.path())))
            im_pb.ParseFromString( entry )
        return im_pb
    def getImage( self, key : str ):
        pbimage = self.getImagePB( key )
        return pbimage.system_time, pbImageToNpImage( pbimage )
    def getNumImages(self):
        return self.env.stat()['entries']
    def getKeys(self):
        keys = None
        with self.env.begin(write=False) as txn:
            keys = [ str(key, encoding=self.encoding) for key, _ in txn.cursor() ]
        if (keys is None) or len(keys)==0:
            raise ValueError("Keyset is empty in image dataset for some reason")
        return keys