text:
import warnings

import numpy
from matplotlib import pyplot
from matplotlib import dates
from matplotlib import gridspec
import seaborn
import pandas

from wqio import utils
from wqio import viz
from wqio import validate

from pandas.plotting import register_matplotlib_converters

register_matplotlib_converters()

SEC_PER_MINUTE = 60.0
MIN_PER_HOUR = 60.0
HOUR_PER_DAY = 24.0
SEC_PER_HOUR = SEC_PER_MINUTE * MIN_PER_HOUR
SEC_PER_DAY = SEC_PER_HOUR * HOUR_PER_DAY

def _wet_first_row(df, wetcol, diffcol):
    # make sure that if the first record is associated with the first
    # storm if it's wet
    firstrow = df.iloc[0]
    if firstrow[wetcol]:
        df.loc[firstrow.name, diffcol] = 1

    return df

def _wet_window_diff(is_wet, ie_periods):
    return (
        is_wet.rolling(int(ie_periods), min_periods=1)
        .apply(lambda window: window.any(), raw=False)
        .diff()
    )

def parse_storm_events(
    data,
    intereventHours,
    outputfreqMinutes,
    precipcol=None,
    inflowcol=None,
    outflowcol=None,
    baseflowcol=None,
    stormcol="storm",
    debug=False,
):
    """Parses the hydrologic data into distinct storms.

    In this context, a storm is defined as starting whenever the
    hydrologic records shows non-zero precipitation or [in|out]flow
    from the BMP after a minimum inter-event dry period duration
    specified in the the function call. The storms ends the observation
    *after* the last non-zero precipitation or flow value.

    Parameters
    ----------
    data : pandas.DataFrame
    intereventHours : float
        The Inter-Event dry duration (in hours) that classifies the
        next hydrlogic activity as a new event.
    precipcol : string, optional (default = None)
        Name of column in `hydrodata` containing precipiation data.
    inflowcol : string, optional (default = None)
        Name of column in `hydrodata` containing influent flow data.
    outflowcol : string, optional (default = None)
        Name of column in `hydrodata` containing effluent flow data.
    baseflowcol : string, optional (default = None)
        Name of column in `hydrodata` containing boolean indicating
        which records are considered baseflow.
    stormcol : string (default = 'storm')
        Name of column in `hydrodata` indentifying distinct storms.
    debug : bool (default = False)
        If True, diagnostic columns will not be dropped prior to
        returning the dataframe of parsed_storms.

    Writes
    ------
    None

    Returns
    -------
    parsed_storms : pandas.DataFrame
        Copy of the origin `hydrodata` DataFrame, but resampled to a
        fixed frequency, columns possibly renamed, and a `storm` column
        added to denote the storm to which each record belongs. Records
        where `storm` == 0 are not a part of any storm.

    """

    # pull out the rain and flow data
    if precipcol is None:
        precipcol = "precip"
        data.loc[:, precipcol] = numpy.nan

    if inflowcol is None:
        inflowcol = "inflow"
        data.loc[:, inflowcol] = numpy.nan

    if outflowcol is None:
        outflowcol = "outflow"
        data.loc[:, outflowcol] = numpy.nan

    if baseflowcol is None:
        baseflowcol = "baseflow"
        data.loc[:, baseflowcol] = False

    # bool column where True means there's rain or flow of some kind
    water_columns = [inflowcol, outflowcol, precipcol]
    cols_to_use = water_columns + [baseflowcol]

    agg_dict = {
        precipcol: numpy.sum,
        inflowcol: numpy.mean,
        outflowcol: numpy.mean,
        baseflowcol: numpy.any,
    }

    freq = pandas.offsets.Minute(outputfreqMinutes)
    ie_periods = int(MIN_PER_HOUR / freq.n * intereventHours)

    # periods between storms are where the cumulative number
    # of storms that have ended are equal to the cumulative
    # number of storms that have started.
    # Stack Overflow: http://tinyurl.com/lsjkr9x
    res = (
        data.resample(freq)
        .agg(agg_dict)
        .loc[:, lambda df: df.columns.isin(cols_to_use)]
        .assign(
            __wet=lambda df: numpy.any(df[water_columns] > 0, axis=1) & ~df[baseflowcol]
        )
        .assign(__windiff=lambda df: _wet_window_diff(df["__wet"], ie_periods))
        .pipe(_wet_first_row, "__wet", "__windiff")
        .assign(__event_start=lambda df: df["__windiff"] == 1)
        .assign(__event_end=lambda df: df["__windiff"].shift(-1 * ie_periods) == -1)
        .assign(__storm=lambda df: df["__event_start"].cumsum())
        .assign(
            storm=lambda df: numpy.where(
                df["__storm"] == df["__event_end"].shift(2).cumsum(),
                0,  # inter-event periods marked as zero
                df["__storm"],  # actual events keep their number
            )
        )
    )

    if not debug:
        res = res.loc[:, res.columns.map(lambda c: not c.startswith("__"))]

    return res

