import copy
import json
import os
import ssl
import time
import urllib.parse
from datetime import datetime
from threading import Lock, Thread

import websocket

from ..config import Config, Settings
from ..onap.audit import Audit
from ..utils import Utils
from .pdp_consts import MATCHING_CONDITIONS, POLICY_NAME, POLICY_VERSION

LOADED_POLICIES = 'loadedPolicies'
REMOVED_POLICIES = 'removedPolicies'
POLICY_VER = 'versionNo'
POLICY_MATCHES = 'matches'

_LOGGER = Utils.get_logger(__file__)

class PolicyListener(Thread):
    """web-socket to PolicyEngine"""
    PDP_API_FOLDER = os.path.basename(os.path.dirname(os.path.realpath(__file__)))
    WS_STARTED = "started"
    WS_START_COUNT = "start_count"
    WS_CLOSE_COUNT = "close_count"
    WS_ERROR_COUNT = "error_count"
    WS_PONG_COUNT = "pong_count"
    WS_MESSAGE_COUNT = "message_count"
    WS_MESSAGE_TIMESTAMP = "message_timestamp"
    WS_STATUS = "status"
    WS_PING_INTERVAL_DEFAULT = 30
    WEB_SOCKET_HEALTH = "web_socket_health"

    def __init__(self, audit, policy_updater):
        """web-socket inside the thread to receive policy notifications from PolicyEngine"""
        Thread.__init__(self, name="policy_receiver", daemon=True)

        self._policy_updater = policy_updater
        self._lock = Lock()
        self._keep_running = True
        self._settings = Settings(Config.FIELD_POLICY_ENGINE)

        self._sleep_before_restarting = 5
        self._web_socket_url = None
        self._web_socket_sslopt = None
        self._tls_wss_ca_mode = None
        self._web_socket = None
        self._ws_ping_interval_in_secs = PolicyListener.WS_PING_INTERVAL_DEFAULT
        self._web_socket_health = {
            PolicyListener.WS_START_COUNT: 0,
            PolicyListener.WS_CLOSE_COUNT: 0,
            PolicyListener.WS_ERROR_COUNT: 0,
            PolicyListener.WS_PONG_COUNT: 0,
            PolicyListener.WS_MESSAGE_COUNT: 0,
            PolicyListener.WS_STATUS: "created"
        }

        Audit.register_item_health(PolicyListener.WEB_SOCKET_HEALTH, self._get_health)
        self.reconfigure(audit)

    def reconfigure(self, audit):
        """configure and reconfigure the web-socket"""
        with self._lock:
            _LOGGER.info(audit.info("web_socket_health {}".format(
                json.dumps(self._get_health(), sort_keys=True))))
            self._sleep_before_restarting = 5
            self._settings.set_config(Config.discovered_config)
            changed, config = self._settings.get_by_key(Config.FIELD_POLICY_ENGINE)

            if not changed:
                self._settings.commit_change()
                return False

            prev_web_socket_url = self._web_socket_url
            prev_web_socket_sslopt = self._web_socket_sslopt
            prev_ws_ping_interval_in_secs = self._ws_ping_interval_in_secs

            self._web_socket_sslopt = None

            resturl = urllib.parse.urljoin(config.get("url", "").lower().rstrip("/") + "/",
                                           config.get("path_notifications", "/pdp/notifications"))

            self._tls_wss_ca_mode = config.get(Config.TLS_WSS_CA_MODE)

            self._ws_ping_interval_in_secs = config.get(Config.WS_PING_INTERVAL_IN_SECS)
            if not self._ws_ping_interval_in_secs or self._ws_ping_interval_in_secs < 60:
                self._ws_ping_interval_in_secs = PolicyListener.WS_PING_INTERVAL_DEFAULT

            if resturl.startswith("https:"):
                self._web_socket_url = resturl.replace("https:", "wss:")

                verify = Config.get_tls_verify(self._tls_wss_ca_mode)
                if verify is False:
                    self._web_socket_sslopt = {'cert_reqs': ssl.CERT_NONE}
                elif verify is True:
                    pass
                else:
                    self._web_socket_sslopt = {'ca_certs': verify}

            else:
                self._web_socket_url = resturl.replace("http:", "ws:")

            log_changed = (
                "changed web_socket_url(%s) or tls_wss_ca_mode(%s)"
                " or ws_ping_interval_in_secs(%s): %s" %
                (self._web_socket_url, self._tls_wss_ca_mode, self._ws_ping_interval_in_secs,
                 self._settings))
            if (self._web_socket_url == prev_web_socket_url
                    and Utils.are_the_same(prev_web_socket_sslopt, self._web_socket_sslopt)
                    and prev_ws_ping_interval_in_secs == self._ws_ping_interval_in_secs):
                _LOGGER.info(audit.info("not {}".format(log_changed)))
                self._settings.commit_change()
                return False

            _LOGGER.info(audit.info(log_changed))
            self._settings.commit_change()

        self._stop_notifications()
        return True

    def run(self):
        """listen on web-socket and pass the policy notifications to policy-updater"""
        _LOGGER.info("starting policy_receiver...")
        websocket.enableTrace(True)
        restarting = False
        while True:
            if not self._get_keep_running():
                break

            self._stop_notifications()

            if restarting:
                with self._lock:
                    sleep_before_restarting = self._sleep_before_restarting
                _LOGGER.info(
                    "going to sleep for %s secs before restarting policy-notifications",
                    sleep_before_restarting)

                time.sleep(sleep_before_restarting)
                if not self._get_keep_running():
                    break

            with self._lock:
                web_socket_url = self._web_socket_url
                sslopt = copy.deepcopy(self._web_socket_sslopt)
                tls_wss_ca_mode = self._tls_wss_ca_mode
                ws_ping_interval_in_secs = self._ws_ping_interval_in_secs

            _LOGGER.info(
                "connecting to policy-notifications at %s with sslopt(%s) tls_wss_ca_mode(%s)"
                " ws_ping_interval_in_secs(%s)",
                web_socket_url, json.dumps(sslopt), tls_wss_ca_mode, ws_ping_interval_in_secs)

            self._web_socket = websocket.WebSocketApp(
                web_socket_url,
                on_open=self._on_ws_open,
                on_message=self._on_pdp_message,
                on_close=self._on_ws_close,
                on_error=self._on_ws_error,
                on_pong=self._on_ws_pong
            )

            _LOGGER.info("waiting for policy-notifications...")
            self._web_socket.run_forever(sslopt=sslopt, ping_interval=ws_ping_interval_in_secs)
            restarting = True

        Audit.register_item_health(PolicyListener.WEB_SOCKET_HEALTH)
        _LOGGER.info("exit policy-receiver")

    def _get_keep_running(self):
        """thread-safe check whether to continue running"""
        with self._lock:
            keep_running = self._keep_running
        return keep_running

    def _stop_notifications(self):
        """close the web-socket == stops the notification service if running."""
        with self._lock:
            if self._web_socket and self._web_socket.sock and self._web_socket.sock.connected:
                self._web_socket.close()
                _LOGGER.info("stopped receiving notifications from PDP")

    def _on_pdp_message(self, *args):
        """received the notification from PDP"""
        self._web_socket_health[PolicyListener.WS_MESSAGE_COUNT] += 1
        self._web_socket_health[PolicyListener.WS_MESSAGE_TIMESTAMP] = str(datetime.utcnow())
        try:
            message = args and args[-1]
            _LOGGER.info("Received notification message: %s", message)
            _LOGGER.info("web_socket_health %s", json.dumps(self._get_health(), sort_keys=True))
            if not message:
                return
            message = json.loads(message)

            if not message or not isinstance(message, dict):
                _LOGGER.warning("unexpected message from PDP: %s", json.dumps(message))
                return

            policies_updated = [
                {POLICY_NAME: policy.get(POLICY_NAME),
                 POLICY_VERSION: policy.get(POLICY_VER),
                 MATCHING_CONDITIONS: policy.get(POLICY_MATCHES, {})}
                for policy in message.get(LOADED_POLICIES, [])
            ]

            policies_removed = [
                {POLICY_NAME: removed_policy.get(POLICY_NAME),
                 POLICY_VERSION: removed_policy.get(POLICY_VER)}
                for removed_policy in message.get(REMOVED_POLICIES, [])
            ]

            if not policies_updated and not policies_removed:
                _LOGGER.info("no policy updated or removed")
                return

            self._policy_updater.policy_update(policies_updated, policies_removed)
        except Exception as ex:
            error_msg = "crash {} {} at {}: {}".format(type(ex).__name__, str(ex),
                                                       "on_pdp_message", json.dumps(message))

            _LOGGER.exception(error_msg)

    def _on_ws_error(self, error):
        """report an error"""
        _LOGGER.exception("policy-notification error %s", str(error))
        self._sleep_before_restarting = 60 if isinstance(error, ssl.SSLError) else 5

        self._web_socket_health[PolicyListener.WS_STATUS] = "error"
        self._web_socket_health[PolicyListener.WS_ERROR_COUNT] += 1
        self._web_socket_health["last_error"] = {
            "error": str(error), "timestamp": str(datetime.utcnow())
        }
        _LOGGER.info("web_socket_health %s", json.dumps(self._get_health(), sort_keys=True))

    def _on_ws_close(self, code, reason):
        """restart web-socket on close"""
        self._web_socket_health["last_closed"] = str(datetime.utcnow())
        self._web_socket_health[PolicyListener.WS_STATUS] = "closed"
        self._web_socket_health[PolicyListener.WS_CLOSE_COUNT] += 1
        _LOGGER.info(
            "lost connection(%s, %s) to PDP web_socket_health %s",
            code, reason, json.dumps(self._get_health(), sort_keys=True))

    def _on_ws_open(self):
        """started web-socket"""
        self._web_socket_health[PolicyListener.WS_STATUS] = PolicyListener.WS_STARTED
        self._web_socket_health[PolicyListener.WS_START_COUNT] += 1
        self._web_socket_health[PolicyListener.WS_STARTED] = datetime.utcnow()
        _LOGGER.info("opened connection to PDP web_socket_health %s",
                     json.dumps(self._get_health(), sort_keys=True))

    def _on_ws_pong(self, pong):
        """pong = response to pinging the server of the web-socket"""
        self._web_socket_health[PolicyListener.WS_PONG_COUNT] += 1
        _LOGGER.info(
            "pong(%s) from connection to PDP web_socket_health %s",
            pong, json.dumps(self._get_health(), sort_keys=True))

    def _get_health(self):
        """returns the healthcheck of the web-socket as json"""
        web_socket_health = copy.deepcopy(self._web_socket_health)
        web_socket_health[Config.WS_PING_INTERVAL_IN_SECS] = self._ws_ping_interval_in_secs
        started = web_socket_health.get(PolicyListener.WS_STARTED)
        if started:
            web_socket_health[PolicyListener.WS_STARTED] = str(started)
            web_socket_health["uptime"] = str(datetime.utcnow() - started)
        return web_socket_health

    def shutdown(self, audit):
        """Shutdown the policy_listener"""
        _LOGGER.info(audit.info("shutdown policy_listener"))
        with self._lock:
            self._keep_running = False

        self._stop_notifications()

        if self.is_alive():
            self.join()


