from ..video.base import BaseVideoFilter
from ..audio.base import BaseAudioFilter
from ..base import BaseFilter
import numpy
from itertools import count
from more_itertools import windowed
from ...util import cached, SourceError, IncompatibleSource, BrokenReference
from fractions import Fraction as QQ
from copy import deepcopy
import weakref
import regex

def tryweakref(obj):
    try:
        return weakref.ref(obj)

    except Exception:
        obj

class Concatenate(BaseVideoFilter, BaseAudioFilter):
    __name__ = "Concatenate"
    allowedtypes = ("audio", "video")
    sourceCount = "+"

    def __init__(self, segments=[], time_base=QQ(1, 10**9), **kwargs):
        self.segments = list(map(tryweakref, segments))
        self.time_base = time_base
        super().__init__(**kwargs)

    def __iter__(self):
        for item in self.segments:
            if isinstance(item, weakref.ref):
                yield item()

            else:
                yield item

    @property
    def dependencies(self):
        dependencies = set(self)

        for source in self:
            if isinstance(source, BaseFilter):
                dependencies.update(source.dependencies)

        return dependencies

    @cached
    def pts_time(self):
        L = []
        T = 0

        for segment in self:
            L.append(segment.pts_time + T)
            T += segment.duration

        if len(L):
            return numpy.concatenate(L)

        return numpy.array((), dtype=numpy.float64)

    @cached
    def pts(self):
        return numpy.int0(self.pts_time/float(self.time_base) + 0.0001)

    @cached
    def duration(self):
        return sum(segment.duration for segment in self)

    @cached
    def durations(self):
        if len(self):
            return numpy.int0(numpy.concatenate([
                segment.durations*segment.time_base/self.time_base
                for segment in self]) + 0.0001)

        return numpy.array((), dtype=numpy.float64)

    @cached
    def sizes(self):
        if len(self):
            return numpy.concatenate([segment.sizes for segment in self])

        return numpy.array((), dtype=numpy.int0)

    @cached
    def framecount(self):
        return sum(segment.framecount for segment in self)

    @property
    def type(self):
        if len(self):
            return self[0].type

    @property
    def codec(self):
        if len(self):
            return self[0].codec

    @property
    def extradata(self):
        if len(self):
            return self[0].extradata

    @property
    def format(self):
        if len(self):
            return self[0].format

    @property
    def bitdepth(self):
        if len(self):
            return self[0].bitdepth

    @property
    def defaultDuration(self):
        if len(self):
            return self[0].defaultDuration

    @property
    def rate(self):
        if len(self):
            return self[0].rate

    @property
    def layout(self):
        if len(self):
            return self[0].layout

    @property
    def channels(self):
        if len(self):
            return self[0].channels

    @property
    def width(self):
        if len(self):
            return self[0].width

    @property
    def height(self):
        if len(self):
            return self[0].height

    @property
    def sar(self):
        if len(self):
            return self[0].sar

    @property
    def prev(self):
        raise AttributeError

    @prev.setter
    def prev(self, value):
        return

    @property
    def cumulativeIndexMap(self):
        return numpy.arange(self.framecount)

    def iterFrames(self, start=0, end=None, whence=None):
        if self.type == "video" and whence is None:
            whence = "framenumber"

        elif self.type in ("audio", "subtitle") and whence is None:
            whence = "seconds"

        N = 0
        T = 0

        for segment in self:
            if whence == "framenumber":
                if start >= N + segment.framecount:
                    N += segment.framecount
                    T += segment.duration
                    continue

                elif N <= start:
                    if end is not None and end < N + segment.framecount:
                        frames = segment.iterFrames(start - N, end - N, whence)

                    else:
                        frames = segment.iterFrames(start - N, None, whence)

                elif end is None or end >= N + segment.framecount:
                    frames = segment.iterFrames()

                elif end > N:
                    frames = segment.iterFrames(0, end - N, whence)

                else:
                    break

            elif whence == "pts":
                if start*segment.time_base >= T + segment.duration:
                    N += segment.framecount
                    T += segment.duration
                    continue

                elif T <= start*segment.time_base:
                    if (end is not None
                            and end*segment.time_base < T + segment.duration):
                        frames = segment.iterFrames(
                            start - int(T/segment.time_base + 0.5),
                            end - int(T/segment.time_base + 0.5), whence)

                    else:
                        frames = segment.iterFrames(
                            start - int(T/segment.time_base + 0.5),
                            None, whence)

                elif (end is None
                      or end*segment.time_base >= T + segment.duration):
                    frames = segment.iterFrames()

                elif end*segment.time_base > T:
                    frames = segment.iterFrames(
                        0, end - int(T/segment.time_base + 0.5), whence)

                else:
                    break

            elif whence == "seconds":
                if start >= T + segment.duration:
                    N += segment.framecount
                    T += segment.duration
                    continue

                elif T <= start:
                    if end is not None and end < T + segment.duration:
                        frames = segment.iterFrames(start - T, end - T, whence)

                    else:
                        frames = segment.iterFrames(start - T, None, whence)

                elif end is None or end >= T + segment.duration:
                    frames = segment.iterFrames()

                elif end > T:
                    frames = segment.iterFrames(0, end - T, whence)

                else:
                    break

            for frame in frames:
                frame.pts = int(
                    (T + frame.pts*frame.time_base)/self.time_base + 0.5)
                frame.time_base = self.time_base
                yield frame

            N += segment.framecount
            T += segment.duration

    def iterPackets(self, start=0, end=None, whence="pts"):
        if self.type == "video" and whence is None:
            whence = "packetnumber"

        elif self.type in ("audio", "subtitle") and whence is None:
            whence = "seconds"

        N = 0
        T = 0

        if whence == "pts":
            K = count(self.frameIndexFromPts(start, "+"))

        elif whence == "seconds":
            K = count(self.frameIndexFromPtsTime(start, "+"))

        else:
            K = count(start)

        for segment in self:
            if whence == "packetnumber":
                if start >= N + segment.framecount:
                    N += segment.framecount
                    T += segment.duration
                    continue

                elif N <= start:
                    if end is not None and end < N + segment.framecount:
                        packets = segment.iterPackets(
                            start - N, end - N, whence)

                    else:
                        packets = segment.iterPackets(start - N, None, whence)

                elif end is None or end >= N + segment.framecount:
                    packets = segment.iterPackets()

                elif end > N:
                    packets = segment.iterPackets(None, end - N, whence)

                else:
                    break

            elif whence == "pts":
                if start*segment.time_base >= T + segment.duration:
                    N += segment.framecount
                    T += segment.duration
                    continue

                elif T <= start*segment.time_base:
                    if (end is not None
                            and end*segment.time_base < T + segment.duration):
                        packets = segment.iterPackets(
                            start - int(T/segment.time_base + 0.5),
                            end - int(T/segment.time_base + 0.5), whence)

                    else:
                        packets = segment.iterPackets(
                            start - int(T/segment.time_base + 0.5),
                            None, whence)

                elif (end is None
                      or end*segment.time_base >= T + segment.duration):
                    packets = segment.iterPackets()

                elif end*segment.time_base > T:
                    packets = segment.iterPackets(
                        None, end - int(T/segment.time_base + 0.5), whence)

                else:
                    break

            elif whence == "seconds":
                if start >= T + segment.duration:
                    N += segment.framecount
                    T += segment.duration
                    continue

                elif T <= start:
                    if end is not None and end < T + segment.duration:
                        packets = segment.iterPackets(
                            start - T, end - T, whence)

                    else:
                        packets = segment.iterPackets(start - T, None, whence)

                elif end is None or end >= T + segment.duration:
                    packets = segment.iterPackets()

                elif end > T:
                    packets = segment.iterPackets(None, end - T, whence)

                else:
                    break

            for packet in packets:
                packet.pts = int(
                    (T + packet.pts*packet.time_base)/self.time_base + 0.5)

                if packet.duration:
                    packet.duration = int(
                        packet.duration*packet.time_base/self.time_base + 0.5)

                packet.time_base = self.time_base

                if self[0].source.codec == "ass":
                    match, = regex.findall(b"\\d,(.+)", packet.data)
                    packet.data = str(next(K)).encode("utf8") + b"," + match

                yield packet

            N += segment.framecount
            T += segment.duration

    def append(self, segment):
        self.segments.append(tryweakref(segment))

        if isinstance(segment, BaseFilter):
            segment.addMonitor(self)

    def insert(self, index, segment):
        self.segments.insert(index, tryweakref(segment))

        if isinstance(segment, BaseFilter):
            segment.addMonitor(self)

    def extend(self, segments):
        k = len(self)
        self.segments.extend(map(tryweakref, segments))

        for segment in self.segments[k:]:
            if isinstance(segment, weakref.ref):
                segment = segment()

            if isinstance(segment, BaseFilter):
                segment.addMonitor(self)

    def clear(self):
        for segment in self:
            if isinstance(segment, BaseFilter):
                segment.removeMonitor(self)

        self.segments.clear()

    def __getitem__(self, index):
        item = self.segments[index]

        if isinstance(index, int) and isinstance(item, weakref.ref):
            item = item()

        elif isinstance(index, slice):
            item = [x() if isinstance(x, weakref.ref) else x for x in item]

        return item

    def __len__(self):
        return len(self.segments)

    def __delitem__(self, index):
        segment = self.segments.pop(index)

        if isinstance(segment, weakref.ref):
            segment = segment()

        if isinstance(segment, BaseFilter) and segment not in self.segments:
            segment.removeMonitor(self)

    def __setitem__(self, index, value):
        oldvalue = self[index]
        self.segments[index] = tryweakref(value)

        if isinstance(value, BaseFilter):
            value.addMonitor(self)

        if isinstance(oldvalue, BaseFilter) and oldvalue not in self.segments:
            oldvalue.removeMonitor(self)

    def __reduce__(self):
        return type(self), (), self.__getstate__(), iter(self)

    def __getstate__(self):
        state = super().__getstate__()

        if self.time_base:
            state["time_base"] = self.time_base

        return state

    def __setstate__(self, state):
        self.time_base = state.get("time_base", QQ(1, 10**9))
        super().__setstate__(state)

    def __deepcopy__(self, memo):
        cls, args, state, segments = self.__reduce__()
        new = cls()
        new.__setstate__(deepcopy(state, memo))
        new.extend(segments)
        return new

    @property
    def source(self):
        raise AttributeError

    @cached
    def keyframes(self):
        kf = set()
        n = 0

        for segment in self:
            kf.update(k + n for k in segment.keyframes)
            n += segment.framecount

        return kf

    @staticmethod
    def QtDlgClass():
        from .qconcatenate import QConcatenate
        return QConcatenate

    def validate(self):
        if len(self) == 0:
            return [SourceError("No sources provided.", self)]

        exceptions = []

        for k, s in enumerate(self):
            if s is None:
                exceptions.append(BrokenReference(
                    f"Broken reference at position {k}.", self))

        for ((j, source1), (k, source2)) in windowed(filter(
                lambda X: X[1] is not None and X[1].prev is not None,
                enumerate(self)), 2):
            if source1.type != source2.type:
                exceptions.append(IncompatibleSource(
                    f"Incompatible sources: '{source1.type}' and "
                    f"'{source2.type}'.", self))

            elif source1.type == "video":
                if ((source1.width, source1.height)
                        != (source2.width, source2.height)):
                    exceptions.append(IncompatibleSource(
                        f"Sources have different resolutions: "
                        f"'{source1.width}×{source1.height}' ({j}) and "
                        f"'{source2.width}×{source2.height}' ({k}).", self))

                if source1.sar != source2.sar:
                    exceptions.append(IncompatibleSource(
                        f"Sources have different sample aspect ratios: "
                        f"'{source1.sar}' ({j}) and '{source2.sar}' ({k}).",
                        self))

            elif source1.type == "audio":
                if source1.rate != source2.rate:
                    exceptions.append(IncompatibleSource(
                        f"Sources have different sampling frequencies: "
                        f"'{source1.rate}' ({j}) and '{source2.rate}' ({k}).",
                        self))

                if source1.layout != source2.layout:
                    exceptions.append(IncompatibleSource(
                        f"Sources have different layouts: '{source1.layout}' "
                        f"({j}) and '{source2.layout}' ({k}).", self))

        return exceptions

    @property
    def canIterPackets(self):
        if len(self) == 0:
            return False

        for k, item in enumerate(self):
            if not item.canIterPackets:
                return False

            if k == 0:
                codec = item.codec
                extradata = item.extradata

            else:
                if item.codec != codec:
                    return False

                if (self.type in ("video", "audio")
                        and item.extradata != extradata):
                    return False

        return True