class Storm(object):
    """ Object representing a storm event

    Parameters
    ----------
    dataframe : pandas.DataFrame
        A datetime-indexed Dataframe containing all of the hydrologic
        data and am interger column indentifying distinct storms.
    stormnumber : int
        The storm we care about.
    precipcol, inflowcol, outflow, tempcol, stormcol : string, optional
        Names for columns representing each hydrologic quantity.
    freqMinutes : float (default = 5)
        The time period, in minutes, between observations.
    volume_conversion : float, optional (default = 1)
        Conversion factor to go from flow to volume for a single
        observation.

    """

    # TODO: rename freqMinutes to periodMinutes
    def __init__(
        self,
        dataframe,
        stormnumber,
        precipcol="precip",
        inflowcol="inflow",
        outflowcol="outflow",
        tempcol="temp",
        stormcol="storm",
        freqMinutes=5,
        volume_conversion=1,
    ):

        self.inflowcol = inflowcol
        self.outflowcol = outflowcol
        self.precipcol = precipcol
        self.tempcol = tempcol
        self.stormnumber = stormnumber
        self.freqMinutes = freqMinutes
        self.volume_conversion = volume_conversion * SEC_PER_MINUTE * self.freqMinutes

        # basic data
        self.data = dataframe[dataframe[stormcol] == self.stormnumber].copy()
        self.hydrofreq_label = "{0} min".format(self.freqMinutes)

        # tease out start/stop info
        self.start = self.data.index[0]
        self.end = self.data.index[-1]
        self._season = utils.getSeason(self.start)

        # storm duration (hours)
        duration = self.end - self.start
        self.duration_hours = duration.total_seconds() / SEC_PER_HOUR

        # antecedent dry period (hours)
        if self.stormnumber > 1:
            prev_storm_mask = dataframe[stormcol] == self.stormnumber - 1
            previous_end = dataframe[prev_storm_mask].index[-1]
            antecedent_timedelta = self.start - previous_end
            self.antecedent_period_days = (
                antecedent_timedelta.total_seconds() / SEC_PER_DAY
            )
        else:
            self.antecedent_period_days = numpy.nan

        # quantities
        self._precip = None
        self._inflow = None
        self._outflow = None

        # starts and stop
        self._precip_start = None
        self._precip_end = None
        self._inflow_start = None
        self._inflow_end = None
        self._outflow_start = None
        self._outflow_end = None

        # peaks
        self._peak_precip_intensity = None
        self._peak_inflow = None
        self._peak_outflow = None

        # times of peaks
        self._peak_precip_intensity_time = None
        self._peak_inflow_time = None
        self._peak_outflow_time = None
        self._peak_lag_hours = None

        # centroids
        self._centroid_precip_time = None
        self._centroid_inflow_time = None
        self._centroid_outflow_time = None
        self._centroid_lag_hours = None

        # totals
        self._total_precip_depth = None
        self._total_inflow_volume = None
        self._total_outflow_volume = None

        self.meta = {
            self.outflowcol: {
                "name": "Flow (calculated, L/s)",
                "ylabel": "Effluent flow (L/s)",
                "color": "CornFlowerBlue",
                "linewidth": 1.5,
                "alpha": 0.5,
                "ymin": 0,
            },
            self.inflowcol: {
                "name": "Inflow (estimated, L/s)",
                "ylabel": "Estimated influent flow (L/s)",
                "color": "Maroon",
                "linewidth": 1.5,
                "alpha": 0.5,
                "ymin": 0,
            },
            self.precipcol: {
                "name": "Precip (mm)",
                "ylabel": "%s Precip.\nDepth (mm)" % self.hydrofreq_label,
                "color": "DarkGreen",
                "linewidth": 1.5,
                "alpha": 0.4,
                "ymin": 0,
            },
            self.tempcol: {
                "name": "Air Temp (deg C)",
                "ylabel": "Air Temperature (deg. C)",
                "color": "DarkGoldenRod",
                "linewidth": 1.5,
                "alpha": 0.5,
                "ymin": None,
            },
        }

        self._summary_dict = None

    @property
    def precip(self):
        if self._precip is None:
            if self.precipcol is not None:
                self._precip = self.data[self.data[self.precipcol] > 0][self.precipcol]
            else:
                self._precip = numpy.array([])
        return self._precip

    @property
    def inflow(self):
        if self._inflow is None:
            if self.inflowcol is not None:
                self._inflow = self.data[self.data[self.inflowcol] > 0][self.inflowcol]
            else:
                self._inflow = numpy.array([])
        return self._inflow

    @property
    def outflow(self):
        if self._outflow is None:
            if self.outflowcol is not None:
                self._outflow = self.data[self.data[self.outflowcol] > 0][
                    self.outflowcol
                ]
            else:
                self._outflow = numpy.array([])
        return self._outflow

    @property
    def has_precip(self):
        return self.precip.shape[0] > 0

    @property
    def has_inflow(self):
        return self.inflow.shape[0] > 0

    @property
    def has_outflow(self):
        return self.outflow.shape[0] > 0

    @property
    def season(self):
        return self._season

    @season.setter
    def season(self, value):
        self._season = value

    # starts and stops
    @property
    def precip_start(self):
        if self._precip_start is None and self.has_precip:
            self._precip_start = self._get_event_time(self.precipcol, "start")
        return self._precip_start

    @property
    def precip_end(self):
        if self._precip_end is None and self.has_precip:
            self._precip_end = self._get_event_time(self.precipcol, "end")
        return self._precip_end

    @property
    def inflow_start(self):
        if self._inflow_start is None and self.has_inflow:
            self._inflow_start = self._get_event_time(self.inflowcol, "start")
        return self._inflow_start

    @property
    def inflow_end(self):
        if self._inflow_end is None and self.has_inflow:
            self._inflow_end = self._get_event_time(self.inflowcol, "end")
        return self._inflow_end

    @property
    def outflow_start(self):
        if self._outflow_start is None and self.has_outflow:
            self._outflow_start = self._get_event_time(self.outflowcol, "start")
        return self._outflow_start

    @property
    def outflow_end(self):
        if self._outflow_end is None and self.has_outflow:
            self._outflow_end = self._get_event_time(self.outflowcol, "end")
        return self._outflow_end

    @property
    def _peak_depth(self):
        if self.has_precip:
            return self.precip.max()

    @property
    def peak_precip_intensity(self):
        if self._peak_precip_intensity is None and self.has_precip:
            self._peak_precip_intensity = (
                self._peak_depth * MIN_PER_HOUR / self.freqMinutes
            )
        return self._peak_precip_intensity

    @property
    def peak_inflow(self):
        if self._peak_inflow is None and self.has_inflow:
            self._peak_inflow = self.inflow.max()
        return self._peak_inflow

    @property
    def peak_outflow(self):
        if self._peak_outflow is None and self.has_outflow:
            self._peak_outflow = self.outflow.max()
        return self._peak_outflow

    @property
    def total_precip_depth(self):
        if self._total_precip_depth is None and self.has_precip:
            self._total_precip_depth = self.data[self.precipcol].sum()
        return self._total_precip_depth

    @property
    def total_inflow_volume(self):
        if self._total_inflow_volume is None and self.has_inflow:
            self._total_inflow_volume = (
                self.data[self.inflowcol].sum() * self.volume_conversion
            )
        return self._total_inflow_volume

    @property
    def total_outflow_volume(self):
        if self._total_outflow_volume is None and self.has_outflow:
            self._total_outflow_volume = (
                self.data[self.outflowcol].sum() * self.volume_conversion
            )
        return self._total_outflow_volume

    @property
    def centroid_precip_time(self):
        if self._centroid_precip_time is None and self.has_precip:
            self._centroid_precip_time = self._compute_centroid(self.precipcol)
        return self._centroid_precip_time

    @property
    def centroid_inflow_time(self):
        if self._centroid_inflow_time is None and self.has_inflow:
            self._centroid_inflow_time = self._compute_centroid(self.inflowcol)
        return self._centroid_inflow_time

    @property
    def centroid_outflow_time(self):
        if self._centroid_outflow_time is None and self.has_outflow:
            self._centroid_outflow_time = self._compute_centroid(self.outflowcol)
        return self._centroid_outflow_time

    @property
    def centroid_lag_hours(self):
        if (
            self._centroid_lag_hours is None
            and self.centroid_outflow_time is not None
            and self.centroid_inflow_time is not None
        ):
            self._centroid_lag_hours = (
                self.centroid_outflow_time - self.centroid_inflow_time
            ).total_seconds() / SEC_PER_HOUR
        return self._centroid_lag_hours

    @property
    def peak_precip_intensity_time(self):
        if self._peak_precip_intensity_time is None and self.has_precip:
            PI_selector = self.data[self.precipcol] == self._peak_depth
            self._peak_precip_intensity_time = self.data[PI_selector].index[0]
        return self._peak_precip_intensity_time

    @property
    def peak_inflow_time(self):
        if self._peak_inflow_time is None and self.has_inflow:
            PInf_selector = self.data[self.inflowcol] == self.peak_inflow
            self._peak_inflow_time = self.data[PInf_selector].index[0]
        return self._peak_inflow_time

    @property
    def peak_outflow_time(self):
        if self._peak_outflow_time is None and self.has_outflow:
            PEff_selector = self.data[self.outflowcol] == self.peak_outflow
            if PEff_selector.sum() > 0:
                self._peak_outflow_time = self.data[PEff_selector].index[0]
        return self._peak_outflow_time

    @property
    def peak_lag_hours(self):
        if (
            self._peak_lag_hours is None
            and self.peak_outflow_time is not None
            and self.peak_inflow_time is not None
        ):

            time_delta = self.peak_outflow_time - self.peak_inflow_time
            self._peak_lag_hours = time_delta.total_seconds() / SEC_PER_HOUR
        return self._peak_lag_hours

    @property
    def summary_dict(self):
        if self._summary_dict is None:
            self._summary_dict = {
                "Storm Number": self.stormnumber,
                "Antecedent Days": self.antecedent_period_days,
                "Start Date": self.start,
                "End Date": self.end,
                "Duration Hours": self.duration_hours,
                "Peak Precip Intensity": self.peak_precip_intensity,
                "Total Precip Depth": self.total_precip_depth,
                "Total Inflow Volume": self.total_inflow_volume,
                "Peak Inflow": self.peak_inflow,
                "Total Outflow Volume": self.total_outflow_volume,
                "Peak Outflow": self.peak_outflow,
                "Peak Lag Hours": self.peak_lag_hours,
                "Centroid Lag Hours": self.centroid_lag_hours,
                "Season": self.season,
            }

        return self._summary_dict

    def is_small(self, minprecip=0.0, mininflow=0.0, minoutflow=0.0):
        """ Determines whether a storm can be considered "small".

        Parameters
        ----------
        minprecip, mininflow, minoutflow : float, optional (default = 0)
            The minimum amount of each hydrologic quantity below which a
            storm can be considered "small".

        Returns
        -------
        storm_is_small : bool
            True if the storm is considered small.

        """

        storm_is_small = (
            (
                self.total_precip_depth is not None
                and self.total_precip_depth < minprecip
            )
            or (
                self.total_inflow_volume is not None
                and self.total_inflow_volume < mininflow
            )
            or (
                self.total_outflow_volume is not None
                and self.total_outflow_volume < minoutflow
            )
        )
        return storm_is_small

    def _get_event_time(self, column, bound):
        index_map = {"start": 0, "end": -1}
        quantity = self.data[self.data[column] > 0]
        if quantity.shape[0] == 0:
            warnings.warn("Storm has no {}".format(column), UserWarning)
        else:
            return quantity.index[index_map[bound]]

    def _get_max_quantity(self, column):
        return self.data[column].max()

    def _compute_centroid(self, column):
        # ordinal time index of storm
        time_idx = [
            dates.date2num(idx.to_pydatetime()) for idx in self.data.index.tolist()
        ]

        centroid = numpy.sum(self.data[column] * time_idx) / numpy.sum(
            self.data[column]
        )

        if numpy.isnan(centroid):
            return None
        else:
            return pandas.Timestamp(dates.num2date(centroid)).tz_convert(None)

    def _plot_centroids(self, ax, yfactor=0.5):

        artists = []
        labels = []
        y_val = yfactor * ax.get_ylim()[1]

        if self.centroid_precip is not None:
            ax.plot(
                [self.centroid_precip],
                [y_val],
                color="DarkGreen",
                marker="o",
                linestyle="none",
                zorder=20,
                markersize=6,
            )
            artists.append(
                pyplot.Line2D(
                    [0],
                    [0],
                    marker=".",
                    markersize=6,
                    linestyle="none",
                    color="DarkGreen",
                )
            )
            labels.append("Precip. centroid")

        if self.centroid_flow is not None:
            ax.plot(
                [self.centroid_flow],
                [y_val],
                color="CornflowerBlue",
                marker="s",
                linestyle="none",
                zorder=20,
                markersize=6,
            )
            artists.append(
                pyplot.Line2D(
                    [0],
                    [0],
                    marker="s",
                    markersize=6,
                    linestyle="none",
                    color="CornflowerBlue",
                )
            )
            labels.append("Effluent centroid")

        if self.centroid_precip is not None and self.centroid_flow is not None:
            ax.annotate(
                "",
                (self.centroid_flow, y_val),
                arrowprops=dict(arrowstyle="-|>"),
                xytext=(self.centroid_precip, y_val),
            )

        return artists, labels

    def plot_hydroquantity(
        self, quantity, ax=None, label=None, otherlabels=None, artists=None
    ):
        """ Draws a hydrologic quantity to a matplotlib axes.

        Parameters
        ----------
        quantity : string
            Column name of the quantity you want to plot.
        ax : matplotlib axes object, optional
            The axes on which the data will be plotted. If None, a new
            one will be created.
        label : string, optional
            How the series should be labeled in the figure legend.
        otherlabels : list of strings, optional
            A list of other legend labels that have already been plotted
            to ``ax``. If provided, ``label`` will be appended. If not
            provided, and new list will be created.
        artists : list of matplotlib artists, optional
            A list of other legend items that have already been plotted
            to ``ax``. If provided, the artist created will be appended.
            If not provided, and new list will be created.

        Returns
        -------
        fig : matplotlib.Figure
            The figure containing the plot.
        labels : list of strings
            Labels to be included in a legend for the figure.
        artists : list of matplotlib artists
            Symbology for the figure legend.

        """

        # setup the figure
        fig, ax = validate.axes(ax)

        if label is None:
            label = quantity

        # select the plot props based on the column
        try:
            meta = self.meta[quantity]
        except KeyError:
            raise KeyError("{} not available".format(quantity))

        # plot the data
        self.data[quantity].fillna(0).plot(
            ax=ax, kind="area", color=meta["color"], alpha=meta["alpha"], zorder=5
        )

        if artists is not None:
            proxy = pyplot.Rectangle(
                (0, 0), 1, 1, facecolor=meta["color"], linewidth=0, alpha=meta["alpha"]
            )
            artists.append(proxy)
        if otherlabels is not None:
            otherlabels.append(label)

        return fig, otherlabels, artists

    def summaryPlot(
        self,
        axratio=2,
        filename=None,
        showLegend=True,
        precip=True,
        inflow=True,
        outflow=True,
        figopts={},
        serieslabels={},
    ):
        """
        Creates a figure showing the hydrlogic record (flow and
            precipitation) of the storm

        Input:
            axratio : optional float or int (default = 2)
                Relative height of the flow axis compared to the
                precipiation axis.

            filename : optional string (default = None)
                Filename to which the figure will be saved.

            **figwargs will be passed on to `pyplot.Figure`

        Writes:
            Figure of flow and precipitation for a storm

        Returns:
            None
        """
        fig = pyplot.figure(**figopts)
        gs = gridspec.GridSpec(
            nrows=2, ncols=1, height_ratios=[1, axratio], hspace=0.12
        )
        rainax = fig.add_subplot(gs[0])
        rainax.yaxis.set_major_locator(pyplot.MaxNLocator(5))
        flowax = fig.add_subplot(gs[1], sharex=rainax)

        # create the legend proxy artists
        artists = []
        labels = []

        # in the label assignment: `serieslabels.pop(item, item)` might
        # seem odd. What it does is looks for a label (value) in the
        # dictionary with the key equal to `item`. If there is no valur
        # for that key in the dictionary the `item` itself is returned.
        # so if there's nothing called "test" in mydict,
        # `mydict.pop("test", "test")` returns `"test"`.
        if inflow:
            fig, labels, artists = self.plot_hydroquantity(
                self.inflowcol,
                ax=flowax,
                label=serieslabels.pop(self.inflowcol, self.inflowcol),
                otherlabels=labels,
                artists=artists,
            )

        if outflow:
            fig, labels, arti = self.plot_hydroquantity(
                self.outflowcol,
                ax=flowax,
                label=serieslabels.pop(self.outflowcol, self.outflowcol),
                otherlabels=labels,
                artists=artists,
            )

        if precip:
            fig, labels, arti = self.plot_hydroquantity(
                self.precipcol,
                ax=rainax,
                label=serieslabels.pop(self.precipcol, self.precipcol),
                otherlabels=labels,
                artists=artists,
            )
            rainax.invert_yaxis()

        if showLegend:
            leg = rainax.legend(
                artists,
                labels,
                fontsize=7,
                ncol=1,
                markerscale=0.75,
                frameon=False,
                loc="lower right",
            )
            leg.get_frame().set_zorder(25)
            _leg = [leg]
        else:
            _leg = None

        seaborn.despine(ax=rainax, bottom=True, top=False)
        seaborn.despine(ax=flowax)
        flowax.set_xlabel("")
        rainax.set_xlabel("")

        if filename is not None:
            fig.savefig(
                filename,
                dpi=300,
                transparent=True,
                bbox_inches="tight",
                bbox_extra_artists=_leg,
            )

        return fig, artists, labels