#
#
# @param intervals int整型二维数组
# @return int整型
#
from collections import defaultdict

class Solution:
    def eraseOverlapIntervals(self, intervals):
        orders = []
        for interval in intervals:
            orders += [(interval[0], True), (interval[1], False)]
        orders.sort()
        use = defaultdict(int)
        for orderTime, orderType in orders:
            if orderType:
                use[orderTime] += 1
            else:
                use[orderTime] -= 1
        using = 0
        ans = 0
        for num in use.values():
            using += num
            ans = max(ans, using)
        return ans - 1

print(Solution().eraseOverlapIntervals([[1, 2], [2, 3], [3, 4], [1, 3]]))

2020.8/商汤-无重叠区间调度问题.py

商汤-无重叠区间调度问题.py


from django.shortcuts import render,get_object_or_404, get_list_or_404, redirect
from django.http import JsonResponse, HttpResponse
from django.utils.translation import ugettext as _
from django.contrib import messages
from django.contrib.auth.decorators import login_required
from contract.models import Contract
from .models import Annotation
from .form import AnnotationForm
from account.models import User
from django.db import IntegrityError
from django.db.models import Q, Sum, Count, Prefetch
import os
import json

####### SERVICETYPE  ################

def annotation_save_form(request,form,template_name, data, user_created=None):
    if request.method == 'POST':
        if form.is_valid():
            obj = form.save(commit=False)
            if user_created:# Se cair aqui é EDIT
                obj.user_created = user_created
            else:# Se cair aqui é CREATE
                obj.user_created = request.user
            obj.user_updated = request.user
            obj.save()

            return redirect('annotation:url_annotations_list',data['contract'].slug)
        else:
            print("algo não está valido.")

    data['form'] = form
    return render(request,template_name,data)

@login_required
def annotation_create(request,contract):
    contract = get_object_or_404(Contract, slug=contract)
    template_name = 'annotation/form.html'
    data = {
            "title": _("Create Annotation"),
            "back":_("Back"),
            "save":_("Save"),
            "clear":_("Clear"),
            "contract": contract,
        }
    if request.method == 'POST':
        form = AnnotationForm(request.POST, request.FILES,contract=contract.id)
    else:
        form = AnnotationForm(contract=contract.id)

    return annotation_save_form(request, form, template_name, data)

@login_required
def annotation_edit(request, slug):
    annotation = get_object_or_404(Annotation, slug=slug)
    contract = get_object_or_404(Contract, slug=annotation.contract.slug)
    template_name='annotation/form.html'
    data = {
            "title": _("Edit"),
            "back":_("Back"),
            "save":_("Save"),
            "clear":_("Clear"),
            "contract": contract,
        }
    user_created = annotation.user_created # Esta linha faz com que o user_created não seja modificado, para mostrar quem criou esta pessoa
    if request.method == 'POST':
        form = AnnotationForm(request.POST, request.FILES, instance=annotation)
    else:
        form = AnnotationForm(instance=annotation)
    return annotation_save_form(request, form, template_name, data, user_created=user_created)

@login_required
def annotations_list(request,contract):
    contract = get_object_or_404(Contract, slug=contract)
    template_name = "annotation/list.html"
    annotations = Annotation.objects.filter(contract=contract.id)
    context = {
        'annotations': annotations,
        'contract': contract,
        'title': _("Registered Annotations"),
        'add': _("Add"),
        'vars': [_("Contract"), _("Provider")],
        'back': _("Back"),
    }
    return render(request,template_name,context)

@login_required
def annotation_detail(request, slug):
    template_name = "annotation/detail.html"
    annotation = get_object_or_404(Annotation,slug=slug)
    contract = get_object_or_404(Contract,slug=annotation.contract.slug)
    context = {
        'annotation': annotation,
        'contract': contract,
        'title': _("Detail Info"),
        'edit': _("Edit"),
        'list_all': _("List All"),
        'vars': [_("Contract"), _("Provider")],
    }
    return render(request, template_name, context)

@login_required
def annotation_delete(request, slug):
    annotation = get_object_or_404(Annotation, slug=slug)
    slug = annotation.contract.slug
    if request.method == 'POST':
       try:
           annotation.delete()
           messages.success(request, _('Completed successful.'))
           return redirect('annotation:url_annotations_list', contract=slug)
       except IntegrityError:
           messages.warning(request, _('You cannot delete. This annotation has an existing deal.'))
           return redirect('annotation:url_annotations_list', contract=slug)

@login_required
def annotation_delete_all(request):
    marc = 0
    if request.method == "POST":
        context = request.POST["checkbox_selected"].split(",")
        context = [str(x) for x in context]
        if context:
            contract_slug = Annotation.objects.prefetch_related('contract').values('contract__slug').distinct()
            contract_slug = contract_slug[0]['contract__slug']
            b = Annotation.objects.filter(slug__in=context)
            for i in b:
                try:
                    i.delete()
                except IntegrityError:
                    marc = 1
    if marc == 0:
        messages.success(request, _('Completed successful.'))
    else:
        messages.warning(request, _('You cannot delete. This annotation has an existing deal.'))

    return redirect('annotation:url_annotations_list', contract=contract_slug)

########### FIM SERVICETYPE ############################

# VIEW PARA TRADUZIR O DATATABLES. USO GERAL
def translate_datables_js(request):
    if request.method == "GET":
        module_dir = os.path.dirname(__file__)  # get current directory
        file_path = os.path.join(module_dir, 'templates/default')
        translate = ""
        with open(file_path+"/translate_data_tables-"+request.LANGUAGE_CODE+".json", 'r') as arquivo:
            for linha in arquivo:
                translate += linha
        obj = json.loads(translate)
    return JsonResponse(obj)


from operator import mod
import pandas as pd
import scipy.stats as stats
from itertools import product
data=pd.read_csv("C:/Users/abomb/Projects/microbiotaPhd/allDiets_collection_total.csv")

def getCombinations(var):
    varList=pd.unique(data[var].tolist())
    output = list(product(varList, varList))
    return output

def defGetDietComb():
    normal=data.loc[data['Type']=='N' ]
    normalList=pd.unique(normal['Diet'].tolist())
    modified=data.loc[(data['Type']=='HF') | (data['Type']=='HS')]
    modifiedList=pd.unique(modified['Diet'].tolist())
    output = list(product(normalList, modifiedList))
    return output

# function to get p values of comparisons statistic and difference in means

def mwTest(compare, variable, treatment):
    comparisons=[]
    pVals=[]
    meanDifference=[]
    groups=defGetDietComb()
    dataSel=data.loc[(data['Treatment']==treatment)]
    for i in groups:
        var1=i[0]
        var2=i[1]
        varComparison=(var1+'_vs_'+var2)
        group1=dataSel.loc[(dataSel[variable]==var1)]
        xVar=group1[compare].dropna()
        xVarMean=xVar.mean()
        group2=dataSel.loc[(dataSel[variable]==var2)]
        yVar=group2[compare].dropna()
        yVarMean=yVar.mean()
        U1, p=stats.mannwhitneyu(xVar, yVar, nan_policy='omit')
        difference=(xVarMean-yVarMean)
        # append lists
        comparisons.append(varComparison)
        pVals.append(p)
        meanDifference.append(difference)
    testResults=pd.DataFrame({'Comparisons': comparisons, 'pVal': pVals, 'Mean_Difference': meanDifference})
    return(testResults)

totalDiet=mwTest(compare='Sum(Total Collected)', variable='Diet', treatment='S')
totalDeitSel=totalDiet.iloc[[1,5,4,8,10,16,20]].reset_index(drop=True)

totalDeitSel.to_csv("C:/Users/abomb/Projects/microbiotaPhd/output/totalDietS.csv", index=False)


import subprocess
import streamlit as st
import torch
import torchvision.transforms as T
import torchvision.models as models
import numpy as np
import pandas as pd
import cv2
import torch
import glob as glob
from PIL import Image, ImageDraw
from py2neo import Graph
import plotly.graph_objects as go
import networkx as nx
import matplotlib.pyplot as plt
# from model import FasterRCNN

def get_scores():
    scores = []
    with open('./example_test_data/score.txt', 'r') as fi:
        for line in fi:
            scores.append(line.strip('\n'))
    fi.close()
    st.divider()
    st.write('### TEST PREDICTIONS COMPLETE')
    col1, col2, col3 = st.columns(3)
    col1.metric(label="Total FPS", value=scores[0][:5])
    col2.metric(label="Frame count", value=scores[1])
    col3.metric(label="Average FPS", value=scores[2][:5])
    st.divider()

def get_classes():
    clas = []
    with open('./example_test_data/clas.txt', 'r') as f:
        for line in f:
            clas.append(line.strip('\n'))
    f.close()
    return clas

def run_inference(image_name):

    input_file = './example_test_data/test_image.jpg'
    weights_file = "./outputs/training/custom_training/last_model.pth"

    command = f"python inference.py --input {input_file} --weights {weights_file}"

    # command = "python inference.py --input /content/drive/MyDrive/faster-rcnn-test/fastercnn-pytorch-training-pipeline/example_test_data/operating-room.jpg --weights outputs/training/custom_training/last_model_state.pth"

    subprocess.call(command.split())
    OUT_DIR_NUM = 0
    with open('./example_test_data/dir_num.txt', 'r') as fil:
        OUT_DIR_NUM = fil.readline()
        print(OUT_DIR_NUM)
    fil.close()
    OUT_DIR_NUM = int(OUT_DIR_NUM)
    output_file = f"./outputs/inference/res_{OUT_DIR_NUM}/.jpg"
    return output_file