import numpy as np
from typing import Dict, List, Optional, Tuple
import torch
from torch import Tensor, nn

from detectron2.data.detection_utils import convert_image_to_rgb
from detectron2.modeling import Backbone
from detectron2.structures import Boxes, ImageList, Instances
from detectron2.utils.events import get_event_storage

from ..postprocessing import detector_postprocess

def permute_to_N_HWA_K(tensor, K: int):
    """
    Transpose/reshape a tensor from (N, (Ai x K), H, W) to (N, (HxWxAi), K)
    """
    assert tensor.dim() == 4, tensor.shape
    N, _, H, W = tensor.shape
    tensor = tensor.view(N, -1, K, H, W)
    tensor = tensor.permute(0, 3, 4, 1, 2)
    tensor = tensor.reshape(N, -1, K)  # Size=(N,HWA,K)
    return tensor

class DenseDetector(nn.Module):
    """
    Base class for dense detector. We define a dense detector as a fully-convolutional model that
    makes per-pixel (i.e. dense) predictions.
    """

    def __init__(
        self,
        backbone: Backbone,
        head: nn.Module,
        head_in_features: Optional[List[str]] = None,
        *,
        pixel_mean,
        pixel_std,
    ):
        """
        Args:
            backbone: backbone module
            head: head module
            head_in_features: backbone features to use in head. Default to all backbone features.
            pixel_mean (Tuple[float]):
                Values to be used for image normalization (BGR order).
                To train on images of different number of channels, set different mean & std.
                Default values are the mean pixel value from ImageNet: [103.53, 116.28, 123.675]
            pixel_std (Tuple[float]):
                When using pre-trained models in Detectron1 or any MSRA models,
                std has been absorbed into its conv1 weights, so the std needs to be set 1.
                Otherwise, you can use [57.375, 57.120, 58.395] (ImageNet std)
        """
        super().__init__()

        self.backbone = backbone
        self.head = head
        if head_in_features is None:
            shapes = self.backbone.output_shape()
            self.head_in_features = sorted(shapes.keys(), key=lambda x: shapes[x].stride)
        else:
            self.head_in_features = head_in_features

        self.register_buffer("pixel_mean", torch.tensor(pixel_mean).view(-1, 1, 1), False)
        self.register_buffer("pixel_std", torch.tensor(pixel_std).view(-1, 1, 1), False)

    @property
    def device(self):
        return self.pixel_mean.device

    def forward(self, batched_inputs: List[Dict[str, Tensor]]):
        """
        Args:
            batched_inputs: a list, batched outputs of :class:`DatasetMapper` .
                Each item in the list contains the inputs for one image.
                For now, each item in the list is a dict that contains:

                * image: Tensor, image in (C, H, W) format.
                * instances: Instances

                Other information that's included in the original dicts, such as:

                * "height", "width" (int): the output resolution of the model, used in inference.
                  See :meth:`postprocess` for details.

        Returns:
            In training, dict[str, Tensor]: mapping from a named loss to a tensor storing the
            loss. Used during training only. In inference, the standard output format, described
            in :doc:`/tutorials/models`.
        """
        images = self.preprocess_image(batched_inputs)
        features = self.backbone(images.tensor)
        features = [features[f] for f in self.head_in_features]
        predictions = self.head(features)

        if self.training:
            assert not torch.jit.is_scripting(), "Not supported"
            assert "instances" in batched_inputs[0], "Instance annotations are missing in training!"
            gt_instances = [x["instances"].to(self.device) for x in batched_inputs]
            return self.forward_training(images, features, predictions, gt_instances)
        else:
            results = self.forward_inference(images, features, predictions)
            if torch.jit.is_scripting():
                return results

            processed_results = []
            for results_per_image, input_per_image, image_size in zip(
                results, batched_inputs, images.image_sizes
            ):
                height = input_per_image.get("height", image_size[0])
                width = input_per_image.get("width", image_size[1])
                r = detector_postprocess(results_per_image, height, width)
                processed_results.append({"instances": r})
            return processed_results

    def forward_training(self, images, features, predictions, gt_instances):
        raise NotImplementedError()

    def preprocess_image(self, batched_inputs: List[Dict[str, Tensor]]):
        """
        Normalize, pad and batch the input images.
        """
        images = [x["image"].to(self.device) for x in batched_inputs]
        images = [(x - self.pixel_mean) / self.pixel_std for x in images]
        images = ImageList.from_tensors(images, self.backbone.size_divisibility)
        return images

    def _transpose_dense_predictions(
        self, predictions: List[List[Tensor]], dims_per_anchor: List[int]
    ) -> List[List[Tensor]]:
        """
        Transpose the dense per-level predictions.

        Args:
            predictions: a list of outputs, each is a list of per-level
                predictions with shape (N, Ai x K, Hi, Wi), where N is the
                number of images, Ai is the number of anchors per location on
                level i, K is the dimension of predictions per anchor.
            dims_per_anchor: the value of K for each predictions. e.g. 4 for
                box prediction, #classes for classification prediction.

        Returns:
            List[List[Tensor]]: each prediction is transposed to (N, Hi x Wi x Ai, K).
        """
        assert len(predictions) == len(dims_per_anchor)
        res: List[List[Tensor]] = []
        for pred, dim_per_anchor in zip(predictions, dims_per_anchor):
            pred = [permute_to_N_HWA_K(x, dim_per_anchor) for x in pred]
            res.append(pred)
        return res

    def _ema_update(self, name: str, value: float, initial_value: float, momentum: float = 0.9):
        """
        Apply EMA update to `self.name` using `value`.

        This is mainly used for loss normalizer. In Detectron1, loss is normalized by number
        of foreground samples in the batch. When batch size is 1 per GPU, #foreground has a
        large variance and using it lead to lower performance. Therefore we maintain an EMA of
        #foreground to stabilize the normalizer.

        Args:
            name: name of the normalizer
            value: the new value to update
            initial_value: the initial value to start with
            momentum: momentum of EMA

        Returns:
            float: the updated EMA value
        """
        if hasattr(self, name):
            old = getattr(self, name)
        else:
            old = initial_value
        new = old * momentum + value * (1 - momentum)
        setattr(self, name, new)
        return new

    def _decode_per_level_predictions(
        self,
        anchors: Boxes,
        pred_scores: Tensor,
        pred_deltas: Tensor,
        score_thresh: float,
        topk_candidates: int,
        image_size: Tuple[int, int],
    ) -> Instances:
        """
        Decode boxes and classification predictions of one featuer level, by
        the following steps:
        1. filter the predictions based on score threshold and top K scores.
        2. transform the box regression outputs
        3. return the predicted scores, classes and boxes

        Args:
            anchors: Boxes, anchor for this feature level
            pred_scores: HxWxA,K
            pred_deltas: HxWxA,4

        Returns:
            Instances: with field "scores", "pred_boxes", "pred_classes".
        """
        # Apply two filtering to make NMS faster.
        # 1. Keep boxes with confidence score higher than threshold
        keep_idxs = pred_scores > score_thresh
        pred_scores = pred_scores[keep_idxs]
        topk_idxs = torch.nonzero(keep_idxs)  # Kx2

        # 2. Keep top k top scoring boxes only
        num_topk = min(topk_candidates, topk_idxs.size(0))
        pred_scores, idxs = pred_scores.topk(num_topk)
        topk_idxs = topk_idxs[idxs]

        anchor_idxs, classes_idxs = topk_idxs.unbind(dim=1)

        pred_boxes = self.box2box_transform.apply_deltas(
            pred_deltas[anchor_idxs], anchors.tensor[anchor_idxs]
        )
        return Instances(
            image_size, pred_boxes=Boxes(pred_boxes), scores=pred_scores, pred_classes=classes_idxs
        )

    def _decode_multi_level_predictions(
        self,
        anchors: List[Boxes],
        pred_scores: List[Tensor],
        pred_deltas: List[Tensor],
        score_thresh: float,
        topk_candidates: int,
        image_size: Tuple[int, int],
    ) -> Instances:
        """
        Run `_decode_per_level_predictions` for all feature levels and concat the results.
        """
        predictions = [
            self._decode_per_level_predictions(
                anchors_i,
                box_cls_i,
                box_reg_i,
                self.test_score_thresh,
                self.test_topk_candidates,
                image_size,
            )
            # Iterate over every feature level
            for box_cls_i, box_reg_i, anchors_i in zip(pred_scores, pred_deltas, anchors)
        ]
        return predictions[0].cat(predictions)  # 'Instances.cat' is not scriptale but this is

    def visualize_training(self, batched_inputs, results):
        """
        A function used to visualize ground truth images and final network predictions.
        It shows ground truth bounding boxes on the original image and up to 20
        predicted object bounding boxes on the original image.

        Args:
            batched_inputs (list): a list that contains input to the model.
            results (List[Instances]): a list of #images elements returned by forward_inference().
        """
        from detectron2.utils.visualizer import Visualizer

        assert len(batched_inputs) == len(
            results
        ), "Cannot visualize inputs and results of different sizes"
        storage = get_event_storage()
        max_boxes = 20

        image_index = 0  # only visualize a single image
        img = batched_inputs[image_index]["image"]
        img = convert_image_to_rgb(img.permute(1, 2, 0), self.input_format)
        v_gt = Visualizer(img, None)
        v_gt = v_gt.overlay_instances(boxes=batched_inputs[image_index]["instances"].gt_boxes)
        anno_img = v_gt.get_image()
        processed_results = detector_postprocess(results[image_index], img.shape[0], img.shape[1])
        predicted_boxes = processed_results.pred_boxes.tensor.detach().cpu().numpy()

        v_pred = Visualizer(img, None)
        v_pred = v_pred.overlay_instances(boxes=predicted_boxes[0:max_boxes])
        prop_img = v_pred.get_image()
        vis_img = np.vstack((anno_img, prop_img))
        vis_img = vis_img.transpose(2, 0, 1)
        vis_name = f"Top: GT bounding boxes; Bottom: {max_boxes} Highest Scoring Results"
        storage.put_image(vis_name, vis_img)