class HydroRecord(object):
    """ Class representing an entire hydrologic record.

    Parameters
    ----------
    hydrodata : pandas.DataFrame
        DataFrame of hydrologic data of the storm. Should contain
        a unique index of type pandas.DatetimeIndex.
    precipcol : string, optional (default = None)
        Name of column in `hydrodata` containing precipiation data.
    inflowcol : string, optional (default = None)
        Name of column in `hydrodata` containing influent flow data.
    outflowcol : string, optional (default = None)
        Name of column in `hydrodata` containing effluent flow data.
    baseflowcol : string, optional (default = None)
        Name of column in `hydrodata` containing boolean indicating
        which records are considered baseflow.
    stormcol : string (default = 'storm')
        Name of column in `hydrodata` indentifying distinct storms.
    minprecip, mininflow, minoutflow : float, optional (default = 0)
        The minimum amount of each hydrologic quantity below which a
        storm can be considered "small".
    outputfreqMinutes : int, optional (default = 10)
        The default frequency (minutes) to which all data will be
        resampled. Precipitation data will be summed up across '
        multiple timesteps during resampling, while flow will be
        averaged.
    intereventHours : int, optional (default = 6)
        The dry duration (no flow or rain) required to signal the end of
        a storm.
    volume_conversion : float, optional (default = 1)
        Conversion factor to go from flow to volume for a single
        observation.
    stormclass : object, optional
        Defaults to wqio.hydro.Storm. Can be a subclass of that in cases
        where custom functionality is needed.
    lowmem : bool (default = False)
        If True, all dry observations are removed from the dataframe.

    """

    # TODO: rename `outputfreqMinutes` to `outputPeriodMinutes`
    def __init__(
        self,
        hydrodata,
        precipcol=None,
        inflowcol=None,
        outflowcol=None,
        baseflowcol=None,
        tempcol=None,
        stormcol="storm",
        minprecip=0.0,
        mininflow=0.0,
        minoutflow=0.0,
        outputfreqMinutes=10,
        intereventHours=6,
        volume_conversion=1,
        stormclass=None,
        lowmem=False,
    ):

        # validate input
        if precipcol is None and inflowcol is None and outflowcol is None:
            msg = "`hydrodata` must have at least a precip or in/outflow column"
            raise ValueError(msg)

        self.stormclass = stormclass or Storm

        # static input
        self._raw_data = hydrodata
        self.precipcol = precipcol
        self.inflowcol = inflowcol
        self.outflowcol = outflowcol
        self.baseflowcol = baseflowcol
        self.stormcol = stormcol
        self.tempcol = tempcol
        self.outputfreq = pandas.offsets.Minute(outputfreqMinutes)
        self.intereventHours = intereventHours
        self.intereventPeriods = MIN_PER_HOUR / self.outputfreq.n * self.intereventHours
        self.minprecip = minprecip
        self.mininflow = mininflow
        self.minoutflow = minoutflow
        self.volume_conversion = volume_conversion
        self.lowmem = lowmem

        # properties
        self._data = None
        self._all_storms = None
        self._storms = None
        self._storm_stats = None

    @property
    def data(self):
        if self._data is None:
            self._data = self._define_storms()
            if self.lowmem:
                self._data = self._data[self._data[self.stormcol] != 0]

        return self._data

    @property
    def all_storms(self):
        if self._all_storms is None:
            self._all_storms = {}
            for storm_number in self.data[self.stormcol].unique():
                if storm_number > 0:
                    this_storm = self.stormclass(
                        self.data,
                        storm_number,
                        precipcol=self.precipcol,
                        inflowcol=self.inflowcol,
                        outflowcol=self.outflowcol,
                        tempcol=self.tempcol,
                        stormcol=self.stormcol,
                        volume_conversion=self.volume_conversion,
                        freqMinutes=self.outputfreq.n,
                    )
                    self._all_storms[storm_number] = this_storm

        return self._all_storms

    @property
    def storms(self):
        if self._storms is None:
            self._storms = {}
            for snum, storm in self.all_storms.items():
                is_small = storm.is_small(
                    minprecip=self.minprecip,
                    mininflow=self.mininflow,
                    minoutflow=self.minoutflow,
                )

                if not is_small:
                    self._storms[snum] = storm

        return self._storms

    @property
    def storm_stats(self):
        col_order = [
            "Storm Number",
            "Antecedent Days",
            "Season",
            "Start Date",
            "End Date",
            "Duration Hours",
            "Peak Precip Intensity",
            "Total Precip Depth",
            "Total Inflow Volume",
            "Peak Inflow",
            "Total Outflow Volume",
            "Peak Outflow",
            "Peak Lag Hours",
            "Centroid Lag Hours",
        ]
        if self._storm_stats is None:
            storm_stats = pandas.DataFrame(
                [self.storms[sn].summary_dict for sn in self.storms]
            )

            self._storm_stats = storm_stats[col_order]

        return self._storm_stats.sort_values(by=["Storm Number"]).reset_index(drop=True)

    def _define_storms(self, debug=False):

        parsed = parse_storm_events(
            self._raw_data,
            self.intereventHours,
            self.outputfreq.n,
            precipcol=self.precipcol,
            inflowcol=self.inflowcol,
            outflowcol=self.outflowcol,
            baseflowcol=self.baseflowcol,
            stormcol="storm",
            debug=debug,
        )
        return parsed

    def getStormFromTimestamp(self, timestamp, lookback_hours=0, smallstorms=False):
        """ Get the storm associdated with a give (sample) date

        Parameters
        ----------
        timestamp : pandas.Timestamp
            The date/time for which to search within the hydrologic
            record.
        lookback_hours : positive int or float, optional (default = 0)
            If no storm is actively occuring at the provided timestamp,
            we can optionally look backwards in the hydrologic record a
            fixed amount of time (specified in hours). Negative values
            are ignored.
        smallstorms : bool, optional (default = False)
            If True, small storms will be included in the search.

        Returns
        -------
        storm_number : int
        storm : wqio.Storm

        """

        # santize date input
        timestamp = validate.timestamp(timestamp)

        # check lookback hours
        if lookback_hours < 0:
            raise ValueError("`lookback_hours` must be greater than 0")

        # initial search for the storm
        storm_number = int(self.data.loc[:timestamp, self.stormcol].iloc[-1])

        # look backwards if we have too
        if (storm_number == 0 or pandas.isnull(storm_number)) and lookback_hours != 0:
            lookback_time = timestamp - pandas.offsets.Hour(lookback_hours)
            storms = self.data.loc[lookback_time:timestamp, [self.stormcol]]
            storms = storms[storms > 0].dropna()

            if storms.shape[0] == 0:
                # no storm
                storm_number = None
            else:
                # storm w/i the lookback period
                storm_number = int(storms.iloc[-1])

        # return storm_number and storms
        if smallstorms:
            return storm_number, self.all_storms.get(storm_number, None)
        else:
            return storm_number, self.storms.get(storm_number, None)

    def histogram(self, valuecol, bins, **factoropts):
        """ Plot a faceted, categorical histogram of storms.

        Parameters
        ----------
        valuecol : str, optional
            The name of the column that should be categorized and plotted.
        bins : array-like, optional
            The right-edges of the histogram bins.
        factoropts : keyword arguments, optional
            Options passed directly to seaborn.factorplot

        Returns
        -------
        fig : seaborn.FacetGrid

        See also
        --------
        viz.categorical_histogram
        seaborn.factorplot

        """

        fg = viz.categorical_histogram(self.storm_stats, valuecol, bins, **factoropts)
        fg.fig.tight_layout()
        return fg

class DrainageArea(object):
    def __init__(self, total_area=1.0, imp_area=1.0, bmp_area=0.0):
        """ A simple object representing the drainage area of a BMP.

        Units are not enforced, so keep them consistent yourself. The
        calculations available assume that the area of the BMP and the
        "total" area are mutually exclusive. In other words,
        the watershed outlet is at the BMP inlet.

        Parameters
        ----------
        total_area : float, optional (default = 1.0)
            The total geometric area of the BMP's catchment
        imp_area : float, optional (default = 1.0)
            The impervious area of the BMP's catchment
        bmp_area : float, optional (default = 0.0)
            The geometric area of the BMP itself.

        """

        self.total_area = float(total_area)
        self.imp_area = float(imp_area)
        self.bmp_area = float(bmp_area)

    def simple_method(self, storm_depth, volume_conversion=1.0, annual_factor=1.0):
        """
        Estimate runoff volume via Bob Pitt's Simple Method.

        Parameters
        ----------
        storm_depth : float
            Depth of the storm.
        volume_conversion : float, optional (default = 1.0)
            Conversion factor to go from [area units] * [depth units] to
            the desired [volume units]. If [area] = m^2, [depth] = mm,
            and [volume] = L, then `volume_conversion` = 1.
        annual_factor : float, optional (default = 1.0)
            The Simple Method's annual correction factor to account for
            small storms that do not produce runoff.

        Returns
        -------
        runoff_volume : float
            The volume of water entering the BMP immediately downstream
            of the drainage area.

        """

        # volumetric run off coneffiecient
        Rv = 0.05 + (0.9 * (self.imp_area / self.total_area))

        # run per unit storm depth
        drainage_conversion = Rv * self.total_area * volume_conversion
        bmp_conversion = self.bmp_area * volume_conversion

        # total runoff based on actual storm depth
        runoff_volume = (
            drainage_conversion * annual_factor + bmp_conversion
        ) * storm_depth
        return runoff_volume

max_stars_repo_name: Geosyntec/wqio

max_issues_repo_name: Geosyntec/wqio

max_forks_repo_name: Geosyntec/wqio


text:
from input.enums import UsefullButton

import itertools
import numpy as np

class ControllerState(dict):
    def __init__(self, button=None, stick=(0.5, 0.5), c_stick=(0.5, 0.5), duration=3, no_op=False):
        dict.__init__(self)
        for item in UsefullButton:
            if button and item.name == button:
                self[button] = True
            else:
                self[item.name] = False

        self.to_press = np.array([button]) if button is not None else np.array([], np.int32)

        self['stick'] = stick
        self['c_stick'] = c_stick
        self['duration'] = duration
        self['id'] = None
        self.no_op = no_op

    def __str__(self):

        string = ""
        for item in UsefullButton:
            string += "%s(%d) " % (item.name, self[item.name])
        string += "S(%s,%s) " % self['stick']
        string += "C(%s,%s)" % self['c_stick']

        return string

    def send_controller(self, pad):
        if not self.no_op:
            for button in np.setdiff1d(pad.previous_state.to_press, self.to_press):
                pad.release_button(button, buffering=True)
            for button in self.to_press:
                pad.press_button(button, buffering=True)
            pad.tilt_stick('MAIN', *self['stick'], buffering=True)
            pad.tilt_stick('C', *self['c_stick'], buffering=True)
            pad.flush()
            pad.previous_state = self

class ActionSpace:
    neutral = [
        (0.5, 0.5)
    ]

    back_front = [
        (0., 0.5),
        (1., 0.5)
    ]

    cardinal = [
        (0.5 * (1 + np.cos(x)), 0.5 * (1 + np.sin(x))) for x in np.arange(0, 4) * np.pi * 0.5
    ]

    diagonal = [
        (0.5 * (1 + np.cos(x)), 0.5 * (1 + np.sin(x))) for x in np.arange(1, 8, 2) * np.pi * 0.25
    ]

    other_angles = [
        (0.5 * (1 + np.cos(x)), 0.5 * (1 + np.sin(x))) for x in np.arange(1, 17, 2) * np.pi * 0.125
    ]

    all_stick_states = neutral + cardinal + diagonal + other_angles

    stick_states_upB = [
        (0.853, 0.853),
        (0.146, 0.853),
    ]

    smash_states = cardinal

    tilt_stick_states = [
        (0.3, 0.5),
        (0.7, 0.5),
        (0.5, 0.3),
        (0.5, 0.7)
    ]

    wave_dash_sticks = [
        (0.023, 0.353),
        (0.977, 0.353),

    ]

    def __getitem__(self, item):
        return self.controller_states[item]

    def __init__(self, short_hop=3):
        self.len = 0

        """
        1 directions with no button
        2 tap b, lazer, plumbers downb
        3 tilts
        4 smashes with c stick and cardinals
        5 jumps, normal, front, back, full hop, short hop
        6 upbs
        7 L with all directions (rolls, spot dodge, shield, air dodges...)
        8 wavedashes
        9 shield grab, jump grab
        10 tap down b, tap side b
        11 noop

        """

        all_states = [ControllerState(stick=stick) for stick in self.all_stick_states] + \
                     [[ControllerState(stick=self.cardinal[-1], button='B', duration=1), ControllerState(duration=1)]] + \
                     [ControllerState(stick=stick, button='A') for stick in (self.tilt_stick_states + self.neutral)] + \
                     [ControllerState(stick=stick, c_stick=c_stick) for stick, c_stick in
                      itertools.product(self.cardinal + self.neutral, self.cardinal)] + \
                     [ControllerState(stick=stick, button='X', duration=short_hop + 1) for stick in
                      (self.back_front + self.neutral)] + \
                     [[ControllerState(stick=stick, button='X', duration=short_hop),
                       ControllerState(stick=stick, duration=1)] for stick in (self.back_front + self.neutral)] + \
                     [ControllerState(stick=upB_stick, button='B') for upB_stick in self.stick_states_upB] + \
                     [ControllerState(stick=stick, button='L') for stick in
                      (self.neutral + self.cardinal + self.diagonal)] + \
                     [[ControllerState(stick=stick, button='X', duration=short_hop),
                       ControllerState(stick=stick, duration=1), ControllerState(stick=stick, button='L', duration=1)] for stick in
                      (self.wave_dash_sticks)] + \
                     [[ControllerState(stick=stick, button='B', duration=1), ControllerState(stick=stick, duration=1)]
                      for stick in (self.back_front)] + \
                     [ControllerState(button='B')] + \
                     [[ControllerState(button='X', duration=1), ControllerState(button='Z', duration=2)]] + \
                     [ControllerState(no_op=True)]

        shield_grab = ControllerState()
        shield_grab.to_press = np.array(['A', 'L'])
        all_states += [shield_grab]

        self.controller_states = np.array(
            all_states
        , dtype=ControllerState)

        self.dim = len(self.controller_states)