# Define the Streamlit app
def main():
    # Set the title and description of the app
    tab1, tab2 = st.tabs(["About", "Demo"])
    with tab1:
        st.title("Knowledge-based Scene Graph Generation in Medical Field")
        im1 = Image.open('./main_img.png')
        st.image(im1, "Example of an image and it's scene graph")
        st.write('''
                    **Scene understanding** in the medical field  based on visual context using object detection and representing it in the form of knowledge graphs in order to derive conclusions or understand contexts.

    This problem statement highlights understanding the hospital environment and deciding the **relationship between the objects detected** in the image.

                    The major contributions of our proposed work are:
    - A Knowledge- based  representation in the form of scene graphs of images and scenarios seen in the medical field.
    - Accurate analysis or understanding of the objects in the scene presented to it as an image.
    - Publication of the accomplished work in domain-specific Conference/Journal.
    - Contribution to future research in Scene understanding and Intent recognition.

    ''')
        st.divider()
        st.subheader("Methodology")
        #im4 = Image.open('G:/knowledge_graphs_faster_rcnn/block1.png')
        im4 = Image.open('./arch_img.png')
        st.image(im4, 'Architecture')
        st.write('''
                We used Faster RCNN model to train our custom dataset. Restricting our project to the medical field, our dataset consists of 11 classes.
                These classes depict the day to day basic scenarios in a hospital - *cardiac_monitor, doctor, face_mask, iv_fluids, oxygen_mask, patient, stethoscope, syringe, ventilator, lightings, patient_bed.* \n
                The nodes and relationships are created using SVO triplets - Subject Verb and Objects. A complete triplet includes the predicate label and the class labels as well as the bounding box coordinates of the subject and object. All inferences will be based on the SVO triplets.\n
                The predicate classes used are - USES, HAS, IN, LAY_ON, RECEIVES, CHECKS, READS, TREATS, INJECTED_TO, WEARS. The knowledge base is populated with the SVO triplets that depicts simple scenarios in the medical field. We have 12 nodes and 21 relationsips as of now.
                ''')
        # Allow the user to upload an image file
        #im5 = Image.open('G:/knowledge_graphs_faster_rcnn/graph (1).png')
        im5 = Image.open('./kb_full.png')
        st.image(im5, 'Knowledge Base')

        st.divider()
        st.subheader('Performance')
        st.write('''
                We evaluate the object detection model using Precision, Recall, mAP for each bounding box detected via IoU (Intersection of Unions).
                Since we used bbox as the means of evaluation, Faster RCNN provides these metrics for each bbox involved. \n
                The method shows the results of both object detection (bounding boxes and object labels) and edge predictions (relationship labels), i.e., a scene graph.
                Relationship labels are represented as directed edges.
                ''')
        results_eval = pd.read_csv(
            './results.csv')
        st.write(results_eval)
        imm = Image.open(
            './map.png')
        st.image(imm, "mAP mean Average Precision")
        imm1 = Image.open(
            './train_loss_epoch.png')
        st.image(imm1, "Training loss")
        st.subheader("Future Work")
        st.write('''
                 Scene graph generation (SGG) is a semantic
understanding task that goes beyond object detection and is
closely linked to visual relationship detection. At present,
scene graphs have shown their potential in different visionlanguage tasks such as image retrieval, image captioning, visual question answering (VQA) and image
generation. The task of scene graph generation has
also received sustained attention in the computer vision
community.\n
We believe that the biggest problem of classic scene graph generation (SGG) comes from noisy datasets. Classic scene graph generation datasets adopt a bounding box-based object grounding, which inevitably causes a number of issues:

- Coarse localization: bounding boxes cannot reach pixel-level accuracy,
- Inability to ground comprehensively: bounding boxes cannot ground backgrounds,
- Tendency to provide trivial information: current datasets usually capture frivolous objects like head to form trivial relations like person-has-head, due to too much freedom given during bounding box annotation.
- Duplicate groundings: the same object could be grounded by multiple separate bounding boxes.
                 ''')
    with tab2:
        st.write("Upload an image related to a scenario in a hospital")
        uploaded_file = st.file_uploader(
            "Choose an image...", type=["jpg", "jpeg", "png"])

        # If an image file is uploaded
        if uploaded_file is not None:
            # Open the image file using PIL
            st.image(uploaded_file)
            image = Image.open(uploaded_file)
            image = image.save(
                './example_test_data/test_image.jpg')

            # frame = np.array(image)
            # Run inference on the image and get the result
            result_image = run_inference(uploaded_file)
            res = Image.open(result_image)
            # Display the resulting image
            st.success("Object Detection complete!")
            st.image(res, caption='Detected objects', use_column_width=True)
            get_scores()
            clas = get_classes()
            graph_query(clas)

def inf_output():
    clas = []
    with open('./example_test_data/final_inf.txt', 'r') as f:
        for line in f:
            clas.append(line.strip('\n'))
    f.close()
    st.subheader('FINAL INFERENCE:')
    im2 = Image.open(
        './example_test_data/test_image.jpg')
    st.image(im2, "Input Image")
    for i in set(clas):
        st.write(f'##### {i}')
    # st.write(str(set(clas)))

def check_room(classes, clas, graph):
    final_inf = []
    op_room = ['doctor', 'iv_fluids', 'ventilator',
               'lightings', 'cardiac_monitor', 'patient_bed']
    if len(classes) < len(op_room):
        check = all(item in classes for item in op_room)
    else:
        check = all(item in op_room for item in classes)
    if check:
        # query3 = "MATCH (d)-[r1]->(n1:"+clas+") RETURN d,r1,n1"
        qu2 = "MATCH (d)-[r1]->(n1:"+clas+") WHERE n1.name ='" + \
            clas+"' AND NOT d.name = 'Icu_room' RETURN d, r1, n1"
        results3 = graph.run(qu2).data()
        if results3:
            st.write('### KNOWLEDGE BASED INFERENCE ON : ')
            st.subheader(f'Nodes labeled {clas}:')
            for row in results3:
                temp = ' '.join([row['d']['name'], row['r1']
                                ['name'], row['n1']['name']])
                final_inf.append(temp)
                st.write(row['d']['name'], row['r1']
                         ['name'], row['n1']['name'])
            with open('./example_test_data/final_inf.txt', 'a') as file:
                for pr in final_inf:
                    file.write(pr)
                    file.write('\n')
                file.close()
            st.divider()
        else:
            st.write(f'No relevant inference found for {clas}')
            st.divider()
        if clas == 'Syringe':
            clas = 'null'
    else:
        qu2 = "MATCH (d)-[r1]->(n1:"+clas+") WHERE n1.name ='" + \
            clas+"' AND NOT d.name = 'Operating_room' RETURN d, r1, n1"
        results3 = graph.run(qu2).data()
        if results3:
            st.write('### KNOWLEDGE BASED INFERENCE ON : ')
            st.subheader(f'Nodes labeled {clas}:')
            for row in results3:
                temp = ' '.join([row['d']['name'], row['r1']
                                ['name'], row['n1']['name']])
                final_inf.append(temp)
                st.write(row['d']['name'], row['r1']
                         ['name'], row['n1']['name'])
            with open('./example_test_data/final_inf.txt', 'a') as file:
                for pr in final_inf:
                    file.write(pr)
                    file.write('\n')
                file.close()
            st.divider()
        else:
            st.write(f'No relevant inference found for {clas}')
            st.divider()
        if clas == 'Syringe':
            clas = 'null'
    if clas == 'Syringe':
        qu5 = "MATCH (d:Syringe)-[r1]->(n1)  RETURN d, r1, n1"
        results6 = graph.run(qu5).data()
        if results6:
            st.write('### KNOWLEDGE BASED INFERENCE ON : ')
            st.subheader(f'Nodes labeled {clas}:')
            for row in results6:
                temp = ' '.join([row['d']['name'], row['r1']
                                ['name'], row['n1']['name']])
                final_inf.append(temp)
                st.write(row['d']['name'], row['r1']
                         ['name'], row['n1']['name'])
            with open('./example_test_data/final_inf.txt', 'a') as file:
                for pr in final_inf:
                    file.write(pr)
                    file.write('\n')
                file.close()
            st.divider()
        else:
            st.write(f'No relevant inference found for {clas}')
            st.divider()

def graph_query(classes):
    graph = Graph("neo4j://localhost:7687", auth=("neo4j", "password"))
    # Run a query to find all nodes labeled "Doctor"
    st.write("Predicted classes: ")
    st.write(classes)
    ans = ', '.join(classes)

    st.write(f'The image contains **{ans}**')
    st.divider()
    new_clas = classes.copy()
    for clas in classes:
        clas = clas.capitalize()
        tab_auto = st.tabs([clas])
        if clas == 'Doctor':
            # query1 = "MATCH (n:Node {label:'"+clas+"'}) RETURN n"
            if clas in new_clas:
                new_clas.remove(clas)

            for nw in new_clas:
                nw = nw.capitalize()
                query1 = "MATCH (d:Doctor)-[r1]->(n1:"+nw+") RETURN d, r1, n1"
                results1 = graph.run(query1).data()
                if results1:
                    st.write(" **Nodes labeled 'Doctor':**")
                    st.subheader(f'Nodes labeled {nw}:')
                    for row in results1:
                        # print(row['d'],row['r1'],row['n1'])
                        st.write(row['d']['name'], row['r1']
                                 ['name'], row['n1']['name'])
                    display_graph(results1)
            query1 = "MATCH (d:Doctor)-[r1]->(n1) RETURN d, r1, n1"
            results11 = graph.run(query1).data()
            if results11:
                st.write("### KNOWLEDGE BASED INFERENCES")
                for row in results11:
                    # print(row['d'],row['r1'],row['n1'])
                    st.write(row['d']['name'], row['r1']
                             ['name'], row['n1']['name'])
                display_graph(results11)

        elif clas == 'Patient':
            if clas in new_clas:
                new_clas.remove(clas)

            for nw in new_clas:
                nw = nw.capitalize()
                query1 = "MATCH (d:Patient)-[r1]->(n1:"+nw+") RETURN d, r1, n1"
                results1 = graph.run(query1).data()
                if results1:
                    st.subheader(f"Nodes labeled {nw}:")
                    for row in results1:
                        st.write(row['d']['name'], row['r1']
                                 ['name'], row['n1']['name'])
                    display_graph(results1)
            query2 = "MATCH (d:Patient)-[r1]->(n1) RETURN d, r1, n1"
            results2 = graph.run(query2).data()
            if results2:
                st.write("### KNOWLEDGE BASED INFERENCES")
                st.subheader("Nodes labeled 'Patients':")
                for row in results2:
                    st.write(row['d']['name'], row['r1']
                             ['name'], row['n1']['name'])
                display_graph(results2)
        else:
            query3 = "MATCH (d)-[r1]->(n1:"+clas+") RETURN d,r1,n1"
            results3 = graph.run(query3).data()
            if results3:
                st.subheader(f'Nodes labeled {clas} (object):')
                for row in results3:
                    st.write(row['d']['name'], row['r1']
                             ['name'], row['n1']['name'])
                display_graph(results3)
                check_room(classes, clas, graph)

            query4 = "MATCH (d:"+clas+")-[r1]->(n1) RETURN d,r1,n1"
            results4 = graph.run(query4).data()
            if results4:
                st.subheader(f'Nodes labeled {clas} (subject):')
                for row in results4:
                    st.write(row['d']['name'], row['r1']
                             ['name'], row['n1']['name'])
                display_graph(results4)
                check_room(classes, clas, graph)
    inf_output()

def display_graph(results):

    fig = plt.figure(figsize=(8, 6))
    G = nx.DiGraph()
    for record in results:
        G.add_edge(record['d']['name'], record['n1']
                   ['name'], label=record['r1']['name'])

    # draw graph
    pos = nx.spring_layout(G)
    nx.draw_networkx_nodes(G, pos)
    nx.draw_networkx_edges(G, pos)
    nx.draw_networkx_labels(G, pos)
    nx.draw_networkx_edge_labels(G, pos, edge_labels={(
        u, v): d['label'] for u, v, d in G.edges(data=True)})
    if G.edges:
        plt.show()
        st.pyplot(fig)

# Run the Streamlit app
if __name__ == '__main__':
    main()


#!/usr/bin/env python
# coding: utf-8

import os
import torch
import numpy as np
import random
from torch import nn
from torchvision import transforms as tfs
from torch.utils.checkpoint import checkpoint, checkpoint_sequential
from tqdm import tqdm
import hiddenlayer as hl
from PIL import Image
from PIL import ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True
import pdb
import cv2

class AverageValueMeter():
    def __init__(self):
        self.N = 0
        self.ave = 0

    def add(self, x):
        self.ave = self.N / (1 + self.N) * self.ave + 1 / (1 + self.N) * x
        self.N += 1

    def reset(self):
        self.N = 0
        self.ave = 0

    def value(self):
        return self.ave, self.N