import tensorflow as tf
from keras.preprocessing import image
import numpy as np

img_width,img_height=125,125

model=tf.keras.models.load_model(r'your path to the savedmodel')

sample='your path to testing image'
print(sample)
#0 label for namonia pics and 1 label for nromal pic
test_img=image.load_img(sample,target_size=(img_width,img_height))
test_img=image.img_to_array(test_img)
test_img=np.expand_dims(test_img,axis=0)
#test_img=test_img/255
result=model.predict(test_img)
print(result)


import os
import numpy as np
from tqdm import tqdm
from pymongo import MongoClient

def create_temporary_data():
    db = MongoClient(os.environ['COMTRIS_MONGODB_URI'])['COMTRIS']

    for i in tqdm(range(400000)):
        document = {}

        performance_list = ["low", "mid", "high"]
        performance = np.random.choice(performance_list)

        # CPU
        cpu_dict = {
            "low": ["i3", "r3"],
            "mid": ["i5", "r5"],
            "high": ["i7", "r7"]
        }
        cpu = np.random.choice(cpu_dict[performance])

        # VGA
        vga_dict = {
            "low": ["GTX 1050", "GTX 1050TI"],
            "mid": ["GTX 1060", "GTX 1070"],
            "high": ["GTX 1080", "GTX 1080TI"]
        }
        vga = np.random.choice(vga_dict[performance])

        # RAM
        ram_dict = {
            "low": ["4GB", "8GB"],
            "mid": ["8GB", "16GB"],
            "high": ["16GB", "32GB"]
        }
        ram = np.random.choice(ram_dict[performance])

        # M/B
        intel_dict = {
            "low": "H410",
            "mid": "B460",
            "high": "Z490"
        }
        amd_dict = {
            "low": "A520",
            "mid": "B550",
            "high": "X570",
        }
        if cpu[0] == 'i':
            mb = intel_dict[performance]
        else:
            mb = amd_dict[performance]

        # SSD
        ssd_dict = {
            "low": "256GB",
            "mid": "512GB",
            "high": "1TB"
        }
        ssd = ssd_dict[performance]

        # POWER
        power_dict = {
            "low": "500W",
            "mid": "600W",
            "high": "700W"
        }
        power = power_dict[performance]

        pc = {
            "CPU": cpu,
            "VGA": vga,
            "RAM": ram,
            "M/B": mb,
            "SSD": ssd,
            "POWER": power,
            "performance": performance
        }

        db['temporary_data'].insert_one(pc)

if __name__ == "__main__":
    create_temporary_data()


# -*- coding: utf-8 -*-
from PIL import Image
import pytesseract
import cv2
import numpy as np
import re
from multiprocessing import Pool, Queue, Lock, Process, freeze_support
import time

#pytesseract.pytesseract.tesseract_cmd = 'C:\\Program Files (x86)\\Tesseract-OCR\\tesseract.exe'
x = 1280.00 / 3840.00
pixel_x = int(x * 3840)
print(x, pixel_x)

#mode0:识别姓名，出生日期，身份证号； mode1：识别所有信息
def idcardocr(imgname, mode=1):
        print(u'进入身份证光学识别流程...')
        if mode==1:
            # generate_mask(x)
            img_data_gray, img_org = img_resize_gray(imgname)
            result_dict = dict()
            name_pic = find_name(img_data_gray, img_org)
            # showimg(name_pic)
            # print 'name'
            result_dict['name'] = get_name(name_pic)
            # print result_dict['name']

            sex_pic = find_sex(img_data_gray, img_org)
            # showimg(sex_pic)
            # print 'sex'
            result_dict['sex'] = get_sex(sex_pic)
            # print result_dict['sex']

            nation_pic = find_nation(img_data_gray, img_org)
            # showimg(nation_pic)
            # print 'nation'
            result_dict['nation'] = get_nation(nation_pic)
            # print result_dict['nation']

            address_pic = find_address(img_data_gray, img_org)
            # showimg(address_pic)
            # print 'address'
            result_dict['address'] = get_address(address_pic)
            # print result_dict['address']

            idnum_pic = find_idnum(img_data_gray, img_org)
            # showimg(idnum_pic)
            # print 'idnum'
            result_dict['idnum'], result_dict['birth'] = get_idnum_and_birth(idnum_pic)
            # print result_dict['idnum']
        elif mode==0:
            # generate_mask(x)
            img_data_gray, img_org = img_resize_gray(imgname)
            result_dict = dict()
            name_pic = find_name(img_data_gray, img_org)
            # showimg(name_pic)
            # print 'name'
            result_dict['name'] = get_name(name_pic)
            # print result_dict['name']

            idnum_pic = find_idnum(img_data_gray, img_org)
            # showimg(idnum_pic)
            # print 'idnum'

            result_dict['idnum'], result_dict['birth'] = get_idnum_and_birth(idnum_pic)
            result_dict['sex']=''
            result_dict['nation']=''
            result_dict['address']=''

        else:
            print(u"模式选择错误！")

        #showimg(img_data_gray)
        return result_dict

def generate_mask(x):
        name_mask_pic = cv2.UMat(cv2.imread('name_mask.jpg'))
        sex_mask_pic = cv2.UMat(cv2.imread('sex_mask.jpg'))
        nation_mask_pic = cv2.UMat(cv2.imread('nation_mask.jpg'))
        birth_mask_pic = cv2.UMat(cv2.imread('birth_mask.jpg'))
        year_mask_pic = cv2.UMat(cv2.imread('year_mask.jpg'))
        month_mask_pic = cv2.UMat(cv2.imread('month_mask.jpg'))
        day_mask_pic = cv2.UMat(cv2.imread('day_mask.jpg'))
        address_mask_pic = cv2.UMat(cv2.imread('address_mask.jpg'))
        idnum_mask_pic = cv2.UMat(cv2.imread('idnum_mask.jpg'))
        name_mask_pic = img_resize_x(name_mask_pic)
        sex_mask_pic = img_resize_x(sex_mask_pic)
        nation_mask_pic = img_resize_x(nation_mask_pic)
        birth_mask_pic = img_resize_x(birth_mask_pic)
        year_mask_pic = img_resize_x(year_mask_pic)
        month_mask_pic = img_resize_x(month_mask_pic)
        day_mask_pic = img_resize_x(day_mask_pic)
        address_mask_pic = img_resize_x(address_mask_pic)
        idnum_mask_pic = img_resize_x(idnum_mask_pic)
        cv2.imwrite('name_mask_%s.jpg'%pixel_x, name_mask_pic)
        cv2.imwrite('sex_mask_%s.jpg' %pixel_x, sex_mask_pic)
        cv2.imwrite('nation_mask_%s.jpg' %pixel_x, nation_mask_pic)
        cv2.imwrite('birth_mask_%s.jpg' %pixel_x, birth_mask_pic)
        cv2.imwrite('year_mask_%s.jpg' % pixel_x, year_mask_pic)
        cv2.imwrite('month_mask_%s.jpg' % pixel_x, month_mask_pic)
        cv2.imwrite('day_mask_%s.jpg' % pixel_x, day_mask_pic)
        cv2.imwrite('address_mask_%s.jpg' %pixel_x, address_mask_pic)
        cv2.imwrite('idnum_mask_%s.jpg' %pixel_x, idnum_mask_pic)

#用于生成模板
def img_resize_x(imggray):
        # print 'dheight:%s' % dheight
        crop = imggray
        size = crop.get().shape
        dheight = int(size[0]*x)
        dwidth = int(size[1]*x)
        crop = cv2.resize(src=crop, dsize=(dwidth, dheight), interpolation=cv2.INTER_CUBIC)
        return crop