ext: py

max_stars_repo_name: villinvic/Georges

max_issues_repo_name: villinvic/Georges

max_forks_repo_name: villinvic/Georges

reason: import numpy


text:
# import bp_pas.params as params
# import bp_pas.corpus as corpus
# import bp_pas.stats as stats
# import sys
# import gflags
# from bp_pas.eval import evaluate

import bp_pas.params as params
import bp_pas.corpus as corpus
import bp_pas.stats as stats
from bp_pas.eval import evaluate
from bp_pas.util.vocab import Vocab

import sys
import gflags
import tensorflow as tf
import numpy as np

from bp_pas.eval import evaluate
from bp_pas.ling.pas import PAS, Predicate, Argument
from bp_pas.ling.sent import Sentence

from bp_pas.model.baseline import Model
from bp_pas.util.feed_dict_helper import separate_pred_dicts

def main(flags):
    # tmp params to add to flags
    unk_threshold = 1
    unk_token = '<UNK>'

    if not flags.use_context:
        flags.context_dims = 0
    if not flags.use_sp:
        flags.sp_dims = 0

    synt_arg_embedding_size = flags.context_dims #+ 1
    sem_arg_embedding_size = flags.sp_dims

    complete_embedding_size = synt_arg_embedding_size + sem_arg_embedding_size # + 1 # for the pred id concat
    num_training_instances = 1

    assert complete_embedding_size % 2 == 0
    print(complete_embedding_size)

    # Load data
    print(flags.train_data)
    ntc = corpus.NTCLoader()
    train_data = ntc.load_corpus(flags.train_data, flags.max_train_instances)
    test_data = ntc.load_corpus(flags.test_data, flags.max_train_instances)
    dev_data = ntc.load_corpus(flags.dev_data, flags.max_train_instances)
    print('{} train sentences.'.format(len(train_data)))
    print('{} test sentences.'.format(len(test_data)))
    print('{} dev sentences.'.format(len(dev_data)))
    stats.corpus_statistics(train_data)
    stats.show_case_dist(train_data)

    # Unpack data from "docs"->sents
    train_data = train_data[0]
    test_data = test_data[0]
    dev_data = dev_data[0]
    print('{} train sentences.'.format(len(train_data)))
    print('{} test sentences.'.format(len(test_data)))
    print('{} dev sentences.'.format(len(dev_data)))

    # Setup vocabulary
    all_sents = train_data + test_data + dev_data
    print('Number of total sentences: {}'.format(len(all_sents)))

    print('Collecting all tokens...')
    word_tokens = list(word.form
                       for sent in all_sents
                       for word in sent)
    print('Total tokens: {}\n'.format(len(word_tokens)))

    print('Building (argument) vocabulary...')
    arg_vocab = Vocab(init_words=word_tokens,
                      unk_token=unk_token,
                      unk_threshold=unk_threshold)
    arg_vocab.freeze()
    arg_vocabulary_size = len(arg_vocab)
    print('Vocabulary size: {}\n'.format(arg_vocabulary_size))

    print('Collecting predicates...')
    pred_vocab = Vocab(init_words=[word.form
                                   for sent in all_sents
                                   for word in sent if word.is_prd])
    pred_vocab.freeze()
    pred_vocabulary_size = len(pred_vocab)
    print('Number of predicates: {}\n'.format(pred_vocabulary_size))

    print('Collecting argument types...')
    arg_types = ['NIL'] + list(set([arg.arg_type
                                    for sent in all_sents
                                    for pas in sent.pas
                                    for arg in pas.args]))
    print('Arg types: ', arg_types)
    num_types = len(arg_types)

    # print(pred_vocab.elems())

    # Convert data structures into numpy ndarrays for placeholders
    train_dicts = [ex
                   for td in train_data
                   for ex in separate_pred_dicts(td, pred_vocab, arg_vocab, arg_types)]

    test_dicts = [separate_pred_dicts(td, pred_vocab, arg_vocab, arg_types)
                  for td in train_data]

    # Just for now
    test_data = train_data
    print(len(train_dicts))
    print(len(test_dicts))
    print('pred ids: ', train_dicts[0][2])
    print('pred ids: ', len(train_dicts[0][2]))
    print('sent ids: ', len(train_dicts[0][0]))
#    assert (len(sent_ids) == len(pred_id_vector))

    # Setup the context embedding
    tf.reset_default_graph()

    # Construct syntactic embedding matrix for contex/arguments
    synt_arg_embeddings = tf.Variable(
        tf.random_uniform([arg_vocabulary_size, synt_arg_embedding_size], -1.0, 1.0, dtype=tf.float64), trainable=False)

    # Construct semantic embedding matrix for selectional preference
    sem_arg_embeddings = tf.Variable(
        tf.random_uniform([arg_vocabulary_size, sem_arg_embedding_size], -1.0, 1.0, dtype=tf.float64), trainable=False)

    # Construct embedding matrix for predicates
    pred_embeddings = tf.Variable(
        tf.random_uniform([pred_vocabulary_size, num_types * complete_embedding_size], -1.0, 1.0, dtype=tf.float64),
        trainable=True)

    # Setup placeholders
    batch_size = 1
    sent_placeholder = tf.placeholder(tf.int32, shape=[None])
    pred_placeholder = tf.placeholder(tf.int32, shape=[None])
    gold_placeholder = tf.placeholder(tf.float64, shape=[num_types, None])
    pred_id_placeholder = tf.placeholder(tf.float64, shape=[None])

    model = Model(pred_embeddings, synt_arg_embeddings, sem_arg_embeddings,
                  sent_placeholder, pred_placeholder, gold_placeholder, flags)

    rep = model.complete_embeddings(sent_placeholder, pred_id_placeholder, sem_arg_embeddings, synt_arg_embeddings, synt_arg_embedding_size)
    scored_mat_opt = model.pred_full_scoring(pred_placeholder, pred_embeddings, rep)

    sess = tf.Session()
    sess.run(tf.global_variables_initializer())

    opt = tf.train.GradientDescentOptimizer(learning_rate=0.001)
    loss = model.loss_old(gold_placeholder, pred_placeholder, pred_embeddings, rep)
    opt_op = opt.minimize(loss)

    def meta_decode(sess, instance):
        #    print('In decode')
        # Collect model-scores labels-x-tokens mats
        int_inputs = separate_pred_dicts(instance, pred_vocab, arg_vocab, arg_types)
        #    print('plen: ', len(int_inputs))
        scored_mats = []
        for sent_ids, pred_ids, pred_id_vector, gold_labels in int_inputs:
            print('pred ids: ', pred_id_vector)
            print('sent ids: ', sent_ids)
            assert(len(sent_ids) == len(pred_id_vector))
            feed_dict = {
                sent_placeholder: sent_ids,
                pred_placeholder: pred_ids,
                gold_placeholder: gold_labels,
                pred_id_placeholder: pred_id_vector
            }
            scored_mats.append(sess.run(scored_mat_opt, feed_dict=feed_dict))
            #        print(gold_labels)
            #        print('\n')
            #        print(scored_mats[-1])
            #        print('------------------------')
            #
            #    print('scored ', len(scored_mats))
        pred_indices = [i for i, w in enumerate(instance.words) if w.is_prd]
        num_preds = len(pred_indices)
        pases = []
        for pred_id, scored_mat in zip(pred_indices, scored_mats):
            #        print(pred_id)
            pred = Predicate(word_index=pred_id,
                             word_form=instance.words[pred_id].form)
            args = []
            for slot, slot_id in zip(arg_types, range(num_types)):
                if slot != 'NIL':
                    max_index = np.argmax(scored_mat[slot_id])
                    #                print(slot, ': ', max_index)
                    args.append(Argument(word_index=max_index,
                                         word_form=instance.words[max_index].form,
                                         arg_type=slot))
                    #                print('\t', max_index)
            pases.append(PAS(pred, args))
            #        print('psize: ', len(pases))
        return Sentence(instance.words, pases)

    num_epochs = 10000
    for epoch in range(num_epochs):
        epoch_loss = 0
        for train_dict in train_dicts:
            sent_ids, pred_ids, pred_id_vector, gold_labels = train_dict
            #        print(gold_labels.shape)
            _, closs = sess.run([opt_op, loss], feed_dict={
                sent_placeholder: sent_ids,
                pred_placeholder: pred_ids,
                pred_id_placeholder: pred_id_vector,
                gold_placeholder: gold_labels
            })
            epoch_loss += closs.mean()
        if epoch > 0 and epoch % 1000 == 0:
            decoded = [pas
                       for tsent in train_data
                       for pas in meta_decode(sess, tsent).pas]  # [0:1]
            #        decoded = [pas
            #                   for tsent, tdict in zip(train_data, train_dicts)
            #                   for pas in meta_decode(sess, tsent)] # #model.decode(sess, tsent, tdict).pas]
            gold = [pas
                    for sent in train_data
                    for pas in sent.pas]  # [0:1]
            evaluate(decoded, gold, verbose=False)
            print('{0:>4}: {1:.3f}'.format(epoch, epoch_loss))

                #
    # print(flags.train_data)
    # ntc = corpus.NTCLoader()
    # train_data = ntc.load_corpus(flags.train_data, flags.max_train_instances)
    # test_data  = ntc.load_corpus(flags.test_data, flags.max_train_instances)
    # dev_data   = ntc.load_corpus(flags.dev_data, flags.max_train_instances)
    # print('{} train docs.'.format(len(train_data)))
    # print('{} test docs.'.format(len(test_data)))
    # print('{} dev docs.'.format(len(dev_data)))
    # stats.corpus_statistics(train_data)
    # stats.show_case_dist(train_data)

    # for w in train_data[0][0]:
    #     print(w)
    #     print(w.chunk_index)
    #     print(w.chunk_head)
    #     print(w.is_prd)
    #     print(w.arg_indices)
    #     print(w.arg_types)
    #     print()

#    print(test_data[0])
#    print(test_data[0][0])

#    re = SampleEval()
#    for d1, d2 in zip(test_data, test_data):
#        result = d1
#        sample = d2
#        re.update_results(y_sys_batch=result, sample=sample)

#    test = [pas for doc in test_data for sent in doc for pas in sent.pas]
#    gold = [pas for doc in test_data for sent in doc for pas in sent.pas]
#    evaluate(test, gold)