def read_as_list(rootpath):
    datalist = os.listdir(rootpath)
    datapath = [os.path.join(rootpath, i) for i in datalist]
    return datapath

def readimglab(img_root, lab_root):
    ip = read_as_list(img_root)
    lp = read_as_list(lab_root)
    return ip, lp

def random_crop(img, lab, crop_size, rc=True):
    wd, ht = crop_size
    w, h = img.size
    wd = min(wd, w)
    ht = min(ht, h)
    wd = wd - wd % 16
    ht = ht - ht % 16
    if rc:
        if w == wd:
            i = 0
        else:
            i = random.randint(0, w - wd)
        if h == ht:
            j = 0
        else:
            j = random.randint(0, h - ht)
    else:
        i = 0
        j = 0

    return img.crop((i, j, i + wd, j + ht)), lab.crop((i, j, i + wd, j + ht))

def img_transforms(img, lab_d, crop_size, rc=True):
    #     img, lab = random_crop(img, lab, crop_size, rc=True)
    # 转化为SAR图像，即，灰度化，加椒盐噪声
    imgray = np.array(img)
    #     imgray = imgray[:, :, np.newaxis]
    #     imgray = np.asarray(img)
    img_tfs = tfs.Compose([
        tfs.ToTensor(),
        #         tfs.Normalize([0.5], [0.225]),
        tfs.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
    ])
    img = img_tfs(imgray)

    lab_d = np.array(lab_d)
    lab = np.ones(lab_d.shape[:2], dtype=np.int)
    b = lab_d[:, :, 0] == 0
    g = lab_d[:, :, 1] == 0
    r = lab_d[:, :, 2] == 128
    lab[~b] = 0
    lab[~g] = 0
    lab[~r] = 0
    return img, torch.from_numpy(lab)

class MapDataset(torch.utils.data.Dataset):
    """
    Map dataset
    """
    def __init__(self, crop_size, transform, img_root, lab_root, rc=True):
        self.crop_size = crop_size
        self.rc = rc
        self.transform = transform
        self.data_list, self.label_list = readimglab(img_root, lab_root)
        print("Read " + str(len(self.data_list)) + " images")

    def __getitem__(self, index):
        img = self.data_list[index]
        lab = self.label_list[index]
        img = Image.open(img)
        #         lab = Image.open(lab)
        lab = cv2.imread(lab)
        img, lab = self.transform(img, lab, self.crop_size, rc=self.rc)
        return img, lab

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

def _fast_hist(label_true, label_pred, n_class):
    mask = (label_true >= 0) & (label_true < n_class)
    hist = np.bincount(n_class * label_true[mask].astype(int) +
                       label_pred[mask],
                       minlength=n_class**2).reshape(n_class, n_class)
    return hist

def label_accuracy_score(label_trues, label_preds, n_class):
    """Returns accuracy score evaluation result.

      - overall accuracy
      - mean accuracy
      - mean IU
      - fwavacc
    """
    hist = np.zeros((n_class, n_class))
    for lt, lp in zip(label_trues, label_preds):
        hist += _fast_hist(lt.flatten(), lp.flatten(), n_class)
    acc = np.diag(hist).sum() / hist.sum()
    with np.errstate(divide='ignore', invalid='ignore'):
        acc_cls = np.diag(hist) / hist.sum(axis=1)
    acc_cls = np.nanmean(acc_cls)
    with np.errstate(divide='ignore', invalid='ignore'):
        iu = np.diag(hist) / (hist.sum(axis=1) + hist.sum(axis=0) -
                              np.diag(hist))
    mean_iu = np.nanmean(iu)
    freq = hist.sum(axis=1) / hist.sum()
    fwavacc = (freq[freq > 0] * iu[freq > 0]).sum()
    return acc, acc_cls, mean_iu, fwavacc

def bilinear_kernel(in_channels, out_channels, kernei_size):
    factor = (kernei_size + 1) // 2
    if kernei_size % 2 == 1:
        center = factor - 1
    else:
        center = factor - 0.5
    og = np.ogrid[:kernei_size, :kernei_size]
    filt = (1 - abs(og[0] - center) / factor) * (1 -
                                                 abs(og[1] - center) / factor)
    filtr = filt.reshape(-1)
    filtr = np.tile(filtr, (in_channels, out_channels))
    weight = filtr.reshape(in_channels, out_channels, filt.shape[0],
                           filt.shape[1]).astype(np.float32)
    #     weight = np.zeros((in_channels, out_channels, kernei_size, kernei_size),
    #                       dtype="float32")
    #     print(weight.shape, filt.shape)
    #     weight[range(in_channels), range(out_channels), :, :] = filt

    return torch.from_numpy(weight)

class MFM(nn.Module):
    def __init__(self):
        super(MFM, self).__init__()

    def forward(self, x):
        channels = x.shape[1]
        xo = x[:, 0, :, :]
        xo = torch.unsqueeze(xo, 1)
        for k in range(0, channels, 2):
            xout = torch.cat([
                torch.unsqueeze(x[:, k, :, :], 1),
                torch.unsqueeze(x[:, k + 1, :, :], 1)
            ],
                             dim=1)
            xout, _ = torch.max(xout, dim=1)
            xout = torch.unsqueeze(xout, 1)
            xo = torch.cat([xo, xout], dim=1)
        xo = xo[:, 1:, :, :]
        return xo