#idcardocr里面resize以高度为依据, 用于get部分
def img_resize(imggray, dheight):
        # print 'dheight:%s' % dheight
        crop = imggray
        size = crop.get().shape
        height = size[0]
        width = size[1]
        width = width * dheight / height
        crop = cv2.resize(src=crop, dsize=(int(width), dheight), interpolation=cv2.INTER_CUBIC)
        return crop

def img_resize_gray(imgorg):

        #imgorg = cv2.imread(imgname)
        crop = imgorg
        size = cv2.UMat.get(crop).shape
        # print size
        height = size[0]
        width = size[1]
        # 参数是根据3840调的
        height = int(height * 3840 * x / width)
        # print height
        crop = cv2.resize(src=crop, dsize=(int(3840 * x), height), interpolation=cv2.INTER_CUBIC)
        return hist_equal(cv2.cvtColor(crop, cv2.COLOR_BGR2GRAY)), crop

def find_name(crop_gray, crop_org):
        template = cv2.UMat(cv2.imread('name_mask_%s.jpg'%pixel_x, 0))
        # showimg(crop_org)
        w, h = cv2.UMat.get(template).shape[::-1]
        res = cv2.matchTemplate(crop_gray, template, cv2.TM_CCOEFF_NORMED)
        min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(res)
        # print(max_loc)
        top_left = (max_loc[0] + w, max_loc[1] - int(20*x))
        bottom_right = (top_left[0] + int(700*x), top_left[1] + int(300*x))
        result = cv2.UMat.get(crop_org)[top_left[1]:bottom_right[1], top_left[0]:bottom_right[0]]
        cv2.rectangle(crop_gray, top_left, bottom_right, 255, 2)
        # showimg(result)
        return cv2.UMat(result)

def find_sex(crop_gray, crop_org):
        template = cv2.UMat(cv2.imread('sex_mask_%s.jpg'%pixel_x, 0))
        # showimg(template)
        w, h = cv2.UMat.get(template).shape[::-1]
        res = cv2.matchTemplate(crop_gray, template, cv2.TM_CCOEFF_NORMED)
        min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(res)
        top_left = (max_loc[0] + w, max_loc[1] - int(20*x))
        bottom_right = (top_left[0] + int(300*x), top_left[1] + int(300*x))
        result = cv2.UMat.get(crop_org)[top_left[1]:bottom_right[1], top_left[0]:bottom_right[0]]
        cv2.rectangle(crop_gray, top_left, bottom_right, 255, 2)
        #showimg(crop_gray)
        return cv2.UMat(result)

def find_nation(crop_gray, crop_org):
        template = cv2.UMat(cv2.imread('nation_mask_%s.jpg'%pixel_x, 0))
        #showimg(template)
        w, h = cv2.UMat.get(template).shape[::-1]
        res = cv2.matchTemplate(crop_gray, template, cv2.TM_CCOEFF_NORMED)
        min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(res)
        top_left = (max_loc[0] + w - int(20*x), max_loc[1] - int(20*x))
        bottom_right = (top_left[0] + int(500*x), top_left[1] + int(300*x))
        result = cv2.UMat.get(crop_org)[top_left[1]:bottom_right[1], top_left[0]:bottom_right[0]]
        cv2.rectangle(crop_gray, top_left, bottom_right, 255, 2)
        #showimg(crop_gray)
        return cv2.UMat(result)

# def find_birth(crop_gray, crop_org):
#         template = cv2.UMat(cv2.imread('birth_mask_%s.jpg'%pixel_x, 0))
#         # showimg(template)
#         w, h = cv2.UMat.get(template).shape[::-1]
#         res = cv2.matchTemplate(crop_gray, template, cv2.TM_CCOEFF_NORMED)
#         #showimg(crop_gray)
#         min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(res)
#         top_left = (max_loc[0] + w, max_loc[1] - int(20*x))
#         bottom_right = (top_left[0] + int(1500*x), top_left[1] + int(300*x))
#         # 提取result需要在rectangle之前
#         date_org = cv2.UMat.get(crop_org)[top_left[1]:bottom_right[1], top_left[0]:bottom_right[0]]
#         date = cv2.cvtColor(date_org, cv2.COLOR_BGR2GRAY)
#         cv2.rectangle(crop_gray, top_left, bottom_right, 255, 2)
#         # cv2.imwrite('date.png',date)
#
#         # 提取年份
#         template = cv2.UMat(cv2.imread('year_mask_%s.jpg'%pixel_x, 0))
#         year_res = cv2.matchTemplate(date, template, cv2.TM_CCOEFF_NORMED)
#         min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(year_res)
#         bottom_right = (max_loc[0]+int(20*x), int(300*x))
#         top_left = (0, 0)
#         year = date_org[top_left[1]:bottom_right[1], top_left[0]:bottom_right[0]]
#         # cv2.imwrite('year.png',year)
#         cv2.rectangle(crop_gray, top_left, bottom_right, 255, 2)
#
#         # 提取月
#         template = cv2.UMat(cv2.imread('month_mask_%s.jpg'%pixel_x, 0))
#         month_res = cv2.matchTemplate(date, template, cv2.TM_CCOEFF_NORMED)
#         min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(month_res)
#         bottom_right = (max_loc[0]+int(40*x), int(300*x))
#         top_left = (max_loc[0] - int(220*x), 0)
#         month = date_org[top_left[1]:bottom_right[1], top_left[0]:bottom_right[0]]
#         # cv2.imwrite('month.png',month)
#         cv2.rectangle(crop_gray, top_left, bottom_right, 255, 2)
#
#         # 提取日
#         template = cv2.UMat(cv2.imread('day_mask_%s.jpg'%pixel_x, 0))
#         day_res = cv2.matchTemplate(date, template, cv2.TM_CCOEFF_NORMED)
#         min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(day_res)
#         bottom_right = (max_loc[0]+int(20*x), int(300*x))
#         top_left = (max_loc[0] - int(220*x), 0)
#         day = date_org[top_left[1]:bottom_right[1], top_left[0]:bottom_right[0]]
#         # cv2.imwrite('day.png',day)
#         cv2.rectangle(crop_gray, top_left, bottom_right, 255, 2)
#         showimg(crop_gray)
#         return cv2.UMat(year), cv2.UMat(month), cv2.UMat(day)

def find_address(crop_gray, crop_org):
        template = cv2.UMat(cv2.imread('address_mask_%s.jpg'%pixel_x, 0))
        # showimg(template)
        #showimg(crop_gray)
        w, h = cv2.UMat.get(template).shape[::-1]
        #t1 = round(time.time()*1000)
        res = cv2.matchTemplate(crop_gray, template, cv2.TM_CCOEFF_NORMED)
        #t2 = round(time.time()*1000)
        #print 'time:%s'%(t2-t1)
        min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(res)
        top_left = (max_loc[0] + w, max_loc[1] - int(20*x))
        bottom_right = (top_left[0] + int(1700*x), top_left[1] + int(550*x))
        result = cv2.UMat.get(crop_org)[top_left[1]:bottom_right[1], top_left[0]:bottom_right[0]]
        cv2.rectangle(crop_gray, top_left, bottom_right, 255, 2)
        #showimg(crop_gray)
        return cv2.UMat(result)

def find_idnum(crop_gray, crop_org):
        template = cv2.UMat(cv2.imread('idnum_mask_%s.jpg'%pixel_x, 0))
        # showimg(template)
        #showimg(crop_gray)
        w, h = cv2.UMat.get(template).shape[::-1]
        res = cv2.matchTemplate(crop_gray, template, cv2.TM_CCOEFF_NORMED)
        min_val, max_val, min_loc, max_loc = cv2.minMaxLoc(res)
        top_left = (max_loc[0] + w, max_loc[1] - int(20*x))
        bottom_right = (top_left[0] + int(2300*x), top_left[1] + int(300*x))
        result = cv2.UMat.get(crop_org)[top_left[1]:bottom_right[1], top_left[0]:bottom_right[0]]
        cv2.rectangle(crop_gray, top_left, bottom_right, 255, 2)
        #showimg(crop_gray)
        return cv2.UMat(result)

def showimg(img):
        cv2.namedWindow("contours", 0);
        cv2.resizeWindow("contours", 1280, 720);
        cv2.imshow("contours", img)
        cv2.waitKey()

#psm model:
#  0    Orientation and script detection (OSD) only.
#  1    Automatic page segmentation with OSD.
#  2    Automatic page segmentation, but no OSD, or OCR.
#  3    Fully automatic page segmentation, but no OSD. (Default)
#  4    Assume a single column of text of variable sizes.
#  5    Assume a single uniform block of vertically aligned text.
#  6    Assume a single uniform block of text.
#  7    Treat the image as a single text line.
#  8    Treat the image as a single word.
#  9    Treat the image as a single word in a circle.
#  10    Treat the image as a single character.
#  11    Sparse text. Find as much text as possible in no particular order.
#  12    Sparse text with OSD.
#  13    Raw line. Treat the image as a single text line,
# 			bypassing hacks that are Tesseract-specific

def get_name(img):
        #    cv2.imshow("method3", img)
        #    cv2.waitKey()
        _, _, red = cv2.split(img) #split 会自动将UMat转换回Mat
        red = cv2.UMat(red)
        red = hist_equal(red)
        red = cv2.adaptiveThreshold(red, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, 151, 50)
        #    red = cv2.medianBlur(red, 3)
        red = img_resize(red, 150)
        img = img_resize(img, 150)
        # showimg(red)
        # cv2.imwrite('name.png', red)
        #    img2 = Image.open('address.png')
        # img = Image.fromarray(cv2.UMat.get(red).astype('uint8'))
        return get_result_vary_length(red, 'chi_sim', img, '-psm 7')
        # return punc_filter(pytesseract.image_to_string(img, lang='chi_sim', config='-psm 13').replace(" ",""))