#    for d1, d2 in zip(test_data, test_data):
#        print(d1)
#        print(len(d1))
#        print(d1[0])
#        s1 = ntc_sent2conll(d1)
#        print(s1)
#        s2 = ntc_sent2conll(d2)
#        re._add_results_gold(s1)
#        re._add_results_sys(s2)
#        break

if __name__ == "__main__":
    params.load_defaults()
    flags = gflags.FLAGS
    argv = flags(sys.argv)
    main(flags)

max_stars_repo_name: cambridgeltl/jp-pred-arg

max_issues_repo_name: cambridgeltl/jp-pred-arg

max_forks_repo_name: cambridgeltl/jp-pred-arg


text:
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Tue Apr  2 14:16:39 2019

@author: q
"""
import os
import random
from os.path import join
import cv2
import numpy as np
import numpy.random as npr

from path_and_config import path
from assist.iou import IoU
from assist.mark_infoget_imgchange import  flip,getDataFromTxt, BBox
def do(a_img_a_box_pic_resized, a_img_marks_rate,name):
    a,b=a_img_a_box_pic_resized, a_img_marks_rate
    a_img_a_box_pic_resized, a_img_marks_rate = np.asarray(a_img_a_box_pic_resized), np.asarray(a_img_marks_rate)

    if np.sum(np.where(a_img_marks_rate[-1] <= 0, 1, 0)) > 0:
        return a,b
    if np.sum(np.where(a_img_marks_rate[-1] >= 1, 1, 0)) > 0:
        return a,b
    cv2.imwrite(join(path.root,"mark_test_%d/%s.jpg" % (pic_size,name)), a_img_a_box_pic_resized[-1])
    landmarks = map(str, list(a_img_marks_rate[-1]))
    f.write(join(path.root, "mark_test_%d/%s.jpg" % (pic_size,name)) + " -2 "+" ".join(landmarks)+"\n")
    return a,b
def rate_to_value(box, point):
    x = box[0] + (box[2]-box[0])*point[0]
    y = box[1] + (box[3]-box[1])*point[1]
    return np.asarray([x, y])
def rate_to_value_marks(box, landmark):
    p = np.zeros((len(landmark), 2))
    for i in range(len(landmark)):
        p[i] = rate_to_value(box,landmark[i])
    return p
def value_to_rate(box, point):
    """相对于label框的偏置的比例"""
    x = (point[0] - box[0]) / (box[2]-box[0])
    y = (point[1] - box[1]) / (box[3]-box[1])
    return np.asarray([x, y])
def value_to_rate_marks(box, landmark):
    """给出标记位置的偏移量，获得偏移率"""
    p = np.zeros((len(landmark), 2))
    for i in range(len(landmark)):
        p[i] = value_to_rate(box,landmark[i])
    return p
def rotate(img, bbox, box,landmark, alpha):
    """
        given a face with bbox and landmark, rotate with alpha
        and return rotated face with bbox, landmark (absolute position)
    """
    center = ((bbox.left+bbox.right+12)/2, (bbox.top+bbox.bottom+12)/2)
    rot_mat = cv2.getRotationMatrix2D(center, alpha, 1)#旋转的矩阵
    print('rot mat',rot_mat)
    #whole image rotate
    #pay attention: 3rd param(col*row)
    img_rotated_by_alpha = cv2.warpAffine(img, rot_mat,(img.shape[1],img.shape[0]))
    print('img_rotated_by_alpha ', img_rotated_by_alpha.shape)
    landmark_ = np.asarray([(rot_mat[0][0]*x+rot_mat[0][1]*y+rot_mat[0][2],
                 rot_mat[1][0]*x+rot_mat[1][1]*y+rot_mat[1][2]) for (x, y) in landmark])
    #crop face
    face = img_rotated_by_alpha[box[1]:box[3]+1,box[0]:box[2]+1]
    return (face, landmark_)
"""准备工作"""
#pic_size = 12#第一次PNet
#pic_size = 24#第二次RNet
pic_size = 48#第三次ONet

expand = True#如果标记的图片数量少，可以通过旋转镜像等来增加mark图片数量
mark_txt = 'txt/trainImageList.txt'
txt_save_dir = os.path.join(path.root,'pic_txt')
mark_save_dir = os.path.join(path.root,'mark_test_%d'%pic_size)
if not os.path.exists(mark_save_dir):
    os.mkdir(mark_save_dir)
f = open(os.path.join(txt_save_dir, 'mark_test_%s.txt'%(pic_size)), 'w')
mark_data = getDataFromTxt(mark_txt, data_path = path.mark)
print('total data is ',len(mark_data))

"""txt中box区域图片生成，txt的生成，"""
img_done_count = 0#计数图片标记点，所有图片的所有label狂的信息提取出了多少
mark_pic_count = 0#最终符合条件的个数mark图像
for (imgPath, box, marks) in mark_data[1:2]:
    a_img_a_box_pic_resized = []#一个照片输出的所有的mark_pic
    a_img_marks_rate = []#所有照片的所有标记点
    img = cv2.imread(imgPath)
    if img is None:
        print(imgPath,box,marks,img_done_count)

    #txt中的数据，这个是第一i部分
    img_h, img_w, img_c = img.shape
    print('img shape',img.shape)
    gt_box = np.array([box.left, box.top, box.right, box.bottom])

    box_pic = img[box.top : box.bottom+1, box.left : box.right+1]
    box_pic = cv2.resize(box_pic, (pic_size, pic_size))
    marks_rate = np.zeros((5, 2))
    for index, a_mark in enumerate(marks):
        rv = ((a_mark[0] - gt_box[0]) / (gt_box[2] - gt_box[0]),
              (a_mark[1] - gt_box[1]) / (gt_box[3] - gt_box[1]))
        marks_rate[index] = rv
    a_img_a_box_pic_resized.append(box_pic)
    a_img_marks_rate.append(marks_rate.reshape(10))
    a_img_a_box_pic_resized, a_img_marks_rate= do(a_img_a_box_pic_resized, a_img_marks_rate,'ori-box')
   #0
    #第二层，每个box随机构造一个box
    if expand:
        img_done_count = img_done_count + 1
        if img_done_count % 100 == 0:
            print(img_done_count,'/%s'%(len(mark_data)), 'images done')
        x1, y1, x2, y2 = gt_box
        gt_box_w = x2 - x1 + 1
        gt_box_h = y2 - y1 + 1
        print('gt',gt_box_w,gt_box_h)
        if max(gt_box_w, gt_box_h) < 40 or x1< 0 or y1 < 0:
            #box size比40还小，或者框小于0了，及不必扩展增加数量了
            continue
        for i in range(1):
            #中心移动
            new_box_size = npr.randint(int(min(gt_box_w, gt_box_h) * 0.8),np.ceil(1.25 * max(gt_box_w, gt_box_h)))
            print('new box size' ,new_box_size)
            delta_x = npr.randint(-gt_box_w * 0.2, gt_box_w * 0.2)
            delta_y = npr.randint(-gt_box_h * 0.2, gt_box_h * 0.2)
            gt_box_center_x = x1 + gt_box_w/2
            gt_box_center_y = y1 + gt_box_h/2
            new_x1 = int(max(gt_box_center_x + delta_x - new_box_size/2, 0))
            new_y1 = int(max(gt_box_center_y + delta_y - new_box_size/2, 0))
            new_x2 = new_x1 + new_box_size
            new_y2 = new_y1 + new_box_size
            if new_x2 > img_w or new_y2 > img_h:
                continue

            new_box = np.array([new_x1, new_y1, new_x2, new_y2])
            new_img = img[new_y1 : new_y2+1, new_x1 : new_x2+1,:]
            new_img_resized = cv2.resize(new_img, (pic_size, pic_size))
            iou = IoU(new_box, np.expand_dims(gt_box, 0))
            if iou > 0.65:
                print(iou)
                a_img_a_box_pic_resized.append(new_img_resized)
                marks_rate = np.zeros((5, 2))
                mmk = np.zeros((5, 2))
                for index, one in enumerate(marks):
                    rv = ((one[0] - new_x1) / new_box_size,(one[1] - new_y1) / new_box_size)
                    mmk[index] = (one[0],one[1])
                    marks_rate[index] = rv
                a_img_marks_rate.append(marks_rate.reshape(10))#1
                a_img_a_box_pic_resized, a_img_marks_rate= do(a_img_a_box_pic_resized, a_img_marks_rate,'rang')

                marks_rate = np.zeros((5, 2))
                landmark_ = a_img_marks_rate[-1].reshape(-1, 2)
                new_rand_box = BBox([new_x1, new_y1, new_x2, new_y2])
                #随机指定下面的进行
                """是否增加一组反转的"""
                if random.choice([0, 1]) >1:
                    face_flipped, landmark_flipped = flip(new_img_resized, landmark_)
                    face_flipped = cv2.resize(face_flipped,(pic_size, pic_size))
                    a_img_a_box_pic_resized.append(face_flipped)
                    a_img_marks_rate.append(landmark_flipped.reshape(10))#3
                    a_img_a_box_pic_resized, a_img_marks_rate= do(a_img_a_box_pic_resized, a_img_marks_rate,'flip')
                """是否增加一组镜逆时针旋转5度的，和一组旋转后镜像的"""
                if random.choice([0, 1]) >= 0:
                    print(gt_box)
                    print(new_box)
                    print('before rote vale\n  ',box.rate_to_value_marks(landmark_))
                    print('before rote rate\n  ',landmark_)
                    #face_rotated_by_alpha, landmark_rotated = rotate(img, box,new_box,rate_to_value_marks(new_box,landmark_),5)
                    face_rotated_by_alpha, landmark_rotated = rotate(img, box,new_box,mmk,5)
                    print('after rote  value',landmark_rotated)
                    #逆时针旋转5度
                    landmark_rotated = value_to_rate_marks(new_box,landmark_rotated)
                    print('after rote  rate\n',landmark_rotated)
                    print(face_rotated_by_alpha.shape)
                    #face_rotated_by_alpha = cv2.resize(face_rotated_by_alpha,(pic_size, pic_size))
                    a_img_a_box_pic_resized.append(face_rotated_by_alpha)#4
                    #a_img_a_box_pic_resized.append(img)#4
                    a_img_marks_rate.append(landmark_rotated.reshape(10))
                    a_img_a_box_pic_resized, a_img_marks_rate= do(a_img_a_box_pic_resized, a_img_marks_rate,'pos5newbox')
                    print('pos5')
                    face_rotated_by_alpha = cv2.resize(face_rotated_by_alpha,(pic_size, pic_size))
                    a_img_a_box_pic_resized.append(face_rotated_by_alpha)#4
                    #a_img_a_box_pic_resized.append(img)#4
                    a_img_marks_rate.append(landmark_rotated.reshape(10))
                    a_img_a_box_pic_resized, a_img_marks_rate= do(a_img_a_box_pic_resized, a_img_marks_rate,'pos5newbox_48')
                    print('pos5-48')

                """是否增加一组镜顺时针旋转5度的，和一组旋转后镜像的"""
                """
                if random.choice([0, 1]) >= 0:
                    #face_rotated_by_alpha, landmark_rotated = rotate(img, box,new_box, rate_to_value_marks(new_box,landmark_),-5)
                    face_rotated_by_alpha, landmark_rotated = rotate(img, box,new_box, mmk,-5)
                    #逆时针旋转5度
                    landmark_rotated = value_to_rate_marks(new_box,landmark_rotated)
                    #face_rotated_by_alpha = cv2.resize(face_rotated_by_alpha,(pic_size, pic_size))
                    a_img_a_box_pic_resized.append(face_rotated_by_alpha)#5
                    a_img_marks_rate.append(landmark_rotated.reshape(10))
                    a_img_a_box_pic_resized, a_img_marks_rate= do(a_img_a_box_pic_resized, a_img_marks_rate,'neg5')
                    print('12')
                    """

f.close()


text:
import torch

from models import SpKBGATModified, SpKBGATConvOnly
from torch.autograd import Variable
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.lines import Line2D
from copy import deepcopy

from preprocess import read_entity_from_id, read_relation_from_id, init_embeddings, build_data
from create_batch import Corpus
from utils import save_model

import random
import argparse
import os
import sys
import logging
import time
import pickle

# %%
# %%from torchviz import make_dot, make_dot_from_trace

def parse_args():
	args = argparse.ArgumentParser()
	# network arguments
	args.add_argument("-data", "--data", default="../data/WN18RR/", help="data directory")
	args.add_argument("-e_g", "--epochs_gat", type=int, default=3600, help="Number of epochs")
	args.add_argument("-e_c", "--epochs_conv", type=int, default=200, help="Number of epochs")
	args.add_argument("-w_gat", "--weight_decay_gat", type=float, default=5e-6, help="L2 reglarization for gat")
	args.add_argument("-w_conv", "--weight_decay_conv", type=float, default=1e-5, help="L2 reglarization for conv")
	args.add_argument("-pre_emb", "--pretrained_emb", type=bool, default=True, help="Use pretrained embeddings")
	args.add_argument("-emb_size", "--embedding_size", type=int, default=50, help="Size of embeddings (if pretrained not used)")
	args.add_argument("-l", "--lr", type=float, default=1e-3)
	args.add_argument("-g2hop", "--get_2hop", type=bool, default=False)
	args.add_argument("-u2hop", "--use_2hop", type=bool, default=False)
	args.add_argument("-p2hop", "--partial_2hop", type=bool, default=False)
	args.add_argument("-outfolder", "--output_folder", default="./checkpoints/wn/out/", help="Folder name to save the models.")

	args.add_argument("-eval_type", "--eval_type", default='last')
	# arguments for GAT
	args.add_argument("-b_gat", "--batch_size_gat", type=int, default=86835, help="Batch size for GAT")
	args.add_argument("-neg_s_gat", "--valid_invalid_ratio_gat", type=int, default=2, help="Ratio of valid to invalid triples for GAT training")
	args.add_argument("-drop_GAT", "--drop_GAT", type=float, default=0.3, help="Dropout probability for SpGAT layer")
	args.add_argument("-alpha", "--alpha", type=float, default=0.2, help="LeakyRelu alphs for SpGAT layer")
	args.add_argument("-out_dim", "--entity_out_dim", type=int, nargs='+', default=[100, 200], help="Entity output embedding dimensions")
	args.add_argument("-h_gat", "--nheads_GAT", type=int, nargs='+', default=[2, 2], help="Multihead attention SpGAT")
	args.add_argument("-margin", "--margin", type=float, default=5, help="Margin used in hinge loss")

	# arguments for convolution network
	args.add_argument("-b_conv", "--batch_size_conv", type=int, default=128, help="Batch size for conv")
	args.add_argument("-alpha_conv", "--alpha_conv", type=float, default=0.2, help="LeakyRelu alphas for conv layer")
	args.add_argument("-neg_s_conv", "--valid_invalid_ratio_conv", type=int, default=40, help="Ratio of valid to invalid triples for convolution training")
	args.add_argument("-o", "--out_channels", type=int, default=500, help="Number of output channels in conv layer")
	args.add_argument("-drop_conv", "--drop_conv", type=float, default=0.0, help="Dropout probability for convolution layer")

	args = args.parse_args()
	return args

args = parse_args()
print(args)
# %%

def load_data(args):
	train_data, validation_data, test_data, entity2id, relation2id, headTailSelector, unique_entities_train = build_data(args.data, is_unweigted=False, directed=True)

	if args.pretrained_emb:
		entity_embeddings, relation_embeddings = init_embeddings(os.path.join(args.data, 'entity2vec.txt'), os.path.join(args.data, 'relation2vec.txt'))
		print("Initialised relations and entities from TransE")
	else:
		entity_embeddings	= np.random.randn(len(entity2id), args.embedding_size)
		relation_embeddings	= np.random.randn(len(relation2id), args.embedding_size)
		print("Initialised relations and entities randomly")

	corpus = Corpus(args, train_data, validation_data, test_data, entity2id, relation2id, headTailSelector, args.batch_size_gat, args.valid_invalid_ratio_gat, unique_entities_train, args.get_2hop)

	return corpus, torch.FloatTensor(entity_embeddings), torch.FloatTensor(relation_embeddings)

Corpus_, entity_embeddings, relation_embeddings = load_data(args)

if(args.get_2hop):
	file = args.data + "/2hop.pickle"
	with open(file, 'wb') as handle:
		pickle.dump(Corpus_.node_neighbors_2hop, handle, protocol=pickle.HIGHEST_PROTOCOL)

if(args.use_2hop):
	print("Opening node_neighbors pickle object")
	file = args.data + "/2hop.pickle"
	with open(file, 'rb') as handle:
		node_neighbors_2hop = pickle.load(handle)

entity_embeddings_copied	= deepcopy(entity_embeddings)
relation_embeddings_copied	= deepcopy(relation_embeddings)

print("Initial entity dimensions {} , relation dimensions {}".format(entity_embeddings.size(), relation_embeddings.size()))

CUDA = torch.cuda.is_available()

def batch_gat_loss(gat_loss_func, train_indices, entity_embed, relation_embed):
	len_pos_triples	= int(train_indices.shape[0] / (int(args.valid_invalid_ratio_gat) + 1))
	pos_triples	= train_indices[:len_pos_triples]
	neg_triples	= train_indices[len_pos_triples:]
	pos_triples	= pos_triples.repeat(int(args.valid_invalid_ratio_gat), 1)

	source_embeds	= entity_embed[pos_triples[:, 0]]
	relation_embeds	= relation_embed[pos_triples[:, 1]]
	tail_embeds	= entity_embed[pos_triples[:, 2]]

	x		= source_embeds + relation_embeds - tail_embeds
	pos_norm	= torch.norm(x, p=1, dim=1)

	source_embeds	= entity_embed[neg_triples[:, 0]]
	relation_embeds	= relation_embed[neg_triples[:, 1]]
	tail_embeds	= entity_embed[neg_triples[:, 2]]

	x		= source_embeds + relation_embeds - tail_embeds
	neg_norm	= torch.norm(x, p=1, dim=1)

	y = -torch.ones(int(args.valid_invalid_ratio_gat) * len_pos_triples).cuda()

	loss = gat_loss_func(pos_norm, neg_norm, y)
	return loss

def train_gat(args):

	# Creating the gat model here.
	####################################

	print("Defining model")

	print("\nModel type -> GAT layer with {} heads used , Initital Embeddings training".format(args.nheads_GAT[0]))
	model_gat = SpKBGATModified(entity_embeddings, relation_embeddings, args.entity_out_dim, args.entity_out_dim, args.drop_GAT, args.alpha, args.nheads_GAT)

	if CUDA:
		model_gat.cuda()

	optimizer	= torch.optim.Adam(model_gat.parameters(), lr=args.lr, weight_decay=args.weight_decay_gat)
	scheduler	= torch.optim.lr_scheduler.StepLR(optimizer, step_size=500, gamma=0.5, last_epoch=-1)
	gat_loss_func	= nn.MarginRankingLoss(margin=args.margin)

	current_batch_2hop_indices = torch.LongTensor([])
	if(args.use_2hop):
		current_batch_2hop_indices = Corpus_.get_batch_nhop_neighbors_all(args, Corpus_.unique_entities_train, node_neighbors_2hop)

	if CUDA: 	current_batch_2hop_indices = Variable(torch.LongTensor(current_batch_2hop_indices)).cuda()
	else: 		current_batch_2hop_indices = Variable(torch.LongTensor(current_batch_2hop_indices))

	epoch_losses = []   # losses of all epochs
	print("Number of epochs {}".format(args.epochs_gat))
	epoch = 0

	for epoch in range(args.epochs_gat):
		print("\nepoch-> ", epoch)
		random.shuffle(Corpus_.train_triples)
		Corpus_.train_indices = np.array(list(Corpus_.train_triples)).astype(np.int32)

		model_gat.train()  # getting in training mode
		start_time = time.time()
		epoch_loss = []

		if len(Corpus_.train_indices) % args.batch_size_gat == 0:
			num_iters_per_epoch = len(Corpus_.train_indices) // args.batch_size_gat
		else:
			num_iters_per_epoch = (len(Corpus_.train_indices) // args.batch_size_gat) + 1

		for iters in range(num_iters_per_epoch):
			start_time_iter = time.time()
			train_indices, train_values = Corpus_.get_iteration_batch(iters)

			if CUDA:
				train_indices	= Variable(torch.LongTensor(train_indices)).cuda()
				train_values	= Variable(torch.FloatTensor(train_values)).cuda()

			else:
				train_indices	= Variable(torch.LongTensor(train_indices))
				train_values	= Variable(torch.FloatTensor(train_values))

			# forward pass
			entity_embed, relation_embed = model_gat(
				Corpus_, Corpus_.train_adj_matrix, train_indices, current_batch_2hop_indices)

			optimizer.zero_grad()

			loss = batch_gat_loss(gat_loss_func, train_indices, entity_embed, relation_embed)

			loss.backward()
			optimizer.step()

			epoch_loss.append(loss.data.item())

			end_time_iter = time.time()

			print("Iteration-> {0}  , Iteration_time-> {1:.4f} , Iteration_loss {2:.4f}".format(
				iters, end_time_iter - start_time_iter, loss.data.item()))

		# scheduler.step()
		print("Epoch {} , average loss {} , epoch_time {}".format(
			epoch, sum(epoch_loss) / len(epoch_loss), time.time() - start_time))
		epoch_losses.append(sum(epoch_loss) / len(epoch_loss))

		save_model(model_gat, args.data, epoch, args.output_folder)
	save_model(model_gat, args.data, epoch, args.output_folder)

def train_conv(args):

	# Creating convolution model here.
	####################################

	print("Defining model")
	model_gat = SpKBGATModified(entity_embeddings, relation_embeddings, args.entity_out_dim, args.entity_out_dim, args.drop_GAT, args.alpha, args.nheads_GAT)
	print("Only Conv model trained")
	model_conv = SpKBGATConvOnly(entity_embeddings, relation_embeddings, args.entity_out_dim, args.entity_out_dim, args.drop_GAT, args.drop_conv, args.alpha, args.alpha_conv, args.nheads_GAT, args.out_channels)

	if CUDA:
		model_conv.cuda()
		model_gat.cuda()

	model_gat.load_state_dict(torch.load('{}/trained_{}.pth'.format(args.output_folder, args.epochs_gat - 1)))

	model_conv.final_entity_embeddings	= model_gat.final_entity_embeddings
	model_conv.final_relation_embeddings	= model_gat.final_relation_embeddings

	Corpus_.batch_size = args.batch_size_conv
	Corpus_.invalid_valid_ratio = int(args.valid_invalid_ratio_conv)

	optimizer	= torch.optim.Adam(model_conv.parameters(), lr=args.lr, weight_decay=args.weight_decay_conv)
	scheduler	= torch.optim.lr_scheduler.StepLR(optimizer, step_size=25, gamma=0.5, last_epoch=-1)
	margin_loss	= torch.nn.SoftMarginLoss()

	epoch_losses = []   # losses of all epochs
	print("Number of epochs {}".format(args.epochs_conv))

	for epoch in range(args.epochs_conv):
		print("\nepoch-> ", epoch)
		random.shuffle(Corpus_.train_triples)
		Corpus_.train_indices = np.array(list(Corpus_.train_triples)).astype(np.int32)

		model_conv.train()  # getting in training mode
		start_time = time.time()
		epoch_loss = []

		if len(Corpus_.train_indices) % args.batch_size_conv == 0:
			num_iters_per_epoch = len(Corpus_.train_indices) // args.batch_size_conv
		else:
			num_iters_per_epoch = (len(Corpus_.train_indices) // args.batch_size_conv) + 1

		for iters in range(num_iters_per_epoch):
			start_time_iter = time.time()
			train_indices, train_values = Corpus_.get_iteration_batch(iters)

			if CUDA:
				train_indices	= Variable(torch.LongTensor(train_indices)).cuda()
				train_values	= Variable(torch.FloatTensor(train_values)).cuda()
			else:
				train_indices	= Variable(torch.LongTensor(train_indices))
				train_values	= Variable(torch.FloatTensor(train_values))

			preds = model_conv(Corpus_, Corpus_.train_adj_matrix, train_indices)
			optimizer.zero_grad()

			loss = margin_loss(preds.view(-1), train_values.view(-1))

			loss.backward()
			optimizer.step()

			epoch_loss.append(loss.data.item())

			end_time_iter = time.time()

			print("Iteration-> {0}  , Iteration_time-> {1:.4f} , Iteration_loss {2:.4f}".format(iters, end_time_iter - start_time_iter, loss.data.item()))

		# scheduler.step()
		print("Epoch {} , average loss {} , epoch_time {}".format(epoch, sum(epoch_loss) / len(epoch_loss), time.time() - start_time))
		epoch_losses.append(sum(epoch_loss) / len(epoch_loss))

		save_model(model_conv, args.data, epoch, args.output_folder + "conv/")

def evaluate_conv(args, unique_entities, eval_type):
	model_conv = SpKBGATConvOnly(entity_embeddings, relation_embeddings, args.entity_out_dim, args.entity_out_dim,
								 args.drop_GAT, args.drop_conv, args.alpha, args.alpha_conv,
								 args.nheads_GAT, args.out_channels)

	model_conv.load_state_dict(torch.load('{0}conv/trained_{1}.pth'.format(args.output_folder, args.epochs_conv - 1)))

	model_conv.cuda()
	model_conv.eval()
	with torch.no_grad():
		Corpus_.get_validation_pred(args, model_conv, unique_entities, eval_type=eval_type)

train_gat(args)
train_conv(args)
evaluate_conv(args, Corpus_.unique_entities_train, args.eval_type)

max_stars_repo_name: MedyG/kg-reeval

max_issues_repo_name: MedyG/kg-reeval

max_forks_repo_name: MedyG/kg-reeval


text:
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

class ToTensor:

    def __call__(self, x):
        return torch.from_numpy(x).type(torch.float)

class RandomResizedCrop(nn.Module):
    def __init__(self, output_size, scale):
        super(RandomResizedCrop, self).__init__()
        self.output_size = output_size
        self.scale = scale

    def forward(self, x):
        if isinstance(x, np.ndarray):
            x = torch.from_numpy(x).type(torch.float)
        if x.dim() == 2:
            x = x.unsqueeze(0)
        length = x.size(-1)
        crop_size = int(length * self.scale)
        index = torch.randint(0, length - crop_size + 1, size=(1,)).item()
        x = x[..., :, index: index + crop_size]
        x = F.interpolate(x, self.output_size, mode="linear", align_corners=False).squeeze()
        return x

class TwoCropTransform:

    def __init__(self, base_transform):
        self.base_transform = base_transform

    def __call__(self, x):
        q = self.base_transform(x)
        k = self.base_transform(x)
        return [q, k]

class Compose:

    def __init__(self, transforms):
        self.transforms = transforms

    def __call__(self, x):
        for t in self.transforms:
            x = t(x)
        return x

max_stars_repo_name: juaduan/babybrainguardian

max_issues_repo_name: conhua/babybrainguardian

max_forks_repo_name: conhua/babybrainguardian


text:
#!/usr/bin/env python
# -*- coding:utf-8 -*-
# Author: Donny You(youansheng@gmail.com)
# COCO det data generator.

import os
import argparse
import shutil
from PIL import Image
import numpy as np

LABEL_DIR = 'label'
IMAGE_DIR = 'image'

class PascalContextGenerator(object):
    def __init__(self, args, image_dir=IMAGE_DIR, label_dir=LABEL_DIR):
        self.args = args
        self.initialized=False
        self.train_label_dir = os.path.join(self.args.save_dir, 'train', label_dir)
        self.val_label_dir = os.path.join(self.args.save_dir, 'val', label_dir)
        if not os.path.exists(self.train_label_dir):
            os.makedirs(self.train_label_dir)
        else:
            self.initialized=True
        if not os.path.exists(self.val_label_dir):
            os.makedirs(self.val_label_dir)

        self.train_image_dir = os.path.join(self.args.save_dir, 'train', image_dir)
        self.val_image_dir = os.path.join(self.args.save_dir, 'val', image_dir)
        if not os.path.exists(self.train_image_dir):
            os.makedirs(self.train_image_dir)

        if not os.path.exists(self.val_image_dir):
            os.makedirs(self.val_image_dir)

    def _class_to_index(self, mask, _mapping, _key):
        # assert the values
        values = np.unique(mask)
        for i in range(len(values)):
            assert (values[i] in _mapping)
        index = np.digitize(mask.ravel(), _mapping, right=True)
        return _key[index].reshape(mask.shape)

    def generate_label(self):
        if self.initialized:
            return
        _image_dir = os.path.join(self.args.ori_root_dir, 'JPEGImages')
        annFile = os.path.join(self.args.ori_root_dir, 'trainval_merged.json')
        _mapping = np.sort(np.array([
            0, 2, 259, 260, 415, 324, 9, 258, 144, 18, 19, 22,
            23, 397, 25, 284, 158, 159, 416, 33, 162, 420, 454, 295, 296,
            427, 44, 45, 46, 308, 59, 440, 445, 31, 232, 65, 354, 424,
            68, 326, 72, 458, 34, 207, 80, 355, 85, 347, 220, 349, 360,
            98, 187, 104, 105, 366, 189, 368, 113, 115]))
        _key = np.array(range(len(_mapping))).astype('uint8')

        from detail import Detail
        train_detail = Detail(annFile, _image_dir, 'train')
        train_ids = train_detail.getImgs()
        for img_id in train_ids:
            mask = Image.fromarray(self._class_to_index(
                train_detail.getMask(img_id), _mapping=_mapping, _key=_key))
            filename = img_id['file_name']
            basename, _ = os.path.splitext(filename)
            if filename.endswith(".jpg"):
                imgpath = os.path.join(_image_dir, filename)
                shutil.copy(imgpath,
                            os.path.join(self.train_image_dir, filename))
                mask_png_name = basename + '.png'
                mask.save(os.path.join(self.train_label_dir, mask_png_name))

        val_detail = Detail(annFile, _image_dir, 'val')
        val_ids = val_detail.getImgs()
        for img_id in val_ids:
            mask = Image.fromarray(self._class_to_index(
                val_detail.getMask(img_id), _mapping=_mapping, _key=_key))
            filename = img_id['file_name']
            basename, _ = os.path.splitext(filename)
            if filename.endswith(".jpg"):
                imgpath = os.path.join(_image_dir, filename)
                shutil.copy(imgpath,
                            os.path.join(self.val_image_dir, filename))
                mask_png_name = basename + '.png'
                mask.save(os.path.join(self.val_label_dir, mask_png_name))

if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument('--save_dir', default=None, type=str,
                        dest='save_dir', help='The directory to save the data.')
    # ori_root_dir: VOCdevkit/VOC2010
    parser.add_argument('--ori_root_dir', default=None, type=str,
                        dest='ori_root_dir', help='The directory of the cityscapes data.')

    args = parser.parse_args()

    pascalcontext_seg_generator = PascalContextGenerator(args)
    pascalcontext_seg_generator.generate_label()

max_stars_repo_name: MendelXu/ANN

max_issues_repo_name: pinglmlcv/ANN

max_forks_repo_name: pinglmlcv/ANN


text:
import numpy as np

def build_blending_matrix(sz, mode='gaussian'):
    assert mode in ['gaussian', 'bump']
    if mode=='gaussian':
        return blend_gaussian(sz)
    else:
        return blend_bump(sz)

def blend_gaussian(sz, sigma=0.2, mu=0.0):
    """
    Gaussian blending.
    """
    zz, yy, xx = np.meshgrid(np.linspace(-1,1,sz[0], dtype=np.float32),
                             np.linspace(-1,1,sz[1], dtype=np.float32),
                             np.linspace(-1,1,sz[2], dtype=np.float32),
                             indexing='ij')

    dd = np.sqrt(zz*zz + yy*yy + xx*xx)
    ww = 1e-4 + np.exp(-( (dd-mu)**2 / ( 2.0 * sigma**2 )))
    return ww

def blend_bump(sz, t=1.5):
    """
    Bump blending (based on [Lee2017]_ with modifications).

    .. [Lee2017] Lee, Kisuk, Jonathan Zung, Peter Li, Viren Jain, and
        H. Sebastian Seung. "Superhuman accuracy on the SNEMI3D connectomics
        challenge." arXiv preprint arXiv:1706.00120, 2017.
    """
    zz, yy, xx = np.meshgrid(np.linspace(0,1,sz[0]+2, dtype=np.float32)[1:-1],
                             np.linspace(0,1,sz[1]+2, dtype=np.float32)[1:-1],
                             np.linspace(0,1,sz[2]+2, dtype=np.float32)[1:-1],
                             indexing='ij')

    dd = -(xx*(1-xx))**(-t)-(yy*(1-yy))**(-t)-(zz*(1-zz))**(-t)
    ww = 1e-4 + np.exp(dd-dd.max())
    return ww


text:
# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
# SPDX-License-Identifier: MIT-0

import argparse
import base64
import csv
import json
import math
import os
import random
import sys
from multiprocessing import Pool

import cv2
import MatterSim
import numpy as np
import torch
import torch.nn as nn
import torchvision
import torchvision.models
import torchvision.transforms.functional as F
from PIL import Image

from timer import Timer

# sys.path.insert(0, "/root/mount/Matterport3DSimulator/models/")
# import bit_pytorch.models as bit_models

csv.field_size_limit(sys.maxsize)

parser = argparse.ArgumentParser()

parser.add_argument(
    "--model",
    type=str,
    default="ResNet-152",
    choices=[
        "ResNet-152",
        "BiT-M-R50x1",
        "BiT-M-R50x3",
        "BiT-M-R50x1",
        "BiT-M-R101x1",
        "BiT-M-R101x3",
        "BiT-M-R152x4",
        "BiT-S-R50x1",
        "BiT-S-R50x3",
        "BiT-S-R101x1",
        "BiT-S-R101x3",
        "BiT-S-R152x4",
    ],
)
parser.add_argument("--img-features-dir", type=str, default="srv/img_features")
parser.add_argument("--models-dir", type=str, default="models/")
parser.add_argument("--output-feature-file", type=str, default="")
parser.add_argument("--seed", type=int, default=1, help="")
parser.add_argument("--batch-size", type=int, default=12, help="")

args = parser.parse_args()

FEATURE_SIZES = {
    "ResNet-152": 2048,
    "BiT-M-R50x1": 2048,
    "BiT-M-R50x3": 6144,
    "BiT-M-R101x1": 2048,
    "BiT-M-R101x3": 6144,
    "BiT-M-R152x4": 8192,
    "BiT-S-R50x1": 2048,
    "BiT-S-R50x3": 6144,
    "BiT-S-R101x1": 2048,
    "BiT-S-R101x3": 6144,
    "BiT-S-R152x4": 8192,
}

NUM_GPUS = 1
MODEL_NAME = args.model
FEATURE_SIZE = FEATURE_SIZES[MODEL_NAME]
BATCH_SIZE = (
    args.batch_size
)  # Some fraction of viewpoint size - batch size 4 equals 11GB memory

if args.output_feature_file == "":
    OUTFILE = "%s-imagenet-pytorch.tsv" % MODEL_NAME
else:
    OUTFILE = args.output_feature_file

OUTFILE = os.path.join(args.img_features_dir, OUTFILE)
MERGED = OUTFILE

if NUM_GPUS != 1:
    OUTFILE = OUTFILE + ".%d"

MODELS = args.models_dir

GRAPHS = "connectivity/"

SEED = args.seed
print("SEED: %d" % SEED)

# --------------------------------------------
# --------------------------------------------

TSV_FIELDNAMES = ["scanId", "viewpointId", "image_w", "image_h", "vfov", "features"]
VIEWPOINT_SIZE = 36  # Number of discretized views from one viewpoint
GPU_ID = 0

# Simulator image parameters
WIDTH = 640
HEIGHT = 480
VFOV = 60

random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
torch.cuda.manual_seed(SEED)

def load_model(model_name):
    if model_name == "ResNet-152":
        resnet_full = torchvision.models.resnet152(pretrained=True)
        resnet = nn.Sequential(*list(resnet_full.children())[:-1])
        return resnet
    elif "BiT" in model_name:  # BiT-M-R50x1
        model = bit_models.KNOWN_MODELS[model_name](head_size=1000, zero_head=True)
        model.load_from(np.load(MODELS + model_name + ".npz"))

        all_layers = list(model.children())
        main_stack = all_layers[:-1]
        last_layer_wihout_fc = all_layers[-1][:-1]
        model_without_fc = main_stack + [last_layer_wihout_fc]
        bit = nn.Sequential(*model_without_fc)
        return bit

def load_viewpointids(job_id=0):
    viewpointIds = []
    with open(GRAPHS + "scans.txt") as f:
        scans = [scan.strip() for scan in f.readlines()]
        for scan in scans:
            with open(GRAPHS + scan + "_connectivity.json") as j:
                data = json.load(j)
                for item in data:
                    if item["included"]:
                        viewpointIds.append((scan, item["image_id"]))
    random.seed(SEED)
    random.shuffle(viewpointIds)
    if NUM_GPUS != 1:
        viewpointIds = viewpointIds[job_id::NUM_GPUS]
    print("%d: Loaded %d viewpoints" % (job_id, len(viewpointIds)))
    return viewpointIds

def transform_img_resnet(im):
    """ Prep opencv 3 channel image for the network """
    np_im = np.array(im, copy=True).astype(np.float32)
    np_im = np_im[..., ::-1]

    np_im = np_im.transpose((2, 0, 1))  # (3, H, W)
    np_im = np.ascontiguousarray(np_im, dtype=np.float32)
    im = torch.from_numpy(np_im)
    im /= 255.0
    return F.normalize(im, mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])

def transform_img_bit(im):
    np_im = np.array(im, copy=True).astype(np.float32)
    np_im = np_im[..., ::-1]

    np_im = np_im.transpose((2, 0, 1))  # (3, H, W)
    np_im = np.ascontiguousarray(np_im, dtype=np.float32)
    im = torch.from_numpy(np_im)
    im /= 255.0
    return F.normalize(im, mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])

def build_tsv(ids):
    job_id = ids[0]
    gpu_id = ids[1]

    print("JOB ID %d GPU ID %d: build_tsv" % (job_id, gpu_id))

    # Set up the simulator
    sim = MatterSim.Simulator()
    sim.setCameraResolution(WIDTH, HEIGHT)
    sim.setCameraVFOV(math.radians(VFOV))
    sim.setDiscretizedViewingAngles(True)
    sim.setBatchSize(1)
    sim.initialize()

    with torch.no_grad():

        device = torch.device("cuda:%d" % gpu_id)

        model = load_model(MODEL_NAME).to(device)
        model.eval()

        count = 0
        t_render = Timer()
        t_net = Timer()

        if NUM_GPUS == 1:
            output_file = OUTFILE
        else:
            output_file = OUTFILE % job_id

        with open(output_file, "wt") as tsvfile:
            writer = csv.DictWriter(tsvfile, delimiter="\t", fieldnames=TSV_FIELDNAMES)

            # Loop all the viewpoints in the simulator
            viewpointIds = load_viewpointids(job_id)

            for scanId, viewpointId in viewpointIds:
                t_render.tic()
                # Loop all discretized views from this location
                blobs = []
                features = np.empty([VIEWPOINT_SIZE, FEATURE_SIZE], dtype=np.float32)

                for ix in range(VIEWPOINT_SIZE):
                    if ix == 0:
                        sim.newEpisode(
                            [scanId], [viewpointId], [0], [math.radians(-30)]
                        )
                    elif ix % 12 == 0:
                        sim.makeAction([0], [1.0], [1.0])
                    else:
                        sim.makeAction([0], [1.0], [0])

                    state = sim.getState()[0]
                    assert state.viewIndex == ix

                    # Transform and save generated image
                    if "ResNet" in MODEL_NAME:
                        transformed_im = transform_img_resnet(state.rgb)
                    elif "BiT" in MODEL_NAME:
                        transformed_im = transform_img_bit(state.rgb)
                    blobs.append(transformed_im)

                t_render.toc()
                t_net.tic()

                # Run as many forward passes as necessary
                assert VIEWPOINT_SIZE % BATCH_SIZE == 0
                forward_passes = VIEWPOINT_SIZE // BATCH_SIZE

                ix = 0
                data = torch.empty(
                    (BATCH_SIZE, 3, HEIGHT, WIDTH), dtype=torch.float32, device=device
                )

                for f in range(forward_passes):
                    for n in range(BATCH_SIZE):
                        # Copy image blob to the net
                        data[n, :, :, :] = blobs[ix]
                        ix += 1
                    # Forward pass
                    features[f * BATCH_SIZE : (f + 1) * BATCH_SIZE, :] = (
                        model(data).squeeze().cpu().detach().numpy()
                    )
                writer.writerow(
                    {
                        "scanId": scanId,
                        "viewpointId": viewpointId,
                        "image_w": WIDTH,
                        "image_h": HEIGHT,
                        "vfov": VFOV,
                        "features": str(base64.b64encode(features), "utf-8"),
                    }
                )
                count += 1
                t_net.toc()
                if count % 100 == 0:
                    print(
                        "Processed %d / %d viewpoints, %.1fs avg render time, %.1fs avg net time, projected %.1f hours"
                        % (
                            count,
                            len(viewpointIds),
                            t_render.average_time,
                            t_net.average_time,
                            (t_render.average_time + t_net.average_time)
                            * len(viewpointIds)
                            / 3600,
                        )
                    )

def merge_tsvs():
    test = [OUTFILE % i for i in range(NUM_GPUS)]
    with open(MERGED, "wt") as tsvfile:
        writer = csv.DictWriter(tsvfile, delimiter="\t", fieldnames=TSV_FIELDNAMES)
        for infile in test:
            print(infile)
            with open(infile, "rt") as tsv_in_files:
                reader = csv.DictReader(
                    tsv_in_files, delimiter="\t", fieldnames=TSV_FIELDNAMES
                )
                for item in reader:
                    try:
                        writer.writerow(item)
                    except Exception as e:
                        print(e)
                        print(item["image_id"])

def read_tsv(infile):
    # Verify we can read a tsv
    in_data = []
    with open(infile, "rt") as tsv_in_file:
        reader = csv.DictReader(tsv_in_file, delimiter="\t", fieldnames=TSV_FIELDNAMES)
        for item in reader:
            item["scanId"] = item["scanId"]
            item["image_h"] = int(item["image_h"])
            item["image_w"] = int(item["image_w"])
            item["vfov"] = int(item["vfov"])
            item["features"] = np.frombuffer(
                base64.b64decode(item["features"]), dtype=np.float32
            ).reshape((VIEWPOINT_SIZE, FEATURE_SIZE))
            in_data.append(item)
    return in_data

if __name__ == "__main__":

    if NUM_GPUS == 1:
        build_tsv()
        data = read_tsv(OUTFILE)
        print("Completed %d viewpoints" % len(data))
    else:
        # JOB, GPU

        ids = [[0, 0], [1, 1], [2, 2], [3, 3]]
        p = Pool(4)
        p.map(build_tsv, ids)
        merge_tsvs()
        data = read_tsv(MERGED)
        print("Completed %d viewpoints" % len(data))

max_stars_repo_name: ayshrv/visitron

max_issues_repo_name: ayshrv/visitron

max_forks_repo_name: ayshrv/visitron


text:
""" Module to manage OctreeSettings """

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import numpy as np
import yaml

class OctreeSettings:
    """ Octree settings variables"""
    def __init__(self,
                 depth=6,
                 full_depth=2,
                 node_displacement=True,
                 node_feature=False,
                 split_label=False,
                 adaptive=False,
                 adaptive_depth=4,
                 threshold_distance=0.866,
                 threshold_normal=0.2,
                 key2xyz=False):
        """ Initializes OctreeSettings
        Args:
          depth: Maximum depth of the octree.
          full_depth: Full layer of the octree.
          node_displacement: Output per-node displacement.
          node_feature: Compute per node feature.
          split_label: Compute per node splitting label.
          adaptive: Build adaptive octree.
          adaptive_depth: Starting depth of adaptive octree.
          threshold_distance: The threshold for simplifying octree.
          threshold_normal: The threshold for simplying octree.
          key2xyz: Convert the key to xyz when serialized.
        """
        self.depth = depth
        self.full_depth = full_depth
        self.node_displacement = node_displacement
        self.node_feature = node_feature
        self.split_label = split_label
        self.adaptive = adaptive
        self.adaptive_depth = adaptive_depth
        self.threshold_distance = threshold_distance
        self.threshold_normal = threshold_normal
        self.key2xyz = key2xyz

    @classmethod
    def from_dict(cls, dictionary):
        """ Creates OctreeSettings from Dictionary
        Args:
          dictionary: dictionary with keys as member variables.
        """
        return cls(
            depth=dictionary['depth'],
            full_depth=dictionary['full_depth'],
            node_displacement=dictionary['node_displacement'],
            node_feature=dictionary['node_feature'],
            split_label=dictionary['split_label'],
            adaptive=dictionary['adaptive'],
            adaptive_depth=dictionary['adaptive_depth'],
            threshold_distance=dictionary['threshold_distance'],
            threshold_normal=dictionary['threshold_normal'],
            key2xyz=dictionary['key2xyz'])

    @classmethod
    def from_yaml(cls, yml_filepath):
        """ Creates OctreeSettings from YAML
        Args:
          yml_filepath: Path to yml config file.
        """

        with open(yml_filepath, 'r') as yml_file:
            config = yaml.load(yml_file)

        #Octree Settings
        parameters = config['octree_settings']
        return cls.from_dict(parameters)

    def write_yaml(self, yml_filepath):
        """ Writes ScannerSettings to YAML
        Args:
          yml_filepath: Filepath to output settings
        """
        data = {'octree_settings': {
            'depth': self.depth,
            'full_depth': self.full_depth,
            'node_displacement': self.node_displacement,
            'node_feature': self.node_feature,
            'split_label': self.split_label,
            'adaptive': self.adaptive,
            'adaptive_depth': self.adaptive_depth,
            'threshold_distance': self.threshold_distance,
            'threshold_normal': self.threshold_normal,
            'key2xyz': self.key2xyz}}

        with open(yml_filepath, 'w') as yml_file:
            yaml.dump(data, yml_file, default_flow_style=False)

    def __eq__(self, other):
        if self.__class__ is other.__class__:
            return (self.depth == other.depth and
                    self.full_depth == other.full_depth and
                    self.node_displacement == other.node_displacement and
                    self.node_feature == other.node_feature and
                    self.split_label == other.split_label and
                    self.adaptive == other.adaptive and
                    self.adaptive_depth == other.adaptive_depth and
                    np.isclose(self.threshold_distance, other.treshold_distance) and
                    np.isclose(self.threshold_normal, other.threshold_normal) and
                    self.key2xyz == other.key2xyz)
        return NotImplemented

    def __ne__(self, other):
        return not self == other

max_stars_repo_name: FrozenSilent/O-CNN

max_issues_repo_name: FrozenSilent/O-CNN

max_forks_repo_name: FrozenSilent/O-CNN