class FCN(nn.Module):
    """
    每一层的卷积深度都是64，应该逐层增加，可以为 64,128,256,512，递增
    """
    def __init__(self, in_ch, num_classes):
        super(FCN, self).__init__()
        self.downsample = nn.Conv2d(in_ch, 3, 3, stride=2,
                                    padding=1)  # 相当于一个 preconv 吧
        self.CBM1_1 = nn.Sequential(
            nn.Conv2d(3, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            MFM(),
        )
        self.CBM1_2 = nn.Sequential(
            nn.Conv2d(64 // 2, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            MFM(),
        )
        self.pool1 = nn.MaxPool2d(2, 2, ceil_mode=False)
        self.CBM2_1 = nn.Sequential(
            nn.Conv2d(64 // 2, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            MFM(),
        )
        self.CBM2_2 = nn.Sequential(
            nn.Conv2d(64 // 2, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            MFM(),
        )
        self.pool2 = nn.MaxPool2d(2, 2, ceil_mode=False)

        self.CBM3_1 = nn.Sequential(
            nn.Conv2d(64 // 2, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            MFM(),
        )
        self.CBM3_2 = nn.Sequential(
            nn.Conv2d(64 // 2, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            MFM(),
        )
        self.pool3 = nn.MaxPool2d(2, 2, ceil_mode=False)

        self.CBM4_1 = nn.Sequential(
            nn.Conv2d(64 // 2, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            MFM(),
        )
        self.CBM4_2 = nn.Sequential(
            nn.Conv2d(64 // 2, 64, 3, padding=1),
            nn.BatchNorm2d(64),
            MFM(),
        )

        #         self.upsample_pool1 = nn.ConvTranspose2d(
        #             64 // 2, 64, 4, stride=2,padding=1, bias=False)
        #         self.upsample_pool1.weight.data = bilinear_kernel(64 // 2, 64, 4)
        self.upsample_pool1 = nn.Conv2d(64 // 2, 64, 3, stride=1, padding=1)

        self.upsample_pool2 = nn.ConvTranspose2d(64 // 2,
                                                 64,
                                                 4,
                                                 stride=2,
                                                 padding=1,
                                                 bias=False)
        self.upsample_pool2.weight.data = bilinear_kernel(64 // 2, 64, 4)

        self.upsample_pool3 = nn.ConvTranspose2d(64 // 2,
                                                 64,
                                                 4,
                                                 stride=4,
                                                 padding=0,
                                                 bias=False)
        self.upsample_pool3.weight.data = bilinear_kernel(64 // 2, 64, 4)

        self.upsample_pool4 = nn.ConvTranspose2d(64 // 2,
                                                 64,
                                                 8,
                                                 stride=8,
                                                 padding=0,
                                                 bias=False)
        self.upsample_pool4.weight.data = bilinear_kernel(64 // 2, 64, 8)

        self.fcn1 = nn.Conv2d(64 * 4, 1024, 1, stride=1, padding=0)
        self.fcn2 = nn.Conv2d(1024, num_classes, 1, stride=1, padding=0)

        self.upsample_out_1 = nn.ConvTranspose2d(num_classes,
                                                 num_classes,
                                                 4,
                                                 stride=2,
                                                 padding=1,
                                                 bias=False)
        self.upsample_out_1.weight.data = bilinear_kernel(
            num_classes, num_classes, 4)

#         self.upsample_out_2 = nn.ConvTranspose2d(
#             num_classes, num_classes, 4, stride=2, padding=1,bias=False)
#         self.upsample_out_2.weight.data = bilinear_kernel(
#             num_classes, num_classes, 4)

    def forward(self, x, construct=False):
        h = x
        h = self.downsample(h)
        #         print("downsample",h.shape)#(128,128)
        h = self.CBM1_1(h)
        h = self.CBM1_2(h)
        p1 = h
        h = self.pool1(h)
        h = self.CBM2_1(h)
        h = self.CBM2_2(h)
        p2 = h
        h = self.pool2(h)
        h = self.CBM3_1(h)
        h = self.CBM3_2(h)
        p3 = h
        h = self.pool3(h)
        h = self.CBM4_1(h)
        h = self.CBM4_2(h)

        #         print("p123h",p1.shape, p2.shape, p3.shape, h.shape)
        p1u = self.upsample_pool1(p1)
        p2u = self.upsample_pool2(p2)
        p3u = self.upsample_pool3(p3)
        p4u = self.upsample_pool4(h)
        #         print("p1234u",p1u.shape, p2u.shape, p3u.shape, p4u.shape)
        h = torch.cat((p1u, p2u, p3u, p4u), dim=1)
        if construct:
            return h
        h = self.fcn1(h)
        h = self.fcn2(h)
        h = self.upsample_out_1(h)
        #         h = self.upsample_out_2(h)
        return h

def get_upsampling_weight(in_channels, out_channels, kernel_size):
    """Make a 2D bilinear kernel suitable for upsampling"""
    factor = (kernel_size + 1) // 2
    if kernel_size % 2 == 1:
        center = factor - 1
    else:
        center = factor - 0.5
    og = np.ogrid[:kernel_size, :kernel_size]
    filt = (1 - abs(og[0] - center) / factor) * (1 -
                                                 abs(og[1] - center) / factor)
    weight = np.zeros((in_channels, out_channels, kernel_size, kernel_size),
                      dtype=np.float64)
    weight[range(in_channels), range(out_channels), :, :] = filt
    return torch.from_numpy(weight).float()

class bnFCN8s(nn.Module):
    def __init__(self, encoder, n_class, th_ch=1):
        super().__init__()
        self.n_class = n_class
        self.encoder = encoder
        self.relu = nn.ReLU(inplace=True)
        self.deconv1 = nn.ConvTranspose2d(512,
                                          512,
                                          kernel_size=3,
                                          stride=2,
                                          padding=1,
                                          dilation=1,
                                          output_padding=1)
        self.bn1 = nn.BatchNorm2d(512)
        self.deconv2 = nn.ConvTranspose2d(512,
                                          256,
                                          kernel_size=3,
                                          stride=2,
                                          padding=1,
                                          dilation=1,
                                          output_padding=1)
        self.bn2 = nn.BatchNorm2d(256)
        self.deconv3 = nn.ConvTranspose2d(256,
                                          128,
                                          kernel_size=3,
                                          stride=2,
                                          padding=1,
                                          dilation=1,
                                          output_padding=1)
        self.bn3 = nn.BatchNorm2d(128)
        self.deconv4 = nn.ConvTranspose2d(128,
                                          64,
                                          kernel_size=3,
                                          stride=2,
                                          padding=1,
                                          dilation=1,
                                          output_padding=1)
        self.bn4 = nn.BatchNorm2d(64)
        self.deconv5 = nn.ConvTranspose2d(64,
                                          32,
                                          kernel_size=3,
                                          stride=2,
                                          padding=1,
                                          dilation=1,
                                          output_padding=1)
        self.bn5 = nn.BatchNorm2d(32)
        self.classifier = nn.Conv2d(32, n_class, kernel_size=1)

        self.relu2 = nn.ReLU(inplace=True)
        self.deconv12 = nn.ConvTranspose2d(512,
                                           512,
                                           kernel_size=3,
                                           stride=2,
                                           padding=1,
                                           dilation=1,
                                           output_padding=1)
        self.bn12 = nn.BatchNorm2d(512)
        self.deconv22 = nn.ConvTranspose2d(512,
                                           256,
                                           kernel_size=3,
                                           stride=2,
                                           padding=1,
                                           dilation=1,
                                           output_padding=1)
        self.bn22 = nn.BatchNorm2d(256)
        self.deconv32 = nn.ConvTranspose2d(256,
                                           128,
                                           kernel_size=3,
                                           stride=2,
                                           padding=1,
                                           dilation=1,
                                           output_padding=1)
        self.bn32 = nn.BatchNorm2d(128)
        self.deconv42 = nn.ConvTranspose2d(128,
                                           64,
                                           kernel_size=3,
                                           stride=2,
                                           padding=1,
                                           dilation=1,
                                           output_padding=1)
        self.bn42 = nn.BatchNorm2d(64)
        self.deconv52 = nn.ConvTranspose2d(64,
                                           32,
                                           kernel_size=3,
                                           stride=2,
                                           padding=1,
                                           dilation=1,
                                           output_padding=1)
        self.bn52 = nn.BatchNorm2d(32)
        self.classifier2 = nn.Conv2d(32, th_ch, kernel_size=1)

    def forward(self, x):
        output = self.encoder(x)
        x5 = output['x5']  # size=(N, 512, x.H/32, x.W/32)
        x4 = output['x4']  # size=(N, 512, x.H/16, x.W/16)
        x3 = output['x3']  # size=(N, 256, x.H/8,  x.W/8)

        score = self.relu(self.deconv1(x5))  # size=(N, 512, x.H/16, x.W/16)
        score = self.bn1(score +
                         x4)  # element-wise add, size=(N, 512, x.H/16, x.W/16)
        score = self.relu(self.deconv2(score))  # size=(N, 256, x.H/8, x.W/8)
        score = self.bn2(score +
                         x3)  # element-wise add, size=(N, 256, x.H/8, x.W/8)
        score = self.bn3(self.relu(
            self.deconv3(score)))  # size=(N, 128, x.H/4, x.W/4)
        score = self.bn4(self.relu(
            self.deconv4(score)))  # size=(N, 64, x.H/2, x.W/2)
        score = self.bn5(self.relu(
            self.deconv5(score)))  # size=(N, 32, x.H, x.W)
        score = self.classifier(score)  # size=(N, n_class, x.H/1, x.W/1)

        score2 = self.relu2(self.deconv12(x5))  # size=(N, 512, x.H/16, x.W/16)
        score2 = self.bn12(
            score2 + x4)  # element-wise add, size=(N, 512, x.H/16, x.W/16)
        score2 = self.relu2(
            self.deconv22(score2))  # size=(N, 256, x.H/8, x.W/8)
        score2 = self.bn22(score2 +
                           x3)  # element-wise add, size=(N, 256, x.H/8, x.W/8)
        score2 = self.bn32(self.relu2(
            self.deconv32(score2)))  # size=(N, 128, x.H/4, x.W/4)
        score2 = self.bn42(self.relu2(
            self.deconv42(score2)))  # size=(N, 64, x.H/2, x.W/2)
        score2 = self.bn52(self.relu2(
            self.deconv52(score2)))  # size=(N, 32, x.H, x.W)
        score2 = self.classifier2(score2)  # size=(N, n_class, x.H/1, x.W/1)

        return score, score2  # size=(N, n_class, x.H/1, x.W/1)

from torchvision.models.vgg import VGG

class VGGNet(VGG):
    def __init__(self,
                 model='vgg16',
                 batch_norm=False,
                 requires_grad=True,
                 remove_fc=True,
                 show_params=False,
                 in_channels=3):
        super().__init__(
            make_layers(cfg[model], batch_norm, in_channels=in_channels))
        if batch_norm:
            self.ranges = ranges[model + 'bn']
        else:
            self.ranges = ranges[model]

        if not requires_grad:
            for param in super().parameters():
                param.requires_grad = False

        if remove_fc:  # delete redundant fully-connected layer params, can save memory
            del self.classifier

        if show_params:
            for name, param in self.named_parameters():
                print(name, param.size())

    def forward(self, x):
        output = {}

        # get the output of each maxpooling layer (5 maxpool in VGG net)
        for idx in range(len(self.ranges)):
            for layer in range(self.ranges[idx][0], self.ranges[idx][1]):
                x = self.features[layer](x)
            output["x%d" % (idx + 1)] = x

        return output

    def _initialize_weights(self):
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight,
                                        mode='fan_out',
                                        nonlinearity='relu')
                if m.bias is not None:
                    nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.constant_(m.weight, 1)
                nn.init.constant_(m.bias, 0)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.constant_(m.bias, 0)

def make_layers(cfg, batch_norm=False, in_channels=3):
    layers = []
    in_channels = in_channels
    for v in cfg:
        if v == 'M':
            layers += [nn.MaxPool2d(kernel_size=2, stride=2)]
        else:
            conv2d = nn.Conv2d(in_channels, v, kernel_size=3, padding=1)
            if batch_norm:
                layers += [conv2d, nn.BatchNorm2d(v), nn.ReLU(inplace=True)]
            else:
                layers += [conv2d, nn.ReLU(inplace=True)]
            in_channels = v
    return nn.Sequential(*layers)

cfg = {
    'vgg11': [64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'vgg13':
    [64, 64, 'M', 128, 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M'],
    'vgg16': [
        64, 64, 'M', 128, 128, 'M', 256, 256, 256, 'M', 512, 512, 512, 'M',
        512, 512, 512, 'M'
    ],
    'vgg19': [
        64, 64, 'M', 128, 128, 'M', 256, 256, 256, 256, 'M', 512, 512, 512,
        512, 'M', 512, 512, 512, 512, 'M'
    ],
}

ranges = {
    'vgg11': ((0, 3), (3, 6), (6, 11), (11, 16), (16, 21)),
    'vgg11bn': ((0, 4), (4, 8), (8, 15), (15, 22), (22, 29)),
    'vgg13': ((0, 5), (5, 10), (10, 15), (15, 20), (20, 25)),
    'vgg13bn': ((0, 7), (7, 14), (14, 21), (21, 28), (28, 35)),
    'vgg16': ((0, 5), (5, 12), (12, 22), (22, 32), (32, 42)),
    'vgg16bn': ((0, 7), (7, 14), (14, 24), (24, 34), (34, 44)),
    'vgg19': ((0, 5), (5, 10), (10, 19), (19, 28), (28, 37)),
    'vgg19bn': ((0, 7), (7, 14), (14, 27), (27, 40), (40, 53))
}

# ## 边缘检测

class Edgenet(nn.Module):
    def __init__(self, ):
        """
        算子:
            Sobel 一阶微分算子, 包括纵横双斜角,共四个
            Laplace 二阶微分算子,共一个
        """
        super(Edgenet, self).__init__()
        f_x = torch.tensor([
            [-1, -2, -1],
            [0, 0, 0],
            [1, 2, 1],
        ])
        f_y = torch.tensor([
            [-1, 0, 1],
            [-2, 0, 2],
            [-1, 0, 1],
        ])
        f_tl_br = torch.tensor([
            [-2, -1, 0],
            [-1, 0, 1],
            [0, 1, 2],
        ])
        f_tr_bl = torch.tensor([
            [0, -1, -2],
            [1, 0, -1],
            [2, 1, 0],
        ])
        f_operator = [
            f_x,
            f_y,
            f_tl_br,
            f_tr_bl,
        ]
        self.conv_operator = nn.ModuleList()
        for i, operator_init in enumerate(f_operator):
            self.conv_operator.append(
                nn.Conv2d(
                    1,
                    1,
                    3,
                    padding=1,
                    bias=False,
                ))
            self.conv_operator[i].weight.data[:, :] = operator_init

    def forward(self, lab):
        edge = 0
        lab_l = lab.float()
        for i, conv_op in enumerate(self.conv_operator):
            edge += torch.abs(conv_op(lab_l))
        edge[edge > 0] = 1

        return edge.squeeze().long()

class Sobelnet(nn.Module):
    def __init__(self, in_ch=1, out_ch=1, coef_sobel=1.0):
        """
        算子:
            Sobel 一阶微分算子, 包括纵横双斜角,共四个
            Laplace 二阶微分算子,共一个
        """
        super(Sobelnet, self).__init__()
        f_x = torch.tensor([
            [-1, -2, -1],
            [0, 0, 0],
            [1, 2, 1],
        ])
        f_y = torch.tensor([
            [-1, 0, 1],
            [-2, 0, 2],
            [-1, 0, 1],
        ])
        f_tl_br = torch.tensor([
            [-2, -1, 0],
            [-1, 0, 1],
            [0, 1, 2],
        ])
        f_tr_bl = torch.tensor([
            [0, -1, -2],
            [1, 0, -1],
            [2, 1, 0],
        ])
        f_laplace = torch.tensor([
            [1, 1, 1],
            [1, 8, 1],
            [1, 1, 1],
        ])

        f_operator = [
            f_x,
            f_y,
            f_tl_br,
            f_tr_bl,
            #             f_laplace,
        ]
        self.in_ch = in_ch
        self.out_ch = out_ch
        self.coef_sobel = coef_sobel

        self.conv_operator = nn.ModuleList()
        for i, operator_init in enumerate(f_operator):
            self.conv_operator.append(
                nn.Conv2d(
                    self.in_ch,
                    self.out_ch,
                    3,
                    padding=1,
                    bias=False,
                ))
            self.conv_operator[i].weight.data[:, :] = operator_init

    def forward(self, output_th, pred_seg, lab):
        lab_1 = lab.float() * 1.8 - 0.9
        Res_p = pred_seg - output_th
        loss = 0
        coploss = []
        for i, conv_op in enumerate(self.conv_operator):
            copl = conv_op(lab_1)
            copr = conv_op(Res_p)
            coploss.append(copl)
            #             loss += torch.mean(((1 - lab) * pred_seg.unsqueeze(1) + lab *
            #                                (1 - pred_seg.unsqueeze(1))**0.75 )*
            #                                (conv_op(lab_1) - conv_op(Res_p))**2)
            loss += torch.mean(
                ((1 - lab) * (0.25 + pred_seg.unsqueeze(1)) + lab *
                 (1 - pred_seg.unsqueeze(1))**0.75) * (copl - copr)**2)

#             loss += torch.mean((conv_op(lab_1) - conv_op(Res_p))**2)

        return loss, coploss

class ThresholdLoss(nn.Module):
    def __init__(self, threshold_rate_0, threshold_rate_1):
        super(ThresholdLoss, self).__init__()
        self.threshold_rate_0 = threshold_rate_0
        self.threshold_rate_1 = threshold_rate_1

    def forward(self, output_th, pred_seg, lab):
        lab = lab.float()
        #         print(output_th.shape, pred_seg.shape, lab.shape)
        targets = lab * (pred_seg * self.threshold_rate_1 - 0.) + (1 - lab) * (
            (1 - (1 - pred_seg) * self.threshold_rate_0) + 0.)
        loss = torch.mean((output_th - targets)**2)
        return loss

if __name__ == "__main__":
    device = torch.device("cuda")
    img_root = "road_dataset/image_train/"
    lab_root = "road_dataset/label_train/"

    input_size = (256, 256)
    batch_size = 8
    NUM_CLASSES = 2
    load_path = "FCN8s_SplitNet_EdgeLossV2.pth"
    save_path = "FCN8s_SplitNet_EdgeLossV2.pth"
    EPOCH_start = 0
    EPOCH = 1000
    base_LR = 1e-4
    decay_step = 100
    decay_rate = 0.5
    wd = 1e-3
    mIU_benchmark = 0

    train_data = MapDataset(input_size,
                            img_transforms,
                            img_root,
                            lab_root,
                            rc=False)
    train_dataloader = torch.utils.data.DataLoader(train_data,
                                                   batch_size,
                                                   True,
                                                   num_workers=32,
                                                   drop_last=True)

    img_root_t = "road_dataset/image_val/"
    lab_root_t = "road_dataset/label_val/"

    test_data = MapDataset(input_size,
                           img_transforms,
                           img_root_t,
                           lab_root_t,
                           rc=False)
    test_dataloader = torch.utils.data.DataLoader(test_data,
                                                  8,
                                                  shuffle=False,
                                                  num_workers=8,
                                                  drop_last=True)

    encoder_one = VGGNet('vgg16', batch_norm=True, in_channels=3)
    net_one = bnFCN8s(encoder_one, NUM_CLASSES, 1).to(device)
    thloss = Sobelnet().to(device)
    edgenet = Edgenet().to(device)
    if load_path is not None:
        if os.path.isfile(load_path):
            try:
                checkpoint = torch.load(load_path)
                net_one.load_state_dict(checkpoint['state_one'])
                mIU_benchmark = checkpoint['mIU']
                print("Load last checkpoint OK ")
                print("mIU=", mIU_benchmark)
            except:
                print("Can't Load the checkpoint QAQ")
                mIU_benchmark = 0
        else:
            EPOCH_start = 0
            print("Can't find the checkpoint ,start train from epoch 0 ...")

    his = hl.History()
    canv = hl.Canvas()
    optm_one = torch.optim.Adam(net_one.parameters(),
                                lr=base_LR,
                                weight_decay=wd)
    scheduler_lr_one = torch.optim.lr_scheduler.StepLR(optm_one,
                                                       step_size=decay_step,
                                                       gamma=decay_rate)

    from torch.utils.tensorboard import SummaryWriter
    writer = SummaryWriter(log_dir=os.path.join("./log"))
    # Loss = nn.NLLLoss()
    Loss = nn.CrossEntropyLoss()

    l_seg = AverageValueMeter()
    l_th = AverageValueMeter()
    l_all = AverageValueMeter()
    train_miou_seg = AverageValueMeter()
    train_miou_th = AverageValueMeter()

    test_acc = AverageValueMeter()
    test_mean_iu = AverageValueMeter()
    test_acc_e = AverageValueMeter()
    test_mean_iu_e = AverageValueMeter()

    best_mIU = 0
    best_mIU_2 = mIU_benchmark
    EPOCH_start = 0
    for ep in tqdm(range(EPOCH_start, EPOCH_start + EPOCH),
                   dynamic_ncols=True):
        train_length = len(train_dataloader)
        l_seg.reset()
        l_th.reset()
        l_all.reset()
        train_miou_seg.reset()
        train_miou_th.reset()
        net_one.train()
        for i, data in tqdm(enumerate(train_dataloader),
                            total=train_length,
                            leave=False):
            if i > 10:
                break
            img, lab = data
            img, lab = img.to(device), lab.to(device)
            seg_output, th_output = net_one(img)
            loss_seg = Loss(seg_output, lab)
            seg_output_p = torch.softmax(seg_output, dim=1)[:, 1]
            loss_th, cl = thloss(th_output,
                                 seg_output_p.detach().unsqueeze(1),
                                 lab.unsqueeze(1))
            loss = loss_seg + loss_th
            optm_one.zero_grad()
            loss.backward()
            optm_one.step()

            l_seg.add(loss_seg.item())
            l_th.add(loss_th.item())
            l_all.add(loss.item())

            seg_pred = seg_output.max(dim=1)[1]
            th_pred = torch.zeros_like(seg_pred).squeeze()
            th_pred[seg_output_p > th_output.squeeze()] = 1

            _, _, train_mean_iu_i, _ = label_accuracy_score(
                lab.data.cpu().numpy(),
                seg_pred.data.cpu().numpy(), NUM_CLASSES)
            _, _, train_mean_iu_i_2, _ = label_accuracy_score(
                lab.data.cpu().numpy(),
                th_pred.data.cpu().numpy(), NUM_CLASSES)

            train_miou_seg.add(train_mean_iu_i)
            train_miou_th.add(train_mean_iu_i_2)

        scheduler_lr_one.step()  # 动态学习率

        with torch.no_grad():
            test_acc.reset()
            test_mean_iu.reset()
            test_acc_e.reset()
            test_mean_iu_e.reset()
            net_one.eval()
            for _ in tqdm(range(1), ):
                for i, data in tqdm(enumerate(test_dataloader),
                                    total=test_dataloader.__len__(),
                                    leave=False):
                    if i > 5:
                        break
                    img, lab = data
                    img, lab = img.to(device), lab.to(device)

                    seg_output, th_output = net_one(img)
                    seg_output_p = torch.softmax(seg_output, dim=1)[:, 1]
                    seg_pred = seg_output.max(dim=1)[1]
                    th_pred = torch.zeros_like(seg_pred).squeeze()
                    th_pred[seg_output_p > th_output.squeeze()] = 1

                    test_acc_i, _, test_mean_iu_i, _ = label_accuracy_score(
                        lab.data.cpu().numpy(),
                        seg_pred.data.cpu().numpy(), NUM_CLASSES)
                    test_acc_i_e, _, test_mean_iu_i_e, _ = label_accuracy_score(
                        lab.data.cpu().numpy(),
                        th_pred.data.cpu().numpy(), NUM_CLASSES)

                    test_acc.add(test_acc_i)
                    test_mean_iu.add(test_mean_iu_i)
                    test_acc_e.add(test_acc_i_e)
                    test_mean_iu_e.add(test_mean_iu_i_e)

        #画图
        writer.add_scalars(
            "TrainLoss",
            {
                "loss_seg": l_seg.value()[0],
                "loss_th": l_th.value()[0],
                "loss_all": l_all.value()[0],
            },
            ep,
        )
        writer.add_scalars(
            "miou",
            {
                "train_miou_seg": train_miou_seg.value()[0],
                "train_miou_th": train_miou_th.value()[0],
                "test_miou": test_mean_iu.value()[0],
                "test_miou2": test_mean_iu_e.value()[0],
                "best_mIU": best_mIU,
                "best_mIU2": best_mIU_2,
            },
            ep,
        )
        image_batch = torch.zeros((3, 1, input_size[0], input_size[0])).float()
        image_batch[0, 0] = seg_pred[0].float()
        image_batch[1, 0] = th_pred[0].float()
        image_batch[2, 0] = lab[0].float()
        writer.add_images("test_img",
                          image_batch,
                          global_step=ep,
                          dataformats="NCHW")

        if test_mean_iu_e.value()[0] >= best_mIU_2:
            best_mIU = train_miou_seg.value()[0]
            best_mIU_2 = test_mean_iu_e.value()[0]
            state = {
                'state_one': net_one.state_dict(),
                'EPOCH_start': ep,
                'mIU': best_mIU_2
            }
            torch.save(state, save_path)


#!/usr/bin/env python
# -*- coding: utf-8 -*-

# This operation for LoSoTo implements basic plotting

import logging
from operations_lib import *

logging.debug('Loading PLOT module.')

def make_tec_screen_plots(pp, tec_screen, residuals, station_positions,
    source_names, times, height, order, beta_val, r_0, prefix = 'frame_',
    remove_gradient=True, show_source_names=False, min_tec=None, max_tec=None):
    """Makes plots of TEC screens

    Keyword arguments:
    pp -- array of piercepoint locations
    tec_screen -- array of TEC screen values at the piercepoints
    residuals -- array of TEC screen residuals at the piercepoints
    source_names -- array of source names
    times -- array of times
    height -- height of screen (m)
    order -- order of screen (e.g., number of KL base vectors to keep)
    r_0 -- scale size of phase fluctuations (m)
    beta_val -- power-law index for phase structure function (5/3 =>
        pure Kolmogorov turbulence)
    prefix -- prefix for output file names
    remove_gradient -- fit and remove a gradient from each screen
    show_source_names -- label sources on screen plots
    min_tec -- minimum TEC value for plot range
    max_tec -- maximum TEC value for plot range
    """
    from pylab import kron, concatenate, pinv, norm, newaxis, normalize
    import matplotlib.pyplot as plt
    from mpl_toolkits.axes_grid1.inset_locator import inset_axes
    import numpy as np
    import os
    from operations.tecscreen import calc_piercepoint
    import progressbar

    root_dir = os.path.dirname(prefix)
    if root_dir == '':
        root_dir = './'
    prestr = os.path.basename(prefix) + 'screen_'
    try:
        os.makedirs(root_dir)
    except OSError:
        pass

    N_stations = station_positions.shape[0]
    N_sources = len(source_names)
    N_times = len(times)

    A = concatenate([kron(np.eye(N_sources),
        np.ones((N_stations,1))), kron(np.ones((N_sources,1)),
        np.eye(N_stations))], axis=1)

    N_piercepoints = N_sources * N_stations
    P = np.eye(N_piercepoints) - np.dot(np.dot(A, pinv(np.dot(A.T, A))), A.T)

    x, y, z = station_positions[0, :]
    east = np.array([-y, x, 0])
    east = east / norm(east)

    north = np.array([ -x, -y, (x*x + y*y)/z])
    north = north / norm(north)

    up = np.array([x ,y, z])
    up = up / norm(up)

    T = concatenate([east[:, newaxis], north[:, newaxis]], axis=1)

    pp1 = np.dot(pp[0, :, :], T)
    lower = np.amin(pp1, axis=0)
    upper = np.amax(pp1, axis=0)
    extent = upper - lower

    lower = lower - 0.05 * extent
    upper = upper + 0.05 * extent

    extent = upper - lower

    Nx = 25
    Ny = int(extent[1] / extent[0] * np.float(Nx))
    xr = np.arange(lower[0], upper[0], extent[0]/Nx)
    yr = np.arange(lower[1], upper[1], extent[1]/Ny)
    screen = np.zeros((Nx, Ny, N_times))
    gradient = np.zeros((Nx, Ny, N_times))

    residuals = residuals.transpose([0, 2, 1]).reshape(N_piercepoints, N_times)
    fitted_tec1 = tec_screen.transpose([0, 2, 1]).reshape(N_piercepoints, N_times) + residuals

    logging.info('Calculating TEC screen images...')
    pbar = progressbar.ProgressBar(maxval=N_times).start()
    ipbar = 0
    for k in range(N_times):
        D = np.resize(pp[k, :, :], (N_piercepoints, N_piercepoints, 3))
        D = np.transpose(D, (1, 0, 2)) - D
        D2 = np.sum(D**2, axis=2)
        C = -(D2 / r_0**2)**(beta_val / 2.0) / 2.0
        f = np.dot(pinv(C), tec_screen[:, k, :].reshape(N_piercepoints))
        for i, x in enumerate(xr[0: Nx]):
            for j, y in enumerate(yr[0: Ny]):
                p = calc_piercepoint(np.dot(np.array([x, y]), np.array([east, north])), up, height)
                d2 = np.sum(np.square(pp[k, :, :] - p[0]), axis=1)
                c = -(d2 / ( r_0**2 ))**(beta_val / 2.0) / 2.0
                screen[i, j, k] = np.dot(c, f)

        # Fit and remove a gradient.
        # Plot gradient in lower-left corner with its own color bar?
        if remove_gradient:
            xs, ys = np.indices(screen.shape[0:2])
            zs = screen[:, :, k]
            XYZ = []
            for xf, yf in zip(xs.flatten().tolist(), ys.flatten().tolist()):
                XYZ.append([xf, yf, zs[xf, yf]])
            XYZ = np.array(XYZ)
            a, b, c = fitPLaneLTSQ(XYZ)
            grad_plane = a * xs + b * ys + c
            gradient[:, :, k] = grad_plane
            screen[:, :, k] = screen[:, :, k] - grad_plane
            screen[:, :, k] = screen[:, :, k] - np.mean(screen[:, :, k])
            for t in range(fitted_tec1.shape[0]):
                xs_pt = np.where(np.array(xr) > pp1[t, 0])[0][0]
                ys_pt = np.where(np.array(yr) > pp1[t, 1])[0][0]
                grad_plane_pt = a * xs_pt + b * ys_pt + c
                fitted_tec1[t, k] = fitted_tec1[t, k] - grad_plane_pt
            fitted_tec1[:, k] = fitted_tec1[:, k] - np.mean(fitted_tec1[:, k])
        pbar.update(ipbar)
        ipbar += 1
    pbar.finish()
    if min_tec is None:
        vmin = np.min([np.amin(screen), np.amin(fitted_tec1)])
    else:
        vmin = min_tec
    if max_tec is None:
        vmax = np.max([np.amax(screen), np.amax(fitted_tec1)])
    else:
        vmax = max_tec

    logging.info('Plotting TEC screens...')
    fig, ax = plt.subplots(figsize=[7, 7])
    pbar = progressbar.ProgressBar(maxval=N_times).start()
    ipbar = 0
    for k in range(N_times):
        plt.clf()
        im = plt.imshow(screen.transpose([1, 0, 2])[:, :, k],
            cmap = plt.cm.jet,
            origin = 'lower',
            interpolation = 'nearest',
            extent = (xr[0]/1000.0, xr[-1]/1000.0, yr[0]/1000.0, yr[-1]/1000.0),
            vmin=vmin, vmax=vmax)

        sm = plt.cm.ScalarMappable(cmap=plt.cm.jet,
            norm=normalize(vmin=vmin, vmax=vmax))
        sm._A = []
        cbar = plt.colorbar(im)
        cbar.set_label('TECU', rotation=270)

        x = []
        y = []
        s = []
        c = []
        for j in range(fitted_tec1.shape[0]):
            x.append(pp1[j, 0] / 1000.0)
            y.append(pp1[j, 1] / 1000.0)
            xs = np.where(np.array(xr) > pp1[j, 0])[0][0]
            ys = np.where(np.array(yr) > pp1[j, 1])[0][0]
            fit_screen_diff = abs(fitted_tec1[j, k] - screen[xs, ys, k])
            s.append(max(20*fit_screen_diff/0.01, 10))
            c.append(sm.to_rgba(fitted_tec1[j, k]))

        plt.scatter(x, y, s=s, c=c)
        if show_source_names:
            labels = source_names
            for label, xl, yl in zip(labels, x[0::N_stations], y[0::N_stations]):
                plt.annotate(
                    label,
                    xy = (xl, yl), xytext = (-2, 2),
                    textcoords = 'offset points', ha = 'right', va = 'bottom')

        plt.title('Screen {0}'.format(k))
        plt.xlim(xr[-1]/1000.0, xr[0]/1000.0)
        plt.ylim(yr[0]/1000.0, yr[-1]/1000.0)
        plt.xlabel('Projected Distance along RA (km)')
        plt.ylabel('Projected Distance along Dec (km)')

        if remove_gradient:
            axins = inset_axes(ax, width="15%", height="10%", loc=2)
            axins.imshow(gradient.transpose([1, 0, 2])[:, : ,k],
                cmap = plt.cm.jet,
                origin = 'lower',
                interpolation = 'nearest',
                extent = (xr[0]/1000.0, xr[-1]/1000.0, yr[0]/1000.0, yr[-1]/1000.0),
                vmin=vmin, vmax=vmax)
            plt.xticks(visible=False)
            plt.yticks(visible=False)
            axins.set_xlim(xr[-1]/1000.0, xr[0]/1000.0)
            axins.set_ylim(yr[0]/1000.0, yr[-1]/1000.0)

        plt.savefig(root_dir+'/'+prestr+'frame%0.3i.png' % k)
        pbar.update(ipbar)
        ipbar += 1
    pbar.finish()
    plt.close(fig)

def fitPLaneLTSQ(XYZ):
    """Fits a plane to an XYZ point cloud

    Returns (a, b, c), where Z = aX + bY + c

    Keyword arguments:
    XYZ -- point cloud
    """
    import numpy as np
    [rows, cols] = XYZ.shape
    G = np.ones((rows, 3))
    G[:, 0] = XYZ[:, 0]  #X
    G[:, 1] = XYZ[:, 1]  #Y
    Z = XYZ[:, 2]
    (a, b, c), resid, rank, s = np.linalg.lstsq(G, Z)
    return (a, b, c)

def run( step, parset, H ):

    import matplotlib.pyplot as plt
    import numpy as np
    from h5parm import solFetcher

    solsets = getParSolsets( step, parset, H )
    soltabs = getParSoltabs( step, parset, H )
    ants = getParAxis( step, parset, H, 'ant' )
    pols = getParAxis( step, parset, H, 'pol' )
    dirs = getParAxis( step, parset, H, 'dir' )

    plotType = parset.getString('.'.join(["LoSoTo.Steps", step, "PlotType"]), '' )
    axesToPlot = parset.getStringVector('.'.join(["LoSoTo.Steps", step, "Axes"]), '' )
    minZ, maxZ = parset.getDoubleVector('.'.join(["LoSoTo.Steps", step, "MinMax"]), [0,0] )
    prefix = parset.getString('.'.join(["LoSoTo.Steps", step, "Prefix"]), '' )

    if plotType.lower() in ['1d', '2d']:
        for soltab in openSoltabs( H, soltabs ):

            sf = solFetcher(soltab)
            logging.info("Plotting soltab: "+soltab._v_name)

            sf.setSelection(ant=ants, pol=pols, dir=dirs)

            # some checks
            for axis in axesToPlot:
                if axis not in sf.getAxesNames():
                    logging.error('Axis \"'+axis+'\" not found.')
                    return 1

            if (len(axesToPlot) != 2 and plotType == '2D') or \
               (len(axesToPlot) != 1 and plotType == '1D'):
                logging.error('Wrong number of axes.')
                return 1

            for vals, coord in sf.getValuesIter(returnAxes=axesToPlot):
                # TODO: implement flag control, using different color?

                title = ''
                for axis in coord:
                    if axis in axesToPlot: continue
                    title += str(coord[axis])+'_'
                title = title[:-1]

                if plotType == '2D':
                    fig = plt.figure()
                    ax = plt.subplot(111)
                    plt.title(title)
                    plt.ylabel(axesToPlot[0])
                    plt.xlabel(axesToPlot[1])
                    p = ax.imshow(coord[axesToPlot[1]], coord[axesToPlot[0]], vals)
                    if not (minZ == 0 and maxZ == 0):
                        plt.zlim(zmin=minZ, zmax=maxZ)
                    plt.savefig(title+'.png')
                    logging.info("Saving "+prefix+title+'.png')

                if plotType == '1D':
                    fig = plt.figure()
                    ax = plt.subplot(111)
                    plt.title(title)
                    plt.ylabel(sf.getType())
                    if not (minZ == 0 and maxZ == 0):
                        plt.ylim(ymin=minZ, ymax=maxZ)
                    plt.xlabel(axesToPlot[0])
                    p = ax.plot(coord[axesToPlot[0]], vals)
                    plt.savefig(prefix+title+'.png')
                    logging.info("Saving "+prefix+title+'.png')

    elif plotType.lower() == 'tecscreen':
        # Plot various TEC-screen properties

        for st_scr in openSoltabs(H, soltabs):

            # Check if soltab is a tecscreen table
            full_name = st_scr._v_parent._v_name + '/' + st_scr._v_name
            if st_scr._v_title != 'tecscreen':
                logging.warning('Solution table {0} is not a tecscreen solution '
                    'table. Skipping.'.format(full_name))
                continue
            logging.info('Using input solution table: {0}'.format(full_name))

            # Plot TEC screens as images
            solset = st_scr._v_parent
            sf_scr = solFetcher(st_scr)
            r, axis_vals = sf_scr.getValues()
            source_names = axis_vals['dir']
            station_names = axis_vals['ant']
            station_dict = H.getAnt(solset)
            station_positions = []
            for station in station_names:
                station_positions.append(station_dict[station])
            times = axis_vals['time']

            tec_screen, axis_vals = sf_scr.getValues()
            times = axis_vals['time']
            residuals = sf_scr.getValues(weight=True, retAxesVals=False)
            height = st_scr._v_attrs['height']
            order = st_scr._v_attrs['order']
            beta_val = st_scr._v_attrs['beta']
            r_0 = st_scr._v_attrs['r_0']
            pp = sf_scr.t.piercepoint

            if (minZ == 0 and maxZ == 0):
                min_tec = None
                max_tec = None
            else:
                min_tec = minZ
                max_tec = maxZ

            make_tec_screen_plots(pp, tec_screen, residuals,
                np.array(station_positions), np.array(source_names), times,
                height, order, beta_val, r_0, prefix=prefix,
                remove_gradient=True, show_source_names=False, min_tec=min_tec,
                max_tec=max_tec)

    return 0


import os
import unittest
from typing import List

from cdm.objectmodel import CdmCorpusDefinition, CdmManifestDefinition, CdmEntityDefinition
from cdm.utilities import ResolveOptions, AttributeResolutionDirectiveSet
from tests.cdm.projection.attribute_context_util import AttributeContextUtil
from tests.common import async_test, TestHelper, TestUtils

class ProjectionFKTest(unittest.TestCase):
    res_opts_combinations = [
        [],
        ['referenceOnly'],
        ['normalized'],
        ['structured'],
        ['referenceOnly', 'normalized'],
        ['referenceOnly', 'structured'],
        ['normalized', 'structured'],
        ['referenceOnly', 'normalized', 'structured']
    ]

    # The path between TestDataPath and TestName.
    tests_subpath = os.path.join('Cdm', 'Projection', 'TestProjectionFK')

    @async_test
    async def test_entity_attribute(self):
        test_name = 'test_entity_attribute'
        entity_name = 'SalesEntityAttribute'

        for res_opt in self.res_opts_combinations:
            await self._load_entity_for_resolution_option_and_save(test_name, entity_name, res_opt)

    @async_test
    async def test_entity_attribute_proj(self):
        test_name = 'test_entity_attribute_proj'
        entity_name = 'SalesEntityAttribute'

        for res_opt in self.res_opts_combinations:
            await self._load_entity_for_resolution_option_and_save(test_name, entity_name, res_opt)

    @async_test
    async def test_source_with_EA(self):
        test_name = 'test_source_with_EA'
        entity_name = 'SalesSourceWithEA'

        for res_opt in self.res_opts_combinations:
            await self._load_entity_for_resolution_option_and_save(test_name, entity_name, res_opt)

    @async_test
    async def test_source_with_EA_proj(self):
        test_name = 'test_source_with_EA_proj'
        entity_name = 'SalesSourceWithEA'

        for res_opt in self.res_opts_combinations:
            await self._load_entity_for_resolution_option_and_save(test_name, entity_name, res_opt)

    @async_test
    async def test_group_FK(self):
        test_name = 'test_group_FK'
        entity_name = 'SalesGroupFK'

        for res_opt in self.res_opts_combinations:
            await self._load_entity_for_resolution_option_and_save(test_name, entity_name, res_opt)

    @async_test
    async def test_group_FK_proj(self):
        test_name = 'test_group_FK_proj'
        entity_name = 'SalesGroupFK'

        for res_opt in self.res_opts_combinations:
            await self._load_entity_for_resolution_option_and_save(test_name, entity_name, res_opt)

    @async_test
    async def test_nested_FK_proj(self):
        test_name = 'test_nested_FK_proj'
        entity_name = 'SalesNestedFK'

        for res_opt in self.res_opts_combinations:
            await self._load_entity_for_resolution_option_and_save(test_name, entity_name, res_opt)

    @async_test
    async def test_polymorphic(self):
        test_name = 'test_polymorphic'
        entity_name = 'PersonPolymorphicSource'

        for res_opt in self.res_opts_combinations:
            await self._load_entity_for_resolution_option_and_save(test_name, entity_name, res_opt)

    @async_test
    async def test_polymorphic_proj(self):
        test_name = 'test_polymorphic_proj'
        entity_name = 'PersonPolymorphicSource'

        for res_opt in self.res_opts_combinations:
            await self._load_entity_for_resolution_option_and_save(test_name, entity_name, res_opt)

    @async_test
    async def test_polymorphic_FK_proj(self):
        test_name = 'test_polymorphic_FK_proj'
        entity_name = 'PersonPolymorphicSourceFK'

        for res_opt in self.res_opts_combinations:
            await self._load_entity_for_resolution_option_and_save(test_name, entity_name, res_opt)

    @async_test
    async def test_array_source(self):
        test_name = 'test_array_source'
        entity_name = 'SalesArraySource'

        for res_opt in self.res_opts_combinations:
            await self._load_entity_for_resolution_option_and_save(test_name, entity_name, res_opt)

    @async_test
    async def test_array_source_proj(self):
        test_name = 'test_array_source_proj'
        entity_name = 'SalesArraySource'

        for res_opt in self.res_opts_combinations:
            await self._load_entity_for_resolution_option_and_save(test_name, entity_name, res_opt)

    @async_test
    async def test_foreign_key(self):
        test_name = 'test_foreign_key'
        entity_name = 'SalesForeignKey'

        for res_opt in self.res_opts_combinations:
            await self._load_entity_for_resolution_option_and_save(test_name, entity_name, res_opt)

    @async_test
    async def test_foreign_key_proj(self):
        test_name = 'test_foreign_key_proj'
        entity_name = 'SalesForeignKey'

        for res_opt in self.res_opts_combinations:
            await self._load_entity_for_resolution_option_and_save(test_name, entity_name, res_opt)

    @async_test
    async def test_foreign_key_always(self):
        test_name = 'test_foreign_key_always'
        entity_name = 'SalesForeignKeyAlways'

        for res_opt in self.res_opts_combinations:
            await self._load_entity_for_resolution_option_and_save(test_name, entity_name, res_opt)

    @async_test
    async def test_composite_key_proj(self):
        self.maxDiff = None
        test_name = 'test_composite_key_proj'
        entity_name = 'SalesCompositeKey'

        for res_opt in self.res_opts_combinations:
            await self._load_entity_for_resolution_option_and_save(test_name, entity_name, res_opt)

    async def _load_entity_for_resolution_option_and_save(self, test_name: str, entity_name: str, res_opts: List[str]):
        corpus = TestHelper.get_local_corpus(self.tests_subpath, test_name)
        manifest = await corpus.fetch_object_async('local:/default.manifest.cdm.json')

        expected_output_path = TestHelper.get_expected_output_folder_path(self.tests_subpath, test_name)
        file_name_suffix = TestUtils.get_resolution_option_name_suffix(res_opts)

        ent_sales_foreign_key_projection = await corpus.fetch_object_async('local:/{}.cdm.json/{}'.format(entity_name, entity_name), manifest)
        self.assertIsNotNone(ent_sales_foreign_key_projection)
        resolved_sales_foreign_key_projection = await self._save_resolved(corpus, manifest, test_name, ent_sales_foreign_key_projection, res_opts)
        self.assertIsNotNone(resolved_sales_foreign_key_projection)
        AttributeContextUtil.validate_attribute_context(self, corpus, expected_output_path, '{}{}'.format(entity_name, file_name_suffix), resolved_sales_foreign_key_projection)

    async def _save_resolved(
        self,
        corpus: 'CdmCorpusDefinition',
        manifest: 'CdmManifestDefinition',
        test_name: str,
        input_entity: 'CdmEntityDefinition',
        resolution_options: List[str]
    ) -> 'CdmEntityDefinition':
        ro_hash_set = set()
        for i in range(len(resolution_options)):
            ro_hash_set.add(resolution_options[i])

        file_name_suffix = TestUtils.get_resolution_option_name_suffix(resolution_options)

        resolved_entity_name = 'Resolved_{}{}.cdm.json'.format(input_entity.entity_name, file_name_suffix)

        ro = ResolveOptions(input_entity.in_document, directives=AttributeResolutionDirectiveSet(ro_hash_set))

        resolved_folder = corpus.storage.fetch_root_folder('output')
        resolved_entity = await input_entity.create_resolved_entity_async(resolved_entity_name, ro, resolved_folder)

        return resolved_entity


import PIL
import torch
import torchvision
from torchvision import models, transforms

import numpy as np
import matplotlib.pyplot as plt

def imshow_v2(tensor, labels, nclass):
    """

    :param tensor:
    :param labels:
    :param nclass:

    """
    plt.figure(figsize=(10,6))
    for i, img in enumerate(tensor):
        plt.subplot(4,4,i+1)
        plt.axis('off')
        # print labels
        plt.title(nclass[labels[i]])
        img = img / 2 + 0.5     # unnormalize
        npimg = img.numpy()
        plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.tight_layout()
    plt.show()

def grayscale_to_pil(tensor, nrow, padding=2, pad_value=0):
    # make the mini-batch of images into a grid
    nmaps = tensor.size(0)
    xmaps = min(nrow, nmaps)
    ymaps = int(math.ceil(float(nmaps) / xmaps))
    height, width = int(tensor.size(1) + padding), int(tensor.size(2) + padding)
    num_channels = 1
    grid = tensor.new_full((num_channels, height * ymaps + padding, width * xmaps + padding), pad_value)
    print(grid.shape)
    k = 0
    for y in range(ymaps):
        for x in range(xmaps):
            if k >= nmaps:
                break
            # Tensor.copy_() is a valid method but seems to be missing from the stubs
            # https://pytorch.org/docs/stable/tensors.html#torch.Tensor.copy_
            grid.narrow(1, y * height + padding, height - padding).narrow(  # type: ignore[attr-defined]
                2, x * width + padding, width - padding
            ).copy_(tensor[k])
            k = k + 1
    return grid

def get_inter_act_map(x, model):
    with torch.no_grad():
        for name, mod in model.named_children():
            x = mod(x.cuda())
            if x.size(1) == 3:
                yield name, make_grid(x, 16).cpu().permute(1,2,0)
            else:
                yield name, grayscale_to_pil(intmap, 8)[0].cpu()

if __name__ == '__main__':
    NUM_CLASSES = 10
    device = torch.device('cpu')

    # Define categories of dataset and point to direction to generate data flow to train.
    classes = ['airplane', 'bird', 'car', 'cat', 'deer', 'dog', 'horse', 'monkey', 'ship', 'truck']
    samples = torchvision.datasets.ImageFolder('../dataset/samples',
                                               transform=transforms.Compose([
                                                   transforms.Resize(100),
                                                   transforms.RandomCrop(96, 4),
                                                   transforms.RandomRotation((-180,180), resample=PIL.Image.BILINEAR),
                                                   transforms.ToTensor()
                                               ]))
    sample_loader = torch.utils.data.DataLoader(samples, batch_size=1,
                                                shuffle=False, num_workers=2)

    # Load a existed net into CPU device to visualize
    # Loading net shouldn't to require NetTrainer module
    checkpoint = torch.load('../classifier/resnet18-STL10-checkpoint.pth', map_location=device)
    net = models.resnet18(num_classes=NUM_CLASSES)
    net.load_state_dict(checkpoint['model_state_dict'])

    net.eval()

    dataiter = iter(sample_loader)
    data, _ = dataiter.next()
    # plt.imshow(np.transpose(data[0].numpy(), (1, 2, 0)))
    # plt.show()
    # imshow_v2(data, net(data).argmax(dim=1), classes)

    for name, layer in net._modules.items():
        pass