def get_sex(img):
        _, _, red = cv2.split(img)
        red = cv2.UMat(red)
        red = hist_equal(red)
        red = cv2.adaptiveThreshold(red, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 151, 50)
        #    red = cv2.medianBlur(red, 3)
        #    cv2.imwrite('address.png', img)
        #    img2 = Image.open('address.png')
        red = img_resize(red, 150)
        # cv2.imwrite('sex.png', red)
        # img = Image.fromarray(cv2.UMat.get(red).astype('uint8'))
        return get_result_fix_length(red, 1, 'sex', '-psm 10')
        # return pytesseract.image_to_string(img, lang='sex', config='-psm 10').replace(" ","")

def get_nation(img):
        _, _, red = cv2.split(img)
        red = cv2.UMat(red)
        red = hist_equal(red)
        red = cv2.adaptiveThreshold(red, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 151, 50)
        red = img_resize(red, 150)
        # cv2.imwrite('nation.png', red)
        # img = Image.fromarray(cv2.UMat.get(red).astype('uint8'))
        return get_result_fix_length(red, 1, 'nation', '-psm 10')
        # return pytesseract.image_to_string(img, lang='nation', config='-psm 13').replace(" ","")

# def get_birth(year, month, day):
#         _, _, red = cv2.split(year)
#         red = cv2.UMat(red)
#         red = hist_equal(red)
#         red = cv2.adaptiveThreshold(red, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 151, 50)
#         red = img_resize(red, 150)
#         # cv2.imwrite('year_red.png', red)
#         year_red = red
#
#         _, _, red = cv2.split(month)
#         red = cv2.UMat(red)
#         red = hist_equal(red)
#         red = cv2.adaptiveThreshold(red, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 151, 50)
#         #red = cv2.erode(red,kernel,iterations = 1)
#         red = img_resize(red, 150)
#         # cv2.imwrite('month_red.png', red)
#         month_red = red
#
#         _, _, red = cv2.split(day)
#         red = cv2.UMat(red)
#         red = hist_equal(red)
#         red = cv2.adaptiveThreshold(red, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 151, 50)
#         red = img_resize(red, 150)
#         # cv2.imwrite('day_red.png', red)
#         day_red = red
#         # return pytesseract.image_to_string(img, lang='birth', config='-psm 7')
#         return get_result_fix_length(year_red, 4, 'eng', '-c tessedit_char_whitelist=0123456789 -psm 13'), \
#                get_result_vary_length(month_red, 'eng', '-c tessedit_char_whitelist=0123456789 -psm 13'), \
#                get_result_vary_length(day_red, 'eng', '-c tessedit_char_whitelist=0123456789 -psm 13')

def get_address(img):
        #_, _, red = cv2.split(img)
        #red = cv2.medianBlur(red, 3)
        _, _, red = cv2.split(img)
        red = cv2.UMat(red)
        red = hist_equal(red)
        red = cv2.adaptiveThreshold(red, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 151, 50)
        red = img_resize(red, 300)
        img = img_resize(img, 300)
        # cv2.imwrite('address_red.png', red)
        #img = Image.fromarray(cv2.UMat.get(red).astype('uint8'))
        return punc_filter(get_result_vary_length(red,'chi_sim', img, '-psm 6'))
        #return punc_filter(pytesseract.image_to_string(img, lang='chi_sim', config='-psm 3').replace(" ",""))

def get_idnum_and_birth(img):
        _, _, red = cv2.split(img)
        red = cv2.UMat(red)
        red = hist_equal(red)
        red = cv2.adaptiveThreshold(red, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY, 151, 50)
        red = img_resize(red, 150)
        # cv2.imwrite('idnum_red.png', red)
        idnum_str = get_result_fix_length(red, 18, 'idnum', '-psm 8')
        return idnum_str, idnum_str[6:14]

def get_result_fix_length(red, fix_length, langset, custom_config=''):
    cv2.fastNlMeansDenoising(red, red, 4, 7, 35)
    rec, red = cv2.threshold(red, 127, 255, cv2.THRESH_BINARY_INV)
    image, contours, hierarchy = cv2.findContours(red, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    # print(len(contours))
    # 描边一次可以减少噪点
    cv2.drawContours(red, contours, -1, (0, 255, 0), 1)
    color_img = cv2.cvtColor(red, cv2.COLOR_GRAY2BGR)
    # for x, y, w, h in contours:
    #     imgrect = cv2.rectangle(color_img, (x, y), (x + w, y + h), (0, 255, 0), 2)
    # showimg(imgrect)

    h_threshold = 54
    numset_contours = []
    calcu_cnt = 1
    for cnt in contours:
        x, y, w, h = cv2.boundingRect(cnt)
        if h > h_threshold:
            numset_contours.append((x, y, w, h))
    while len(numset_contours) != fix_length:
        if calcu_cnt > 50:
            print(u'计算次数过多！目前阈值为：', h_threshold)
            break
        numset_contours = []
        calcu_cnt += 1
        if len(numset_contours) > fix_length:
            h_threshold += 1
            contours_cnt = 0
            for cnt in contours:
                x, y, w, h = cv2.boundingRect(cnt)
                if h > h_threshold:
                    contours_cnt += 1
                    numset_contours.append((x, y, w, h))
        if len(numset_contours) < fix_length:
            h_threshold -= 1
            contours_cnt = 0
            for cnt in contours:
                x, y, w, h = cv2.boundingRect(cnt)
                if h > h_threshold:
                    contours_cnt += 1
                    numset_contours.append((x, y, w, h))
    result_string = ''
    numset_contours.sort(key=lambda num: num[0])
    for x, y, w, h in numset_contours:
        result_string += pytesseract.image_to_string(cv2.UMat.get(color_img)[y:y + h, x:x + w], lang=langset, config=custom_config)
    # print(new_r)
    return result_string

def get_result_vary_length(red, langset, org_img, custom_config=''):
    # cv2.fastNlMeansDenoising(red, red, 4, 7, 35)
    rec, red = cv2.threshold(red, 127, 255, cv2.THRESH_BINARY_INV)
    image, contours, hierarchy = cv2.findContours(red, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
    # print(len(contours))
    # 描边一次可以减少噪点
    cv2.drawContours(red, contours, -1, (255, 255, 255), 1)
    color_img = cv2.cvtColor(red, cv2.COLOR_GRAY2BGR)
    numset_contours = []
    height_list=[]
    width_list=[]
    for cnt in contours:
        x, y, w, h = cv2.boundingRect(cnt)
        height_list.append(h)
        # print(h,w)
        width_list.append(w)
    height_list.remove(max(height_list))
    width_list.remove(max(width_list))
    height_threshold = 0.70*max(height_list)
    width_threshold = 1.4 * max(width_list)
    # print('height_threshold:'+str(height_threshold)+'width_threshold:'+str(width_threshold))
    big_rect=[]
    for cnt in contours:
        x, y, w, h = cv2.boundingRect(cnt)
        if h > height_threshold and w < width_threshold:
            # print(h,w)
            numset_contours.append((x, y, w, h))
            big_rect.append((x, y))
            big_rect.append((x + w, y + h))
    big_rect_nparray = np.array(big_rect, ndmin=3)
    x, y, w, h = cv2.boundingRect(big_rect_nparray)
    # imgrect = cv2.rectangle(color_img, (x, y), (x + w, y + h), (0, 255, 0), 2)
    # showimg(imgrect)
    # showimg(cv2.UMat.get(org_img)[y:y + h, x:x + w])

    result_string = ''
    result_string += pytesseract.image_to_string(cv2.UMat.get(org_img)[y:y + h, x:x + w], lang=langset,
                                                 config=custom_config)
    # numset_contours.sort(key=lambda num: num[0])
    # for x, y, w, h in numset_contours:
    #     result_string += pytesseract.image_to_string(cv2.UMat.get(color_img)[y:y + h, x:x + w], lang=langset, config=custom_config)
    return punc_filter(result_string)

def punc_filter(str):
        temp = str
        xx      =   u"([\u4e00-\u9fff0-9]+)"
        pattern =   re.compile(xx)
        results =   pattern.findall(temp)
        string = ""
        for result in results:
            string += result
        return string

#这里使用直方图拉伸，不是直方图均衡
def hist_equal(img):
        # clahe_size = 8
        # clahe = cv2.createCLAHE(clipLimit=1.0, tileGridSize=(clahe_size, clahe_size))
        # result = clahe.apply(img)
        #test

        #result = cv2.equalizeHist(img)

        image = img.get() #UMat to Mat
        # result = cv2.equalizeHist(image)
        lut = np.zeros(256, dtype = image.dtype )#创建空的查找表
        #lut = np.zeros(256)
        hist= cv2.calcHist([image], #计算图像的直方图
                            [0], #使用的通道
                            None, #没有使用mask
                           [256], #it is a 1D histogram
                           [0,256])
        minBinNo, maxBinNo = 0, 255
        #计算从左起第一个不为0的直方图柱的位置
        for binNo, binValue in enumerate(hist):
            if binValue != 0:
                minBinNo = binNo
                break
        #计算从右起第一个不为0的直方图柱的位置
        for binNo, binValue in enumerate(reversed(hist)):
            if binValue != 0:
                maxBinNo = 255-binNo
                break
        #print minBinNo, maxBinNo
        #生成查找表
        for i,v in enumerate(lut):
            if i < minBinNo:
                lut[i] = 0
            elif i > maxBinNo:
                lut[i] = 255
            else:
                lut[i] = int(255.0*(i-minBinNo)/(maxBinNo-minBinNo)+0.5)
        #计算,调用OpenCV cv2.LUT函数,参数 image --  输入图像，lut -- 查找表
        #print lut
        result = cv2.LUT(image, lut)
        #print type(result)
        #showimg(result)
        return cv2.UMat(result)

if __name__=="__main__":
    idocr = idcardocr(cv2.UMat(cv2.imread('testimages/zrh.jpg')))
    print(idocr)
    # for i in range(15):
    #     idocr = idcardocr(cv2.UMat(cv2.imread('testimages/%s.jpg'%(i+1))))
    #     print(idocr['idnum'])


#! /usr/bin/python3
# This code is being started from /etc/xdg/lxsession/LXDE-pi/autostart. Backup options was trough /home/pi/.bashrc
import time
import board
from busio import I2C
import adafruit_bme680
import adafruit_tsl2561
import numpy as np
import RPi.GPIO as GPIO
import cv2
import requests
import base64
import schedule
import os
import subprocess
from getmac import get_mac_address

URL = "insert url to the server here!"
code_version = "2"
framecounter = 0
GPIO.setup(16, GPIO.IN, pull_up_down = GPIO.PUD_DOWN)
input = GPIO.input(16)
i2c = I2C(board.SCL, board.SDA)
bme680 = adafruit_bme680.Adafruit_BME680_I2C(i2c, 0x76, debug=False)
tsl2561 = adafruit_tsl2561.TSL2561(i2c, 0x29)
tsl2561.gain = 1 # 1 = 16x
tsl2561.integration_time = 2 # 2 = 402ms, 1 = 101ms, 0 = 13.7ms
wlan_check_flag = False
mac_address = get_mac_address(interface="wlan0")

def wlan_check():
    global wlan_check_flag
    if wlan_check_flag:
        print("long term issue, reboot")
        subprocess.call(['logger "wlan down, forcing a reboot"'], shell=True)
        wlan_check_flag = False
        subprocess.call(['sudo reboot'], shell=True)
    else:
        print("trying to recover connection")
        subprocess.call(['logger "Wlan is down, try a reconnect"'], shell=True)
        wlan_check_flag = True
        subprocess.call(['sudo /sbin/ifdown wlan0 && sleep 10 && sudo/sbin/ifup --force wlan0'], shell=True)

def take_picture():
    print("smile!")
    cap = cv2.VideoCapture(0)

    for x in range(5):
        ret, img = cap.read()
        time.sleep(1)
    cap.release()
    #save latest high res image
    cv2.imwrite("/home/pi/Desktop/code/high_res.jpg", img)

    height = 75
    scale = img.shape[0]/height
    width = int(img.shape[1] / scale)
    dim = (width, height)
    img = cv2.resize(img, dim, interpolation = cv2.INTER_AREA)
    _, img_arr = cv2.imencode('.jpg', img)
    img_bytes = img_arr.tobytes()
    img_b64 = base64.b64encode(img_bytes)
    try:
        x = requests.post(URL, data={"device_ID":mac_address, "version": code_version, "img":img_b64})
        wlan_check_flag = False
        print(x.text)
    except requests.exceptions.ConnectionError as err:
        print ("request failed:", err)
        wlan_check()

def read_sensors():
    global framecounter
    print("starting measurements")
    counter = 0
    for x in range(55000): # this gives about 60sec measurement
        time.sleep(0.001)
        if GPIO.input(16) == 0:
            counter = counter+1
    percentage = counter / 550 # devide by the same 55000 and multiply by 100 to get %

    percentage = "null" if percentage is None else round(percentage,2)

    temp = bme680.temperature
    temp = "null" if temp is None else round(temp,2)
    humid = bme680.humidity
    humid = "null" if humid is None else round(humid,2)
    press = bme680.pressure
    press = "null" if press is None else round(press,2)
    voc = bme680.gas
    voc = "null" if voc is None else round(voc,2)
    lux = tsl2561.lux
    lux = "null" if lux is None else round(lux,2)
    framecounter = framecounter+1

    print("Measurements: ",temp, humid, press, voc, lux, percentage)
    try:
        x = requests.post(URL, data={"device_ID":mac_address, "version": code_version, "temperature":temp, "humidity":humid, "pressure": press, "voc": voc, "illumination": lux, "dust": percentage, "counter": framecounter})
        wlan_check_flag = False
        print(x.text)
    except requests.exceptions.ConnectionError as err:
        print ("request failed:", err)
        wlan_check()

schedule.every(5).minutes.do(read_sensors)
schedule.every().day.at("02:00").do(take_picture)
schedule.every().day.at("04:00").do(take_picture)

print("mumo logger code started! (device ID: "+mac_address+")")
time.sleep(60) # wait a minute for wifi to start up.
read_sensors()
take_picture()
while True:
    schedule.run_pending()
    time.sleep(20)

cap.release()
GPIO.cleanup()


import numpy as np
import socket
import struct
import os
import time
import random
from PIL import Image
from VECA.environment import GeneralEnvironment
from tasks.COGNIANav.utils import *
from tasks.COGNIANav.constants import AUDIO
from constants import RECORD
import pickle
from moviepy.editor import *
from moviepy.audio.AudioClip import AudioArrayClip

WAV_K = 10
class Environment(GeneralEnvironment):
    def __init__(self, num_envs, port = 8870):
        GeneralEnvironment.__init__(self, num_envs, port)
        self.name = 'COGNIANav'
        self.SIM = 'VECA'
        self.mode = 'CONT'
        self.observation_space = {
            'image': (6, 84, 84),
            #'audio': (2, 66 * 13),
            'audio': (2, 2940*WAV_K) if AUDIO == 'RAW' else (2, 66*13) if AUDIO == 'STFT' else (2, 1) ,
            'touch': 1, # Simple tactile
            #'obj': 3,
            'real_audio': (2, 2940),
            'real_touch': 16480,
            'Recimage': (3, 224, 224)
        }
        self.wav_buffer = np.zeros([WAV_K] + list(self.observation_space['real_audio']))
        self.VEC_OBJ = True
        #self.action_space = 2
        '''
        if VEC_OBJ == False:
            with open('navigation/data/data.pkl', 'rb') as F:
                self.obj_data = pickle.load(F)
        '''
        if RECORD:
            self.imgsRec = []
            self.wavsRec = []

    def collect_observations(self, ignore_agent_dim = True):
        data = super().collect_observations(ignore_agent_dim = True)
        rewards, done, info, Gobjs, Bobjs = [], [], [], [], []
        GposAL, BposAL, GposL, BposL = [], [], [], []
        objimgs, objwavs, objtouchs = [], [], []
        imgs, wavs, real_audios, touchs, real_touchs = [], [], [], [], []

        IMG_C, IMG_H, IMG_W = self.observation_space['image']
        WAV_C, _ = self.observation_space['audio']

        if RECORD:
            IMG_REC_C, IMG_REC_H, IMG_REC_W = self.observation_space['Recimage']
            imgRec = np.reshape(np.array(list(reversed(data['Recimg'][0]))), self.observation_space['Recimage']).astype(np.uint8)
            imgRec = np.transpose(imgRec, [1, 2, 0])
            wavRec = np.reshape(np.array(data['Recwav'][0]), [WAV_C, -1])
            self.imgsRec.append(imgRec)
            self.wavsRec.append(wavRec)

        #print(data.keys())
        mask = np.zeros(self.num_envs, dtype = np.bool)
        for i in range(self.num_envs):
            img = list(reversed(data['img'][i]))
            if 'wav' in data.keys():
                wav = data['wav'][i]
            else:
                wav = np.zeros(self.observation_space['real_audio'])
            touch = data['touch'][i]
            if 'real_touch' in data.keys():
                real_touch = data['real_touch'][i]
            else:
                real_touch = np.zeros([self.observation_space['real_touch']])
            doneA = data['done'][i][0]
            reward = data['reward'][i][0]
            GposA, BposA = data['goodpos'][i], data['badpos'][i]
            Gpos, Bpos = list(data['goodcampos'][i]), list(data['badcampos'][i])
            Gpos[0], Gpos[1], Gpos[2] = np.tanh(Gpos[0] - 0.5), Gpos[1], np.tanh(Gpos[2])
            Bpos[0], Bpos[1], Bpos[2] = np.tanh(Bpos[0] - 0.5), Bpos[1], np.tanh(Bpos[2])
            Gobj, Bobj = str(data['goodobj'][i]), str(data['badobj'][i])
            img = np.reshape(np.array(img), [IMG_C, IMG_H, IMG_W]) / 255.0
            wav = np.reshape(np.array(wav), [WAV_C, -1]) / 32768.0
            self.wav_buffer = np.concatenate([self.wav_buffer[1:], [wav]], axis = 0)
            objimg, objwav, objtouch = data['objimg'][i], data['objwav'][i], data['objtouch'][i]
            #print("objwav", objwav)
            #print("objwav", objwav)
            #objimg, objwav, objtouch = [0,0], 0, 0
            imgs.append(img)
            if AUDIO == 'STFT':
                wavs.append(wav2freqSTFT(self.wav_buffer))
            if AUDIO == 'RAW':
                wavs.append(np.reshape(np.transpose(self.wav_buffer, [1, 0, 2]), [2, -1]))
            if AUDIO == 'SIMPLE':
                wavs.append(objwav * np.ones([2, 1]))
            if objwav[0] == 1:
                objwav = [1, 0, 0]
            elif objwav[0] == -1:
                objwav = [0, 1, 0]
            else:
                objwav = [0, 0, 1]
            real_audios.append(wav)
            touchs.append(touch)
            real_touchs.append(real_touch)
            rewards.append(reward)
            GposAL.append(GposA)
            BposAL.append(BposA)
            GposL.append(Gpos)
            BposL.append(Bpos)
            Gobjs.append(Gobj)
            Bobjs.append(Bobj)
            objimgs.append(objimg)
            objwavs.append(objwav)
            objtouchs.append(objtouch)
            if doneA: done.append(True)
            else: done.append(False)
        self.reset(done)
        #info = (data['pos'], data['campos'])
        GposAL = np.array(GposAL)
        BposAL = np.array(BposAL)
        posAL = np.stack([GposAL, BposAL], axis = 1)
        GposL = np.array(GposL)
        BposL = np.array(BposL)
        posL = np.stack([GposL, BposL], axis = 1)
        objs = [Gobjs, Bobjs]
        objimgs = np.reshape(np.array(objimgs), [self.num_envs, 2])
        objwavs = np.reshape(np.array(objwavs), [self.num_envs, 3])
        objtouchs = np.reshape(np.array(objtouchs), [self.num_envs])
        #info = (posL, posAL, objs)

        imgs = np.array(imgs)
        wavs = np.array(wavs)
        touchs = np.array(touchs)
        real_touchs = np.array(real_touchs)
        obs = {'image': imgs, 'audio': wavs, 'touch': touchs, 'real_audio': real_audios, 'real_touch': real_touchs}
        infos = {'pos': np.stack([GposL, BposL], axis=1), 'objimg': objimgs, 'objwav': objwavs, 'objtouch': objtouchs}

        rewards = np.array(rewards)
        done = np.array(done)

        return (obs, rewards, done, infos)

    def send_action(self, action):
        super().send_action(action)

    def step(self, action, ignore_agent_dim = True):
        self.send_action(action)
        return self.collect_observations(ignore_agent_dim)

    def reset(self, mask = None):
        self.wav_buffer = np.zeros_like(self.wav_buffer)
        super().reset(mask)

    def write_record(self):
        image_clip = ImageSequenceClip(self.imgsRec, fps=15)

        print(self.imgsRec[0].shape, self.imgsRec[0].dtype, self.imgsRec[0].max())
        audios = np.concatenate(self.wavsRec, axis = 1)
        audios = (np.transpose(audios) / 32768.).astype(np.float32)
        audioclip = AudioArrayClip(audios, fps=44100)
        print(audios.shape, audios.dtype, audios.max())

        print(len(self.imgsRec) / 15, audios.shape[0] / 44100)

        video_clip = image_clip.set_audio(audioclip)
        video_clip.write_videofile("result.mp4", fps=15, temp_audiofile="temp-audio.m4a", remove_temp=True, codec="libx264", audio_codec="aac")


#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time    : 2021/5/12 14:44
# @Author  : PPq
from __future__ import print_function
import argparse
import os
import shutil
import time
from torch.nn import DataParallel
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim
import torch.utils.data
import torchvision.transforms as transforms
import torchvision.datasets as datasets
from torch.utils.data import DataLoader
from core import IDR
from core.utils import init_log
from dataloader.CASIA_NIR_VIS_IDR import CASIA_NIR_VIS
import numpy as np
from load_imglist import ImageList

parser = argparse.ArgumentParser(description='PyTorch Light CNN Training')
parser.add_argument('--arch', '-a', metavar='ARCH', default='LightCNN')
parser.add_argument('--cuda', '-c', default=True)
parser.add_argument('-j', '--workers', default=4, type=int, metavar='N',
                    help='number of data loading workers (default: 16)')
parser.add_argument('--epochs', default=70, type=int, metavar='N',
                    help='number of total epochs to run')
parser.add_argument('--start-epoch', default=0, type=int, metavar='N',
                    help='manual epoch number (useful on restarts)')
parser.add_argument('-b', '--batch-size', default=128, type=int,
                    metavar='N', help='mini-batch size (default: 128)')
parser.add_argument('--lr', '--learning-rate', default=1e-2, type=float,
                    metavar='LR', help='initial learning rate')
parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
                    help='momentum')
parser.add_argument('--weight-decay', '--wd', default=1e-4, type=float,
                    metavar='W', help='weight decay (default: 1e-4)')
parser.add_argument('--print-freq', '-p', default=20, type=int,
                    metavar='N', help='print frequency (default: 100)')
parser.add_argument('--model', default='IDRnet', type=str, metavar='Model',
                    help='model type: IDRnet')
parser.add_argument('--resume', default='/root/NIR_VIS_Face_Recognition-master/model/LightCNN_9Layers_checkpoint.pth.tar', type=str, metavar='PATH',
                    help='path to latest checkpoint (default: none)')
parser.add_argument('--root_path', default='', type=str, metavar='PATH',
                    help='path to root path of images (default: none)')
parser.add_argument('--train_dir',
                    default='/root/NIR_VIS_Face_Recognition-master/dataset/CASIA_NIR_VIS_2.0/NIR-VIS-2.0', type=str,
                    metavar='PATH',
                    help='path to training list (default: none)')
parser.add_argument('--val_dir', default='/root/NIR_VIS_Face_Recognition-master/dataset/CASIA_NIR_VIS_2.0/NIR-VIS-2.0',
                    type=str, metavar='PATH',
                    help='path to validation list (default: none)')
parser.add_argument('--save_path', default='model/tmp/', type=str, metavar='PATH',
                    help='path to save checkpoint (default: none)')
parser.add_argument('--num_classes', default=839, type=int,
                    metavar='N', help='number of classes')

def main():
    global args
    args = parser.parse_args()

    # create Light CNN for face recognition
    if args.model == 'IDRnet':
        model = IDR.IDRnet(num_classes=args.num_classes)
    else:
        print('Error model type\n')

    if args.cuda:
        model = torch.nn.DataParallel(model).cuda()

    '''
    print(model.module)
    for name, value in model.named_parameters():
        print(name, value.shape)
    '''

    # large lr for last fc parameters
    basic_params = []
    for name, value in model.module.basic_layer.named_parameters():
        if 'bias' in name:
            basic_params += [{'params': value, 'lr': 2 * args.lr, 'weight_decay': 0}]  #2
        else:
            basic_params += [{'params': value, 'lr': 1 * args.lr}] #1

    basic_opt = torch.optim.SGD(basic_params, args.lr,
                                momentum=args.momentum,
                                weight_decay=args.weight_decay)
    # large lr for last fc parameters
    feat_params = []
    for name, value in model.module.feature_layer.named_parameters():
        if 'bias' in name:
            if 'fc' in name: #fc2.bias
                feat_params += [{'params': value, 'lr': 0.1 * args.lr, 'weight_decay': 0}]
            else: #bias
                feat_params += [{'params': value, 'lr': 0.1 * args.lr, 'weight_decay': 0}]
        else:
            if 'fc' in name: #fc.weight
                feat_params += [{'params': value, 'lr': 0.1 * args.lr}]
            else:  #weight
                feat_params += [{'params': value, 'lr': 0.1 * args.lr}]
    feat_opt = torch.optim.SGD(feat_params, args.lr, momentum = args.momentum, weight_decay = args.weight_decay)

    optimizer = [basic_opt, feat_opt]

    # optionally resume from a checkpoint
    if args.resume:
        if os.path.isfile(args.resume):
            print("=> loading checkpoint '{}'".format(args.resume))
            #args.start_epoch = checkpoint['epoch']
            args.start_epoch = 0
            checkpoint = torch.load(args.resume)
            if 'LightCNN9' in args.resume:
                print(here)
                del(checkpoint['state_dict']['module.fc1.filter.weight'])
                del(checkpoint['state_dict']['module.fc1.filter.bias'])
                del(checkpoint['state_dict']['module.fc2.weight'])
                del(checkpoint['state_dict']['module.fc2.bias'])
                name = [key[7:] for key in checkpoint['state_dict'].keys()]
                value = checkpoint['state_dict'].values()
                new_state_dict = dict(zip(name, value))
                #print(list(checkpoint['state_dict'].keys()))
                model_dict = model.module.basic_layer.state_dict()
                model_dict.update(new_state_dict)
                model.module.basic_layer.load_state_dict(model_dict)
            #print("=> loaded checkpoint '{}' (epoch {})".format(args.resume, checkpoint['epoch']))
        else:
            print("=> no checkpoint found at '{}'".format(args.resume))
    weight_init(model)
    cudnn.benchmark = True

    # load image
    all_transform = transforms.Compose([
        transforms.Grayscale(1),
        transforms.ToTensor(),
    ])
    # define trainloader and testloader
    trainset = CASIA_NIR_VIS(root=args.train_dir, transform=all_transform)
    train_loader = torch.utils.data.DataLoader(trainset, batch_size=args.batch_size,
                                               shuffle=True, num_workers=8, drop_last=False)

    valset = CASIA_NIR_VIS(root=args.val_dir, transform=all_transform, test=True)
    val_loader = torch.utils.data.DataLoader(valset, batch_size=32,
                                             shuffle=False, num_workers=8, drop_last=False)

    # define loss function and optimizer
    basic_criterion = nn.CrossEntropyLoss()
    total_critertion = IDRLoss()
    criterion = [basic_criterion, total_critertion]
    if args.cuda:
        for cri in criterion:
            cri.cuda()

    validate(val_loader, model, criterion)

    for epoch in range(args.start_epoch, args.epochs):
        if epoch == 50:
            print('Change lr')
            for param_group in basic_opt.param_groups:
                param_group['lr'] = 1e-5
            for param_group in feat_opt.param_groups:
                param_group['lr'] = 1e-5

        # train for one epoch
        train(train_loader, model, criterion, optimizer, epoch)

        # evaluate on validation set
        prec1 = validate(val_loader, model, criterion)

        save_name = args.save_path + 'lightCNN_' + str(epoch + 1) + '_checkpoint.pth.tar'
        save_checkpoint({
            'epoch': epoch + 1,
            'arch': args.arch,
            'state_dict': model.state_dict(),
            'prec1': prec1,
        }, save_name)

def train(train_loader, model, criterion, optimizer, epoch):
    batch_time = AverageMeter()
    data_time = AverageMeter()
    losses = AverageMeter()
    top1 = AverageMeter()
    top5 = AverageMeter()

    model.train()

    end = time.time()
    for i, (input, target, idt) in enumerate(train_loader):
        data_time.update(time.time() - end)
        input = input.cuda()
        # print(input.shape)
        target = target.cuda()
        idt = idt.cuda()
        input_var = torch.autograd.Variable(input)
        target_var = torch.autograd.Variable(target)
        idt_var = torch.autograd.Variable(idt)

        # compute output
        output, _ = model(input_var, idt_var)
        basic_loss = criterion[0](output, target_var)

        # compute gradient and do SGD step
        optimizer[0].zero_grad()
        basic_loss.backward()
        optimizer[0].step()

        pn = model.module.feature_layer.unique_mfm1.filter.weight
        pv = model.module.feature_layer.unique_mfm2.filter.weight
        w = model.module.feature_layer.shared_mfm.filter.weight
        output, _ = model(input_var, idt_var)
        loss = criterion[1](output, target_var, w, pn, pv)
        optimizer[1].zero_grad()
        basic_loss.backward()
        optimizer[1].step()

        # measure accuracy and record loss
        prec1, prec5 = accuracy(output.data, target, topk=(1, 5))
        # print('\nprec1:{}, prec5:{}'.format(prec1, prec5))
        losses.update(loss.data.item(), input.size(0))
        top1.update(prec1.item(), input.size(0))
        top5.update(prec5.item(), input.size(0))

        # measure elapsed time
        batch_time.update(time.time() - end)
        end = time.time()

        if i % args.print_freq == 0:
            print('Epoch: [{0}][{1}/{2}]\t'
                  'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
                  'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
                  'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
                  'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'
                  'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
                epoch, i, len(train_loader), batch_time=batch_time,
                data_time=data_time, loss=losses, top1=top1, top5=top5))

def validate(val_loader, model, criterion):
    batch_time = AverageMeter()
    losses = AverageMeter()
    top1 = AverageMeter()
    top5 = AverageMeter()

    # switch to evaluate mode
    model.eval()

    end = time.time()
    for i, (input, target, idt) in enumerate(val_loader):
        input = input.cuda()
        #print('input.shape : {}'.format(input.shape))
        target = target.cuda()
        idt = idt.cuda()
        input_var = input.clone().detach()
        target_var = target.clone().detach()
        idt_var = idt.clone().detach()

        # compute output
        output, _ = model(input_var, idt_var)
        #print('output.shape: {}'.format(output.shape))
        pn = model.module.feature_layer.unique_mfm1.filter.weight
        pv = model.module.feature_layer.unique_mfm2.filter.weight
        w = model.module.feature_layer.shared_mfm.filter.weight
        loss = criterion[1](output, target_var, w, pn, pv)

        # measure accuracy and record loss
        prec1, prec5 = accuracy(output.data, target, topk=(1, 5))
        # print('\nprec1:{}, prec5:{}'.format(prec1, prec5))
        losses.update(loss.data.item(), input.size(0))
        top1.update(prec1.item(), input.size(0))
        top5.update(prec5.item(), input.size(0))

    print('\ntop1.sum: {}, top1.count: {}'.format(top1.sum, top1.count))
    print('\nTest set: Average loss: {}, Accuracy: ({})\n'.format(losses.avg, top1.avg))

    return top1.avg

def save_checkpoint(state, filename):
    torch.save(state, filename)

class AverageMeter(object):
    """Computes and stores the average and current value"""

    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

def accuracy(output, target, topk=(1,)):
    """Computes the precision@k for the specified values of k"""
    maxk = max(topk)
    batch_size = target.size(0)

    # (batch, feature)，从每个feature选择topk
    _, pred = output.topk(maxk, 1, True, True)
    # 转置
    pred = pred.t()
    correct = pred.eq(target.view(1, -1).expand_as(pred))

    res = []
    for k in topk:
        # 计算有多少个预测正确
        correct_k = correct[:k].view(-1).float().sum(0)
        res.append(correct_k.mul_(100.0 / batch_size))
    return res

def weight_init(model):
    #print(type(model.named_parameters()))
    for name, param in model.named_parameters():
        if 'mfm' in name and 'weight' in name:
            print(name)
            nn.init.uniform_(param, -1.0/(param.size(0) ** 0.5), 1.0 / (param.size(0) ** 0.5))

class IDRLoss(nn.Module):
    def __init__(self, lamda=2e-4):
        super(IDRLoss, self).__init__()
        self.lamda = lamda
        self.ce_loss = nn.CrossEntropyLoss()

    def forward(self, x, target, w, pn, pv):
        loss1 = self.ce_loss(x, target)
        loss2 = self.lamda * (torch.norm(torch.mm(pn.t(), w), p='fro') ** 2) + \
                        self.lamda * (torch.norm(torch.mm(pv.t(), w), p = 'fro') ** 2)
        #print('loss1 : {}, loss2: {}'.format(loss1, loss2))
        return loss1 + loss2

if __name__ == '__main__':
    main()


import os
import time
import math
import pybullet as pb
import pybullet_data
import functools
import numpy as np

class Robotiq3FGripper:
    _PINCH_LIMIT = 175 # in pinch mode the fingers can only travel part way (until they touch).

    '''PyBullet-specific implementation of the simulated gripper'''
    # names must match URDF.
    jointNames = [
        # finger A
        'finger_middle_joint_1', 'finger_middle_joint_2', 'finger_middle_joint_3',
        # finger B
        'finger_2_joint_1', 'finger_2_joint_2', 'finger_2_joint_3',
        # finger C
        'finger_1_joint_1', 'finger_1_joint_2', 'finger_1_joint_3',
        # scissor joints
        'palm_finger_2_joint', 'palm_finger_1_joint'
    ]

    fingerIndices = [slice(0, 3), slice(3, 6), slice(6, 9)]
    fingerAll = slice(0, 9)
    scissors = slice(9,11)

    def __init__(self, body_id):
        self.body_id = body_id
        self.joints = {}
        numJoints = pb.getNumJoints(body_id)
        names = self.jointNames
        for i in range(numJoints):
            info = pb.getJointInfo(body_id, i)
            jointName = info[1].decode("utf-8")
            if jointName not in names:
                continue

            self.joints[jointName] = info
            jointID = info[0]
            pb.resetJointState(self.body_id, jointID, 0)
            pb.setJointMotorControl2(body_id, jointID, pb.VELOCITY_CONTROL, targetVelocity=0, force=100)
            pb.enableJointForceTorqueSensor(self.body_id, jointID, True)

        self.jointIDs = [self.joints[name][0] for name in Robotiq3FGripper.jointNames]
        self._targets = [0, 0, 0]

        joint1Stops = np.array(range(256)) / 255.
        joint2Stops = np.zeros(256)
        joint2Stops[128:] = np.array(range(128)) / 128.
        joint3Stops = np.zeros(256)
        joint3Stops[:192] = np.array(range(0, -192, -1)) / 192.
        joint3Stops[192:] = np.array(range(-192, -128)) / 192.
        self.jointStops = np.zeros((256, 9))
        self.jointStops[:, 0] = joint1Stops
        self.jointStops[:, 3] = joint1Stops
        self.jointStops[:, 6] = joint1Stops

        self.jointStops[:, 1] = joint2Stops
        self.jointStops[:, 4] = joint2Stops
        self.jointStops[:, 7] = joint2Stops

        self.jointStops[:, 2] = joint3Stops
        self.jointStops[:, 5] = joint3Stops
        self.jointStops[:, 8] = joint3Stops

        # grasp modes: normal, pinch, wide, scissors, narrow
        # urdf joint limit are [0.16, -0.25], but 0.16 causes a self collision in pinch mode,
        # which messes up the FT sensor reading, so we use 0.13 instead
        self.modeStopsB = [0, 0.13, -0.25, -0.25, 0.13]
        self.modeStopsC = [0, -0.13, 0.25, 0.25, -0.13]
        self.last_joint_positions  = np.zeros(len(self.jointNames))
        self.last_joint_speeds  = np.zeros(len(self.jointNames))

    def reset(self):
        self.read()
        self.set_mode(0)
        self.move(0, 255, 255)

    def move(self, position, speed, force):
        self.read()
        self._targets[0] = self._targets[1] = self._targets[2] = position
        # This is a very rough approximation of how the real hand moves. Needs more work.
        speed = 2 + speed // 8 # min 2 to avoid stall
        current = self.positions()[0]
        inc = 0 if position == current else speed if position > current else -speed
        position = current + inc
        joints = self.jointIDs[self.fingerAll]
        position = min(position, self._mode_limit) # account for pinch mode
        positions = self.jointStops[position]
        force = force + 1 # non-zero
        forces = [force * 2] * 3 + [force] * 6 # middle finger has twice the force, to compensate for the other two
        pb.setJointMotorControlArray(self.body_id, joints, pb.POSITION_CONTROL, targetPositions=positions, forces=forces)

    def move_finger(self, finger, position, speed, force):
        self.read()
        self._targets[finger] = position
        # This is a very rough approximation of how the real hand moves. Needs more work.
        joints = self.jointIDs[self.fingerIndices[finger]]
        cnt = len(joints)
        position = min(position, self._mode_limit)  # account for pinch mode
        positions = self.jointStops[position, self.fingerIndices[finger]]
        forces = [force + 1] * cnt
        pb.setJointMotorControlArray(self.body_id, joints, pb.POSITION_CONTROL, targetPositions=positions, forces=forces)

    def stop(self):
        self.read()
        cnt = len(self.jointIDs)
        pb.setJointMotorControlArray(self.body_id, self.jointIDs, pb.VELOCITY_CONTROL, targetVelocities=[0] * cnt, forces=[255] * cnt)
        self._targets[:] = self.positions()

    def set_mode(self, mode):
        self.read()
        self._mode = mode
        ix = int(mode / 2)
        joints = [self.jointIDs[9], self.jointIDs[10]]
        positions = [self.modeStopsB[ix], self.modeStopsC[ix]]
        pb.setJointMotorControlArray(self.body_id, joints, pb.POSITION_CONTROL, targetPositions=positions, forces=[100, 100])

        # in pinch mode the fingers can only travel part way (until they touch).
        self._mode_limit = self._PINCH_LIMIT if mode == 2 else 255

    def read(self):
        newStates = pb.getJointStates(self.body_id, self.jointIDs)
        jointPos = np.array([s[0] for s in newStates])
        self._is_moving = not np.allclose(jointPos, self.last_joint_positions, atol=0.0001)
        self.last_joint_speeds[:] = [s[1] for s in newStates]
        self.last_joint_positions = jointPos

    def mode(self):
        return self._mode

    def positions(self):
        # hackish way of computing the finger position, by considering only the base joint of each finger.
        # Needs to be in sync with the implementaiton of "move" and the position tables
        pos = np.int32(self.last_joint_positions[:9:3] * 256)
        pos = [255 if p >= self._mode_limit - 1 else p for p in pos]
        return pos

    def targets(self):
        return self._targets

    def object_detected(self):
        d = not self._is_moving and not np.allclose(self._targets, self.positions(), atol=1)
        return d

    def is_moving(self):
        return self._is_moving