text:
# Licensed under a 3-clause BSD style license - see LICENSE.rst
import copy
from functools import lru_cache
import numpy as np
import astropy.units as u
from astropy.convolution import Tophat2DKernel
from astropy.coordinates import Angle, SkyCoord
from astropy.io import fits
from astropy.nddata import Cutout2D
from astropy.nddata.utils import overlap_slices
from astropy.utils import lazyproperty
from astropy.wcs import WCS
from astropy.wcs.utils import (
    celestial_frame_to_wcs,
    proj_plane_pixel_scales,
    wcs_to_celestial_frame,
)
from gammapy.utils.array import round_up_to_even, round_up_to_odd
from ..axes import MapAxes
from ..coord import MapCoord, skycoord_to_lonlat
from ..geom import Geom, get_shape, pix_tuple_to_idx
from ..utils import INVALID_INDEX, _check_binsz, _check_width

__all__ = ["WcsGeom"]

def cast_to_shape(param, shape, dtype):
    """Cast a tuple of parameter arrays to a given shape."""
    if not isinstance(param, tuple):
        param = [param]

    param = [np.array(p, ndmin=1, dtype=dtype) for p in param]

    if len(param) == 1:
        param = [param[0].copy(), param[0].copy()]

    for i, p in enumerate(param):

        if p.size > 1 and p.shape != shape:
            raise ValueError

        if p.shape == shape:
            continue

        param[i] = p * np.ones(shape, dtype=dtype)

    return tuple(param)

def get_resampled_wcs(wcs, factor, downsampled):
    """
    Get resampled WCS object.
    """
    wcs = wcs.deepcopy()

    if not downsampled:
        factor = 1.0 / factor

    wcs.wcs.cdelt *= factor
    wcs.wcs.crpix = (wcs.wcs.crpix - 0.5) / factor + 0.5
    return wcs

class WcsGeom(Geom):
    """Geometry class for WCS maps.

    This class encapsulates both the WCS transformation object and the
    the image extent (number of pixels in each dimension).  Provides
    methods for accessing the properties of the WCS object and
    performing transformations between pixel and world coordinates.

    Parameters
    ----------
    wcs : `~astropy.wcs.WCS`
        WCS projection object
    npix : tuple
        Number of pixels in each spatial dimension
    cdelt : tuple
        Pixel size in each image plane.  If none then a constant pixel size will be used.
    crpix : tuple
        Reference pixel coordinate in each image plane.
    axes : list
        Axes for non-spatial dimensions
    """

    _slice_spatial_axes = slice(0, 2)
    _slice_non_spatial_axes = slice(2, None)
    is_hpx = False
    is_region = False

    def __init__(self, wcs, npix, cdelt=None, crpix=None, axes=None):
        self._wcs = wcs
        self._frame = wcs_to_celestial_frame(wcs).name
        self._projection = wcs.wcs.ctype[0][5:]
        self._axes = MapAxes.from_default(axes, n_spatial_axes=2)

        if cdelt is None:
            cdelt = tuple(np.abs(self.wcs.wcs.cdelt))

        # Shape to use for WCS transformations
        wcs_shape = max([get_shape(t) for t in [npix, cdelt]])
        self._npix = cast_to_shape(npix, wcs_shape, int)
        self._cdelt = cast_to_shape(cdelt, wcs_shape, float)

        # By convention CRPIX is indexed from 1
        if crpix is None:
            crpix = tuple(1.0 + (np.array(self._npix) - 1.0) / 2.0)

        self._crpix = crpix

        # define cached methods
        self.get_coord = lru_cache()(self.get_coord)
        self.get_pix = lru_cache()(self.get_pix)

    def __setstate__(self, state):
        for key, value in state.items():
            if key in ["get_coord", "get_pix"]:
                state[key] = lru_cache()(value)

        self.__dict__ = state

    @property
    def data_shape(self):
        """Shape of the Numpy data array matching this geometry."""
        return self._shape[::-1]

    @property
    def axes_names(self):
        """All axes names"""
        return ["lon", "lat"] + self.axes.names

    @property
    def data_shape_axes(self):
        """Shape of data of the non-spatial axes and unit spatial axes."""
        return self.axes.shape[::-1] + (1, 1)

    @property
    def _shape(self):
        npix_shape = tuple([np.max(self.npix[0]), np.max(self.npix[1])])
        return npix_shape + self.axes.shape

    @property
    def _shape_edges(self):
        npix_shape = tuple([np.max(self.npix[0]) + 1, np.max(self.npix[1]) + 1])
        return npix_shape + self.axes.shape

    @property
    def shape_axes(self):
        """Shape of non-spatial axes."""
        return self._shape[self._slice_non_spatial_axes]

    @property
    def wcs(self):
        """WCS projection object."""
        return self._wcs

    @property
    def frame(self):
        """Coordinate system of the projection.

        Galactic ("galactic") or Equatorial ("icrs").
        """
        return self._frame

    def cutout_slices(self, geom, mode="partial"):
        """Compute cutout slices.

        Parameters
        ----------
        geom : `WcsGeom`
            Parent geometry
        mode : {"trim", "partial", "strict"}
            Cutout slices mode.

        Returns
        -------
        slices : dict
            Dictionary containing "parent-slices" and "cutout-slices".
        """
        position = geom.to_image().coord_to_pix(self.center_skydir)
        slices = overlap_slices(
            large_array_shape=geom.data_shape[-2:],
            small_array_shape=self.data_shape[-2:],
            position=position[::-1],
            mode=mode,
        )
        return {
            "parent-slices": slices[0],
            "cutout-slices": slices[1],
        }

    @property
    def projection(self):
        """Map projection."""
        return self._projection

    @property
    def is_allsky(self):
        """Flag for all-sky maps."""
        if np.all(np.isclose(self._npix[0] * self._cdelt[0], 360.0)) and np.all(
            np.isclose(self._npix[1] * self._cdelt[1], 180.0)
        ):
            return True
        else:
            return False

    @property
    def is_regular(self):
        """Is this geometry is regular in non-spatial dimensions (bool)?

        - False for multi-resolution or irregular geometries.
        - True if all image planes have the same pixel geometry.
        """
        if self.npix[0].size > 1:
            return False
        else:
            return True

    @property
    def width(self):
        """Tuple with image dimension in deg in longitude and latitude."""
        dlon = self._cdelt[0] * self._npix[0]
        dlat = self._cdelt[1] * self._npix[1]
        return (dlon, dlat) * u.deg

    @property
    def pixel_area(self):
        """Pixel area in deg^2."""
        # FIXME: Correctly compute solid angle for projection
        return self._cdelt[0] * self._cdelt[1]

    @property
    def npix(self):
        """Tuple with image dimension in pixels in longitude and latitude."""
        return self._npix

    @property
    def axes(self):
        """List of non-spatial axes."""
        return self._axes

    @property
    def ndim(self):
        return len(self.data_shape)

    @property
    def center_coord(self):
        """Map coordinate of the center of the geometry.

        Returns
        -------
        coord : tuple
        """
        return self.pix_to_coord(self.center_pix)

    @property
    def center_pix(self):
        """Pixel coordinate of the center of the geometry.

        Returns
        -------
        pix : tuple
        """
        return tuple((np.array(self.data_shape) - 1.0) / 2)[::-1]

    @property
    def center_skydir(self):
        """Sky coordinate of the center of the geometry.

        Returns
        -------
        pix : `~astropy.coordinates.SkyCoord`
        """
        return SkyCoord.from_pixel(self.center_pix[0], self.center_pix[1], self.wcs)

    @property
    def pixel_scales(self):
        """
        Pixel scale.

        Returns angles along each axis of the image at the CRPIX location once
        it is projected onto the plane of intermediate world coordinates.

        Returns
        -------
        angle: `~astropy.coordinates.Angle`
        """
        return Angle(proj_plane_pixel_scales(self.wcs), "deg")

    @classmethod
    def create(
        cls,
        npix=None,
        binsz=0.5,
        proj="CAR",
        frame="icrs",
        refpix=None,
        axes=None,
        skydir=None,
        width=None,
    ):
        """Create a WCS geometry object.

        Pixelization of the map is set with
        ``binsz`` and one of either ``npix`` or ``width`` arguments.
        For maps with non-spatial dimensions a different pixelization
        can be used for each image plane by passing a list or array
        argument for any of the pixelization parameters.  If both npix
        and width are None then an all-sky geometry will be created.

        Parameters
        ----------
        npix : int or tuple or list
            Width of the map in pixels. A tuple will be interpreted as
            parameters for longitude and latitude axes.  For maps with
            non-spatial dimensions, list input can be used to define a
            different map width in each image plane.  This option
            supersedes width.
        width : float or tuple or list or string
            Width of the map in degrees.  A tuple will be interpreted
            as parameters for longitude and latitude axes.  For maps
            with non-spatial dimensions, list input can be used to
            define a different map width in each image plane.
        binsz : float or tuple or list
            Map pixel size in degrees.  A tuple will be interpreted
            as parameters for longitude and latitude axes.  For maps
            with non-spatial dimensions, list input can be used to
            define a different bin size in each image plane.
        skydir : tuple or `~astropy.coordinates.SkyCoord`
            Sky position of map center.  Can be either a SkyCoord
            object or a tuple of longitude and latitude in deg in the
            coordinate system of the map.
        frame : {"icrs", "galactic"}, optional
            Coordinate system, either Galactic ("galactic") or Equatorial ("icrs").
        axes : list
            List of non-spatial axes.
        proj : string, optional
            Any valid WCS projection type. Default is 'CAR' (cartesian).
        refpix : tuple
            Reference pixel of the projection.  If None this will be
            set to the center of the map.

        Returns
        -------
        geom : `~WcsGeom`
            A WCS geometry object.

        Examples
        --------
        >>> from gammapy.maps import WcsGeom
        >>> from gammapy.maps import MapAxis
        >>> axis = MapAxis.from_bounds(0,1,2)
        >>> geom = WcsGeom.create(npix=(100,100), binsz=0.1)
        >>> geom = WcsGeom.create(npix=(100,100), binsz="0.1deg")
        >>> geom = WcsGeom.create(npix=[100,200], binsz=[0.1,0.05], axes=[axis])
        >>> geom = WcsGeom.create(npix=[100,200], binsz=["0.1deg","0.05deg"], axes=[axis])
        >>> geom = WcsGeom.create(width=[5.0,8.0], binsz=[0.1,0.05], axes=[axis])
        >>> geom = WcsGeom.create(npix=([100,200],[100,200]), binsz=0.1, axes=[axis])
        """
        if skydir is None:
            crval = (0.0, 0.0)
        elif isinstance(skydir, tuple):
            crval = skydir
        elif isinstance(skydir, SkyCoord):
            xref, yref, frame = skycoord_to_lonlat(skydir, frame=frame)
            crval = (xref, yref)
        else:
            raise ValueError(f"Invalid type for skydir: {type(skydir)!r}")

        if width is not None:
            width = _check_width(width)

        binsz = _check_binsz(binsz)

        shape = max([get_shape(t) for t in [npix, binsz, width]])
        binsz = cast_to_shape(binsz, shape, float)

        # If both npix and width are None then create an all-sky geometry
        if npix is None and width is None:
            width = (360.0, 180.0)

        if npix is None:
            width = cast_to_shape(width, shape, float)
            npix = (
                np.rint(width[0] / binsz[0]).astype(int),
                np.rint(width[1] / binsz[1]).astype(int),
            )
        else:
            npix = cast_to_shape(npix, shape, int)

        if refpix is None:
            nxpix = int(npix[0].flat[0])
            nypix = int(npix[1].flat[0])
            refpix = ((nxpix + 1) / 2.0, (nypix + 1) / 2.0)

        # get frame class
        frame = SkyCoord(np.nan, np.nan, frame=frame, unit="deg").frame
        wcs = celestial_frame_to_wcs(frame, projection=proj)
        wcs.wcs.crpix = refpix
        wcs.wcs.crval = crval

        cdelt = (-binsz[0].flat[0], binsz[1].flat[0])
        wcs.wcs.cdelt = cdelt

        wcs.array_shape = npix[0].flat[0], npix[1].flat[0]
        wcs.wcs.datfix()
        return cls(wcs, npix, cdelt=binsz, axes=axes)

    @property
    def footprint(self):
        """Footprint of the geometry"""
        coords = self.wcs.calc_footprint()
        return SkyCoord(coords, frame=self.frame, unit="deg")

    @classmethod
    def from_aligned(cls, geom, skydir, width):
        """Create an aligned geometry from an existing one

        Parameters
        ----------
        geom : `~WcsGeom`
            A reference WCS geometry object.
        skydir : tuple or `~astropy.coordinates.SkyCoord`
            Sky position of map center.  Can be either a SkyCoord
            object or a tuple of longitude and latitude in deg in the
            coordinate system of the map.
        width : float or tuple or list or string
            Width of the map in degrees.  A tuple will be interpreted
            as parameters for longitude and latitude axes.  For maps
            with non-spatial dimensions, list input can be used to
            define a different map width in each image plane.

        Returns
        -------
        geom : `~WcsGeom`
            An aligned WCS geometry object with specified size and center.

        """
        width = _check_width(width) * u.deg
        npix = tuple(np.round(width / geom.pixel_scales).astype(int))
        xref, yref = geom.to_image().coord_to_pix(skydir)
        xref = int(np.floor(-xref + npix[0] / 2.0)) + geom.wcs.wcs.crpix[0]
        yref = int(np.floor(-yref + npix[1] / 2.0)) + geom.wcs.wcs.crpix[1]
        return cls.create(
            skydir=tuple(geom.wcs.wcs.crval),
            npix=npix,
            refpix=(xref, yref),
            frame=geom.frame,
            binsz=tuple(geom.pixel_scales.deg),
            axes=geom.axes,
            proj=geom.projection,
        )

    @classmethod
    def from_header(cls, header, hdu_bands=None, format="gadf"):
        """Create a WCS geometry object from a FITS header.

        Parameters
        ----------
        header : `~astropy.io.fits.Header`
            The FITS header
        hdu_bands : `~astropy.io.fits.BinTableHDU`
            The BANDS table HDU.
        format : {'gadf', 'fgst-ccube','fgst-template'}
            FITS format convention.

        Returns
        -------
        wcs : `~WcsGeom`
            WCS geometry object.
        """
        wcs = WCS(header, naxis=2)
        # TODO: see https://github.com/astropy/astropy/issues/9259
        wcs._naxis = wcs._naxis[:2]

        axes = MapAxes.from_table_hdu(hdu_bands, format=format)
        shape = axes.shape

        if hdu_bands is not None and "NPIX" in hdu_bands.columns.names:
            npix = hdu_bands.data.field("NPIX").reshape(shape + (2,))
            npix = (npix[..., 0], npix[..., 1])
            cdelt = hdu_bands.data.field("CDELT").reshape(shape + (2,))
            cdelt = (cdelt[..., 0], cdelt[..., 1])
        elif "WCSSHAPE" in header:
            wcs_shape = eval(header["WCSSHAPE"])
            npix = (wcs_shape[0], wcs_shape[1])
            cdelt = None
            wcs.array_shape = npix
        else:
            npix = (header["NAXIS1"], header["NAXIS2"])
            cdelt = None

        return cls(wcs, npix, cdelt=cdelt, axes=axes)

    def _make_bands_cols(self):

        cols = []
        if not self.is_regular:
            cols += [
                fits.Column(
                    "NPIX",
                    "2I",
                    dim="(2)",
                    array=np.vstack((np.ravel(self.npix[0]), np.ravel(self.npix[1]))).T,
                )
            ]
            cols += [
                fits.Column(
                    "CDELT",
                    "2E",
                    dim="(2)",
                    array=np.vstack(
                        (np.ravel(self._cdelt[0]), np.ravel(self._cdelt[1]))
                    ).T,
                )
            ]
            cols += [
                fits.Column(
                    "CRPIX",
                    "2E",
                    dim="(2)",
                    array=np.vstack(
                        (np.ravel(self._crpix[0]), np.ravel(self._crpix[1]))
                    ).T,
                )
            ]
        return cols

    def to_header(self):
        header = self.wcs.to_header()
        header.update(self.axes.to_header())
        shape = "{},{}".format(np.max(self.npix[0]), np.max(self.npix[1]))
        for ax in self.axes:
            shape += f",{ax.nbin}"

        header["WCSSHAPE"] = f"({shape})"
        return header

    def get_idx(self, idx=None, flat=False):
        pix = self.get_pix(idx=idx, mode="center")
        if flat:
            pix = tuple([p[np.isfinite(p)] for p in pix])
        return pix_tuple_to_idx(pix)

    def _get_pix_all(
        self, idx=None, mode="center", sparse=False, axis_name=("lon", "lat")
    ):
        """Get idx coordinate array without footprint of the projection applied"""
        pix_all = []

        for name, nbin in zip(self.axes_names, self._shape):
            if mode == "edges" and name in axis_name:
                pix = np.arange(-0.5, nbin, dtype=float)
            else:
                pix = np.arange(nbin, dtype=float)

            pix_all.append(pix)

        # TODO: improve varying bin size coordinate handling
        if idx is not None:
            pix_all = pix_all[self._slice_spatial_axes] + [float(t) for t in idx]

        return np.meshgrid(*pix_all[::-1], indexing="ij", sparse=sparse)[::-1]

    def get_pix(self, idx=None, mode="center"):
        """Get map pix coordinates from the geometry.

        Parameters
        ----------
        mode : {'center', 'edges'}
            Get center or edge pix coordinates for the spatial axes.

        Returns
        -------
        coord : tuple
            Map pix coordinate tuple.
        """
        pix = self._get_pix_all(idx=idx, mode=mode)
        coords = self.pix_to_coord(pix)
        m = np.isfinite(coords[0])
        for _ in pix:
            _[~m] = INVALID_INDEX.float
        return pix

    def get_coord(
        self, idx=None, mode="center", frame=None, sparse=False, axis_name=None
    ):
        """Get map coordinates from the geometry.

        Parameters
        ----------
        mode : {'center', 'edges'}
            Get center or edge coordinates for the spatial axes.
        frame : str or `~astropy.coordinates.Frame`
            Coordinate frame
        sparse : bool
            Compute sparse coordinates
        axis_name : str
            If mode = "edges", the edges will be returned for this axis.

        Returns
        -------
        coord : `~MapCoord`
            Map coordinate object.
        """
        if axis_name is None:
            axis_name = ("lon", "lat")

        if frame is None:
            frame = self.frame

        pix = self._get_pix_all(idx=idx, mode=mode, sparse=sparse, axis_name=axis_name)

        data = self.pix_to_coord(pix)

        coords = MapCoord.create(
            data=data, frame=self.frame, axis_names=self.axes.names
        )
        return coords.to_frame(frame)

    def coord_to_pix(self, coords):
        coords = MapCoord.create(coords, frame=self.frame, axis_names=self.axes.names)

        if coords.size == 0:
            return tuple([np.array([]) for i in range(coords.ndim)])

        # Variable Bin Size
        if not self.is_regular:
            idxs = self.axes.coord_to_idx(coords, clip=True)
            crpix = [t[idxs] for t in self._crpix]
            cdelt = [t[idxs] for t in self._cdelt]
            pix = world2pix(self.wcs, cdelt, crpix, (coords.lon, coords.lat))
            pix = list(pix)
        else:
            pix = self._wcs.wcs_world2pix(coords.lon, coords.lat, 0)

        pix += self.axes.coord_to_pix(coords)
        return tuple(pix)

    def pix_to_coord(self, pix):
        # Variable Bin Size
        if not self.is_regular:
            idxs = pix_tuple_to_idx(pix[self._slice_non_spatial_axes])
            crpix = [t[idxs] for t in self._crpix]
            cdelt = [t[idxs] for t in self._cdelt]
            coords = pix2world(self.wcs, cdelt, crpix, pix[self._slice_spatial_axes])
        else:
            coords = self._wcs.wcs_pix2world(pix[0], pix[1], 0)

        coords = (
            u.Quantity(coords[0], unit="deg", copy=False),
            u.Quantity(coords[1], unit="deg", copy=False),
        )

        coords += self.axes.pix_to_coord(pix[self._slice_non_spatial_axes])
        return coords

    def pix_to_idx(self, pix, clip=False):
        pix = pix_tuple_to_idx(pix)

        idx_non_spatial = self.axes.pix_to_idx(
            pix[self._slice_non_spatial_axes], clip=clip
        )

        if not self.is_regular:
            npix = (self.npix[0][idx_non_spatial], self.npix[1][idx_non_spatial])
        else:
            npix = self.npix

        idx_spatial = []

        for idx, npix_ in zip(pix[self._slice_spatial_axes], npix):
            if clip:
                idx = np.clip(idx, 0, npix_)
            else:
                idx = np.where((idx < 0) | (idx >= npix_), -1, idx)

            idx_spatial.append(idx)

        return tuple(idx_spatial) + idx_non_spatial

    def contains(self, coords):
        idx = self.coord_to_idx(coords)
        return np.all(np.stack([t != INVALID_INDEX.int for t in idx]), axis=0)

    def to_image(self):
        return self._image_geom

    @lazyproperty
    def _image_geom(self):
        npix = (np.max(self._npix[0]), np.max(self._npix[1]))
        cdelt = (np.max(self._cdelt[0]), np.max(self._cdelt[1]))
        return self.__class__(self._wcs, npix, cdelt=cdelt)

    def to_cube(self, axes):
        npix = (np.max(self._npix[0]), np.max(self._npix[1]))
        cdelt = (np.max(self._cdelt[0]), np.max(self._cdelt[1]))
        axes = copy.deepcopy(self.axes) + axes
        return self.__class__(
            self._wcs.deepcopy(),
            npix,
            cdelt=cdelt,
            axes=axes,
        )

    def _pad_spatial(self, pad_width):
        if np.isscalar(pad_width):
            pad_width = (pad_width, pad_width)

        npix = (self.npix[0] + 2 * pad_width[0], self.npix[1] + 2 * pad_width[1])
        wcs = self._wcs.deepcopy()
        wcs.wcs.crpix += np.array(pad_width)
        cdelt = copy.deepcopy(self._cdelt)
        return self.__class__(wcs, npix, cdelt=cdelt, axes=copy.deepcopy(self.axes))

    def crop(self, crop_width):
        if np.isscalar(crop_width):
            crop_width = (crop_width, crop_width)

        npix = (self.npix[0] - 2 * crop_width[0], self.npix[1] - 2 * crop_width[1])
        wcs = self._wcs.deepcopy()
        wcs.wcs.crpix -= np.array(crop_width)
        cdelt = copy.deepcopy(self._cdelt)
        return self.__class__(wcs, npix, cdelt=cdelt, axes=copy.deepcopy(self.axes))

    def downsample(self, factor, axis_name=None):
        if axis_name is None:
            if np.any(np.mod(self.npix, factor) > 0):
                raise ValueError(
                    f"Spatial shape not divisible by factor {factor!r} in all axes."
                    f" You need to pad prior to calling downsample."
                )

            npix = (self.npix[0] / factor, self.npix[1] / factor)
            cdelt = (self._cdelt[0] * factor, self._cdelt[1] * factor)
            wcs = get_resampled_wcs(self.wcs, factor, True)
            return self._init_copy(wcs=wcs, npix=npix, cdelt=cdelt)
        else:
            if not self.is_regular:
                raise NotImplementedError(
                    "Upsampling in non-spatial axes not supported for irregular geometries"
                )
            axes = self.axes.downsample(factor=factor, axis_name=axis_name)
            return self._init_copy(axes=axes)

    def upsample(self, factor, axis_name=None):
        if axis_name is None:
            npix = (self.npix[0] * factor, self.npix[1] * factor)
            cdelt = (self._cdelt[0] / factor, self._cdelt[1] / factor)
            wcs = get_resampled_wcs(self.wcs, factor, False)
            return self._init_copy(wcs=wcs, npix=npix, cdelt=cdelt)
        else:
            if not self.is_regular:
                raise NotImplementedError(
                    "Upsampling in non-spatial axes not supported for irregular geometries"
                )
            axes = self.axes.upsample(factor=factor, axis_name=axis_name)
            return self._init_copy(axes=axes)

    def to_binsz(self, binsz):
        """Change pixel size of the geometry.

        Parameters
        ----------
        binsz : float or tuple or list
            New pixel size in degree.

        Returns
        -------
        geom : `WcsGeom`
            Geometry with new pixel size.
        """
        return self.create(
            skydir=self.center_skydir,
            binsz=binsz,
            width=self.width,
            proj=self.projection,
            frame=self.frame,
            axes=copy.deepcopy(self.axes),
        )

    def solid_angle(self):
        """Solid angle array (`~astropy.units.Quantity` in ``sr``).

        The array has the same dimension as the WcsGeom object.

        To return solid angles for the spatial dimensions only use::

             WcsGeom.to_image().solid_angle()
        """
        return self._solid_angle

    @lazyproperty
    def _solid_angle(self):
        coord = self.get_coord(mode="edges").skycoord

        # define pixel corners
        low_left = coord[..., :-1, :-1]
        low_right = coord[..., 1:, :-1]
        up_left = coord[..., :-1, 1:]
        up_right = coord[..., 1:, 1:]

        # compute side lengths
        low = low_left.separation(low_right)
        left = low_left.separation(up_left)
        up = up_left.separation(up_right)
        right = low_right.separation(up_right)

        # compute enclosed angles
        angle_low_right = low_right.position_angle(up_right) - low_right.position_angle(
            low_left
        )
        angle_up_left = up_left.position_angle(up_right) - low_left.position_angle(
            up_left
        )

        # compute area assuming a planar triangle
        area_low_right = 0.5 * low * right * np.sin(angle_low_right)
        area_up_left = 0.5 * up * left * np.sin(angle_up_left)
        # TODO: for non-negative cdelt a negative solid angle is returned
        #  find out why and fix properly
        return np.abs(u.Quantity(area_low_right + area_up_left, "sr", copy=False))

    def bin_volume(self):
        """Bin volume (`~astropy.units.Quantity`)"""
        return self._bin_volume

    @lazyproperty
    def _bin_volume(self):
        """Cached property of bin volume"""
        value = self.to_image().solid_angle()

        if not self.is_image:
            value = value * self.axes.bin_volume()

        return value

    def separation(self, center):
        """Compute sky separation wrt a given center.

        Parameters
        ----------
        center : `~astropy.coordinates.SkyCoord`
            Center position

        Returns
        -------
        separation : `~astropy.coordinates.Angle`
            Separation angle array (2D)
        """
        coord = self.to_image().get_coord()
        return center.separation(coord.skycoord)

    def cutout(self, position, width, mode="trim", odd_npix=False):
        """
        Create a cutout around a given position.

        Parameters
        ----------
        position : `~astropy.coordinates.SkyCoord`
            Center position of the cutout region.
        width : tuple of `~astropy.coordinates.Angle`
            Angular sizes of the region in (lon, lat) in that specific order.
            If only one value is passed, a square region is extracted.
        mode : {'trim', 'partial', 'strict'}
            Mode option for Cutout2D, for details see `~astropy.nddata.utils.Cutout2D`.
        odd_npix : bool
            Force width to odd number of pixels.

        Returns
        -------
        cutout : `~gammapy.maps.WcsNDMap`
            Cutout map
        """
        width = _check_width(width) * u.deg

        binsz = self.pixel_scales
        width_npix = np.clip((width / binsz).to_value(""), 1, None)
        width = width_npix * binsz

        if odd_npix:
            width = round_up_to_odd(width_npix)

        dummy_data = np.empty(self.to_image().data_shape)
        c2d = Cutout2D(
            data=dummy_data,
            wcs=self.wcs,
            position=position,
            # Cutout2D takes size with order (lat, lon)
            size=width[::-1],
            mode=mode,
        )
        return self._init_copy(wcs=c2d.wcs, npix=c2d.shape[::-1])

    def boundary_mask(self, width):
        """Create a mask applying binary erosion with a given width from geom edges

        Parameters
        ----------
        width : tuple of `~astropy.units.Quantity`
            Angular sizes of the margin in (lon, lat) in that specific order.
            If only one value is passed, the same margin is applied in (lon, lat).

        Returns
        -------
        mask_map : `~gammapy.maps.WcsNDMap` of boolean type
            Boundary mask

        """
        from .ndmap import WcsNDMap

        data = np.ones(self.data_shape, dtype=bool)
        return WcsNDMap.from_geom(self, data=data).binary_erode(
            width=2 * u.Quantity(width), kernel="box"
        )

    def region_mask(self, regions, inside=True):
        """Create a mask from a given list of regions

        The mask is filled such that a pixel inside the region is filled with
        "True". To invert the mask, e.g. to create a mask with exclusion regions
        the tilde (~) operator can be used (see example below).

        Parameters
        ----------
        regions : str, `~regions.Region` or list of `~regions.Region`
            Region or list of regions (pixel or sky regions accepted).
            A region can be defined as a string ind DS9 format as well.
            See http://ds9.si.edu/doc/ref/region.html for details.
        inside : bool
            For ``inside=True``, pixels in the region to True (the default).
            For ``inside=False``, pixels in the region are False.

        Returns
        -------
        mask_map : `~gammapy.maps.WcsNDMap` of boolean type
            Boolean region mask

        Examples
        --------
        Make an exclusion mask for a circular region::

            from regions import CircleSkyRegion
            from astropy.coordinates import SkyCoord, Angle
            from gammapy.maps import WcsNDMap, WcsGeom

            pos = SkyCoord(0, 0, unit='deg')
            geom = WcsGeom.create(skydir=pos, npix=100, binsz=0.1)

            region = CircleSkyRegion(
                SkyCoord(3, 2, unit='deg'),
                Angle(1, 'deg'),
            )

            # the Gammapy convention for exclusion regions is to take the inverse
            mask = ~geom.region_mask([region])

        Note how we made a list with a single region,
        since this method expects a list of regions.
        """
        from gammapy.maps import Map, RegionGeom

        if not self.is_regular:
            raise ValueError("Multi-resolution maps not supported yet")

        geom = RegionGeom.from_regions(regions, wcs=self.wcs)
        idx = self.get_idx()
        mask = geom.contains_wcs_pix(idx)

        if not inside:
            np.logical_not(mask, out=mask)

        return Map.from_geom(self, data=mask)

    def region_weights(self, regions, oversampling_factor=10):
        """Compute regions weights

        Parameters
        ----------
        regions : str, `~regions.Region` or list of `~regions.Region`
            Region or list of regions (pixel or sky regions accepted).
            A region can be defined as a string ind DS9 format as well.
            See http://ds9.si.edu/doc/ref/region.html for details.
        oversampling_factor : int
            Over-sampling factor to compute the region weights

        Returns
        -------
        map : `~gammapy.maps.WcsNDMap` of boolean type
            Weights region mask
        """
        geom = self.upsample(factor=oversampling_factor)
        m = geom.region_mask(regions=regions)
        m.data = m.data.astype(float)
        return m.downsample(factor=oversampling_factor, preserve_counts=False)

    def binary_structure(self, width, kernel="disk"):
        """Get binary structure

        Parameters
        ----------
        width : `~astropy.units.Quantity`, str or float
            If a float is given it interpreted as width in pixels. If an (angular)
            quantity is given it converted to pixels using ``geom.wcs.wcs.cdelt``.
            The width corresponds to radius in case of a disk kernel, and
            the side length in case of a box kernel.
        kernel : {'disk', 'box'}
            Kernel shape

        Returns
        -------
        structure : `~numoy.ndarray`
            Binary structure
        """
        width = u.Quantity(width)

        if width.unit.is_equivalent("deg"):
            width = width / self.pixel_scales

        width = round_up_to_odd(width.to_value(""))

        if kernel == "disk":
            disk = Tophat2DKernel(width[0])
            disk.normalize("peak")
            structure = disk.array
        elif kernel == "box":
            structure = np.ones(width)
        else:
            raise ValueError(f"Invalid kernel: {kernel!r}")

        shape = (1,) * len(self.axes) + structure.shape
        return structure.reshape(shape)

    def __repr__(self):
        lon = self.center_skydir.data.lon.deg
        lat = self.center_skydir.data.lat.deg
        lon_ref, lat_ref = self.wcs.wcs.crval

        return (
            f"{self.__class__.__name__}\n\n"
            f"\taxes       : {self.axes_names}\n"
            f"\tshape      : {self.data_shape[::-1]}\n"
            f"\tndim       : {self.ndim}\n"
            f"\tframe      : {self.frame}\n"
            f"\tprojection : {self.projection}\n"
            f"\tcenter     : {lon:.1f} deg, {lat:.1f} deg\n"
            f"\twidth      : {self.width[0][0]:.1f} x {self.width[1][0]:.1f}\n"
            f"\twcs ref    : {lon_ref:.1f} deg, {lat_ref:.1f} deg\n"
        )

    def to_odd_npix(self, max_radius=None):
        """Create a new geom object with an odd number of pixel and a maximum size.

        This is useful for PSF kernel creation.

        Parameters
        ----------
        max_radius : `~astropy.units.Quantity`
            Max. radius of the geometry (half the width)

        Returns
        -------
        geom : `WcsGeom`
            Geom with odd number of pixels
        """
        if max_radius is None:
            width = self.width.max()
        else:
            width = 2 * u.Quantity(max_radius)

        binsz = self.pixel_scales.max()

        width_npix = (width / binsz).to_value("")
        npix = round_up_to_odd(width_npix)
        return WcsGeom.create(
            skydir=self.center_skydir,
            binsz=binsz,
            npix=npix,
            proj=self.projection,
            frame=self.frame,
            axes=self.axes,
        )

    def to_even_npix(self):
        """Create a new geom object with an even number of pixel and a maximum size.

        Returns
        -------
        geom : `WcsGeom`
            Geom with odd number of pixels
        """
        width = self.width.max()
        binsz = self.pixel_scales.max()

        width_npix = (width / binsz).to_value("")
        npix = round_up_to_even(width_npix)
        return WcsGeom.create(
            skydir=self.center_skydir,
            binsz=binsz,
            npix=npix,
            proj=self.projection,
            frame=self.frame,
            axes=self.axes,
        )

    def is_aligned(self, other, tolerance=1e-6):
        """Check if WCS and extra axes are aligned.

        Parameters
        ----------
        other : `WcsGeom`
            Other geom.
        tolerance : float
            Tolerance for the comparison.

        Returns
        -------
        aligned : bool
            Whether geometries are aligned
        """
        for axis, otheraxis in zip(self.axes, other.axes):
            if axis != otheraxis:
                return False

        # check WCS consistency with a priori tolerance of 1e-6
        return self.wcs.wcs.compare(other.wcs.wcs, cmp=2, tolerance=tolerance)

    def __eq__(self, other):
        if not isinstance(other, self.__class__):
            return NotImplemented

        if not (self.is_regular and other.is_regular):
            raise NotImplementedError(
                "Geom comparison is not possible for irregular geometries."
            )

        # check overall shape and axes compatibility
        if self.data_shape != other.data_shape:
            return False

        for axis, otheraxis in zip(self.axes, other.axes):
            if axis != otheraxis:
                return False

        # check WCS consistency with a priori tolerance of 1e-6
        # cmp=1 parameter ensures no comparison with ancillary information
        # see https://github.com/astropy/astropy/pull/4522/files
        return self.wcs.wcs.compare(other.wcs.wcs, cmp=1, tolerance=1e-6)

    def __ne__(self, other):
        return not self.__eq__(other)

    def __hash__(self):
        return id(self)

def pix2world(wcs, cdelt, crpix, pix):
    """Perform pixel to world coordinate transformation.

    For a WCS projection with a given pixel size (CDELT) and reference pixel
    (CRPIX). This method can be used to perform WCS transformations
    for projections with different pixelizations but the same
    reference coordinate (CRVAL), projection type, and coordinate system.

    Parameters
    ----------
    wcs : `astropy.wcs.WCS`
        WCS transform object.
    cdelt : tuple
        Tuple of X/Y pixel size in deg.  Each element should have the
        same length as ``pix``.
    crpix : tuple
        Tuple of reference pixel parameters in X and Y dimensions.  Each
        element should have the same length as ``pix``.
    pix : tuple
        Tuple of pixel coordinates.
    """
    pix_ratio = [
        np.abs(wcs.wcs.cdelt[0] / cdelt[0]),
        np.abs(wcs.wcs.cdelt[1] / cdelt[1]),
    ]
    pix = (
        (pix[0] - (crpix[0] - 1.0)) / pix_ratio[0] + wcs.wcs.crpix[0] - 1.0,
        (pix[1] - (crpix[1] - 1.0)) / pix_ratio[1] + wcs.wcs.crpix[1] - 1.0,
    )
    return wcs.wcs_pix2world(pix[0], pix[1], 0)

def world2pix(wcs, cdelt, crpix, coord):
    pix_ratio = [
        np.abs(wcs.wcs.cdelt[0] / cdelt[0]),
        np.abs(wcs.wcs.cdelt[1] / cdelt[1]),
    ]
    pix = wcs.wcs_world2pix(coord[0], coord[1], 0)
    return (
        (pix[0] - (wcs.wcs.crpix[0] - 1.0)) * pix_ratio[0] + crpix[0] - 1.0,
        (pix[1] - (wcs.wcs.crpix[1] - 1.0)) * pix_ratio[1] + crpix[1] - 1.0,
    )

reason: import numpy,import astropy,from astropy


text:
from sentence_transformers import LoggingHandler
from sentence_transformers.evaluation import SentenceEvaluator
import torch
import time
import logging
from tqdm import tqdm
from sentence_transformers.util import pytorch_cos_sim
import os
import csv
import numpy as np
from typing import List, Tuple, Dict, Set

logging.basicConfig(format='%(asctime)s - %(message)s',
                    datefmt='%Y-%m-%d %H:%M:%S',
                    level=logging.INFO,
                    handlers=[LoggingHandler()])

class Evaluator(SentenceEvaluator):
    """
    Base class for all evaluators

    Extend this class and implement __call__ for custom evaluators.
    """

    def __call__(self, model, output_path: str = None, epoch: int = -1, steps: int = -1) -> float:
        """
        This is called during training to evaluate the model.
        It returns a score for the evaluation with a higher score indicating a better result.

        :param model:
            the model to evaluate
        :param output_path:
            path where predictions and metrics are written to
        :param epoch
            the epoch where the evaluation takes place.
            This is used for the file prefixes.
            If this is -1, then we assume evaluation on test data.
        :param steps
            the steps in the current epoch at time of the evaluation.
            This is used for the file prefixes.
            If this is -1, then we assume evaluation at the end of the epoch.
        :return: a score for the evaluation with a higher score indicating a better result
        """
        pass

class ImageTextRetrievalEvaluator(Evaluator):
    """
    This class is the base class for Image-Text Retrieval.
    It measures Mean Reciprocal Rank (MRR), Recall@k, mean Recall (mR) [Average over all R@k for image and text retrieval]
    and Normalized Discounted Cumulative Gain (NDCG)
    """

    def __init__(self,
                 images: List[str],  #imageids
                 captions: Dict[str, str],  #cid => caption
                 imageid2captions: Dict[str, Set[str]],  #imageid => Set[cid]
                 tags = None,
                 mrr_at_k: List[int] = [10],
                 ndcg_at_k: List[int] = [10],
                 recall_at_k: List[int] = [1, 5, 10],
                 map_at_k: List[int] = [10],
                 mR_ks = [1, 5, 10],
                 show_progress_bar: bool = False,
                 batch_size: int = 32,
                 name: str = ''):

        self.images = images
        self.tags = tags
        self.captions_ids = list(captions.keys())
        self.captions = [captions[cid] for cid in self.captions_ids]
        self.imageid2captions = imageid2captions
        for imageid in list(imageid2captions.keys()):
            for cid in imageid2captions[imageid]:
                self.imageid2captions[cid] = {imageid}

        self.mrr_at_k = mrr_at_k
        self.ndcg_at_k = ndcg_at_k
        self.recall_at_k = recall_at_k
        self.map_at_k = map_at_k
        self.mR_ks = mR_ks

        self.show_progress_bar = show_progress_bar
        self.batch_size = batch_size
        self.name = name

        if name:
            name = "_" + name

        self.i2t_csv_file: str = "i2t_retrieval_evaluation" + name + "_results.csv"
        self.t2i_csv_file: str = "t2i_retrieval_evaluation" + name + "_results.csv"
        self.csv_headers = ["epoch", "steps"]

        for k in recall_at_k:
            self.csv_headers.append("Recall@{}".format(k))

        for k in mrr_at_k:
            self.csv_headers.append("MRR@{}".format(k))

        for k in ndcg_at_k:
            self.csv_headers.append("NDCG@{}".format(k))

        for k in map_at_k:
            self.csv_headers.append("MAP@{}".format(k))
        self.csv_headers.append("mR")

    def __call__(self, model, output_path: str = None, epoch: int = -1, steps: int = -1):
        raise NotImplementedError("")

    def _eval_i2t(self, i2t_result_list):
        for query_itr in range(len(i2t_result_list)):
            query_id = self.images[query_itr]

            #Sort scores
            if query_id in self.imageid2captions:
                top_hits = sorted(i2t_result_list[query_itr], key=lambda x: x['score'], reverse=True)
                query_imageid2captions = self.imageid2captions[query_id]
            else:
                continue
            if len(query_imageid2captions) == 0:
                continue

            #Recall@k - We count the result correct, if at least one relevant doc is accross the top-k documents
            for k_val in self.recall_at_k:
                self.i2t_num_hits_at_k[k_val].append(0)
                for hit in top_hits[0:k_val]:
                    if hit['captions_id'] in query_imageid2captions:
                        self.i2t_num_hits_at_k[k_val][-1] = 1
                        break

            #MRR@k
            for k_val in self.mrr_at_k:
                self.i2t_MRR[k_val].append(0)
                for rank, hit in enumerate(top_hits[0:k_val]):
                    if hit['captions_id'] in query_imageid2captions:
                        self.i2t_MRR[k_val][-1] = 1.0 / (rank + 1)
                        break

            #NDCG@k
            for k_val in self.ndcg_at_k:
                predicted_relevance = [1 if top_hit['captions_id'] in query_imageid2captions else 0 for top_hit in top_hits[0:k_val]]
                true_relevances = [1] * len(query_imageid2captions)

                ndcg_value = self.compute_dcg_at_k(predicted_relevance, k_val) / self.compute_dcg_at_k(true_relevances, k_val)
                self.i2t_ndcg[k_val].append(ndcg_value)

            #MAP@k
            for k_val in self.map_at_k:
                num_correct = 0
                sum_precisions = 0

                for rank, hit in enumerate(top_hits[0:k_val]):
                    if hit['captions_id'] in query_imageid2captions:
                        num_correct += 1
                        sum_precisions += num_correct / (rank+1)

                avg_precision = sum_precisions / min(k_val, len(query_imageid2captions))
                self.i2t_AveP_at_k[k_val].append(avg_precision)

    def _eval_t2i(self, t2i_result_list):
        for query_itr in range(len(t2i_result_list)):
            query_id = self.captions_ids[query_itr]

            #Sort scores
            if query_id in self.imageid2captions:
                top_hits = sorted(t2i_result_list[query_itr], key=lambda x: x['score'], reverse=True)
                query_imageid2captions = self.imageid2captions[query_id]
            else:
                continue
            if len(query_imageid2captions) == 0:
                continue

            #Recall@k - We count the result correct, if at least one relevant doc is accross the top-k documents
            for k_val in self.recall_at_k:
                self.t2i_num_hits_at_k[k_val].append(0)
                for hit in top_hits[0:k_val]:
                    if hit['image_id'] in query_imageid2captions:
                        self.t2i_num_hits_at_k[k_val][-1] = 1
                        break

            #MRR@k
            for k_val in self.mrr_at_k:
                self.t2i_MRR[k_val].append(0)
                for rank, hit in enumerate(top_hits[0:k_val]):
                    if hit['image_id'] in query_imageid2captions:
                        self.t2i_MRR[k_val][-1] = 1.0 / (rank + 1)
                        break

            #NDCG@k
            for k_val in self.ndcg_at_k:
                predicted_relevance = [1 if top_hit['image_id'] in query_imageid2captions else 0 for top_hit in top_hits[0:k_val]]
                true_relevances = [1] * len(query_imageid2captions)

                ndcg_value = self.compute_dcg_at_k(predicted_relevance, k_val) / self.compute_dcg_at_k(true_relevances, k_val)
                self.t2i_ndcg[k_val].append(ndcg_value)

            #MAP@k
            for k_val in self.map_at_k:
                num_correct = 0
                sum_precisions = 0

                for rank, hit in enumerate(top_hits[0:k_val]):
                    if hit['image_id'] in query_imageid2captions:
                        num_correct += 1
                        sum_precisions += num_correct / (rank+1)

                avg_precision = sum_precisions / min(k_val, len(query_imageid2captions))
                self.t2i_AveP_at_k[k_val].append(avg_precision)

    def _average_results(self):
            #Compute averages
            for k in self.i2t_num_hits_at_k:
                self.i2t_num_hits_at_k[k] = np.mean(self.i2t_num_hits_at_k[k])

            for k in self.i2t_ndcg:
                self.i2t_ndcg[k] = np.mean(self.i2t_ndcg[k])

            for k in self.i2t_MRR:
                self.i2t_MRR[k] = np.mean(self.i2t_MRR[k])

            for k in self.i2t_AveP_at_k:
                self.i2t_AveP_at_k[k] = np.mean(self.i2t_AveP_at_k[k])

            for k in self.t2i_num_hits_at_k:
                self.t2i_num_hits_at_k[k] = np.mean(self.t2i_num_hits_at_k[k])

            for k in self.t2i_ndcg:
                self.t2i_ndcg[k] = np.mean(self.t2i_ndcg[k])

            for k in self.t2i_MRR:
                self.t2i_MRR[k] = np.mean(self.t2i_MRR[k])

            for k in self.t2i_AveP_at_k:
                self.t2i_AveP_at_k[k] = np.mean(self.t2i_AveP_at_k[k])

            self.mR = (sum(self.t2i_num_hits_at_k[k] for k in self.mR_ks) + sum(self.i2t_num_hits_at_k[k] for k in self.mR_ks)) / (2*len(self.mR_ks))

    def _log_results(self, delta_time):
            # Output
            logging.info("IMAGE TO TEXT")
            for k in self.i2t_num_hits_at_k:
                logging.info("Recall@{}: {:.2f}%".format(k, self.i2t_num_hits_at_k[k]*100))

            for k in self.i2t_MRR:
                logging.info("MRR@{}: {:.4f}".format(k, self.i2t_MRR[k]))

            for k in self.i2t_ndcg:
                logging.info("NDCG@{}: {:.4f}".format(k, self.i2t_ndcg[k]))

            for k in self.i2t_AveP_at_k:
                logging.info("MAP@{}: {:.4f}".format(k, self.i2t_AveP_at_k[k]))

            logging.info("TEXT TO IMAGE")
            for k in self.t2i_num_hits_at_k:
                logging.info("Recall@{}: {:.2f}%".format(k, self.t2i_num_hits_at_k[k]*100))

            for k in self.t2i_MRR:
                logging.info("MRR@{}: {:.4f}".format(k, self.t2i_MRR[k]))

            for k in self.t2i_ndcg:
                logging.info("NDCG@{}: {:.4f}".format(k, self.t2i_ndcg[k]))

            for k in self.t2i_AveP_at_k:
                logging.info("MAP@{}: {:.4f}".format(k, self.t2i_AveP_at_k[k]))

            logging.info("mR: {:.4f}".format(self.mR))

            logging.info("Images: {}".format(len(self.images)))
            logging.info("Queries: {}\n".format(len(self.captions)))
            logging.info(f"Time: {delta_time:.2f}")

    def _write_output(self, output_path, epoch, steps):
        if output_path is not None:
            csv_path = os.path.join(output_path, self.i2t_csv_file)
            if not os.path.isfile(csv_path):
                fOut = open(csv_path, mode="w", encoding="utf-8")
                fOut.write(",".join(self.csv_headers))
                fOut.write("\n")

            else:
                fOut = open(csv_path, mode="a", encoding="utf-8")

            output_data = [epoch, steps]
            for k in self.recall_at_k:
                output_data.append(self.i2t_num_hits_at_k[k])

            for k in self.mrr_at_k:
                output_data.append(self.i2t_MRR[k])

            for k in self.ndcg_at_k:
                output_data.append(self.i2t_ndcg[k])

            for k in self.map_at_k:
                output_data.append(self.i2t_AveP_at_k[k])

            output_data.append(self.mR)

            fOut.write(",".join(map(str,output_data)))
            fOut.write("\n")
            fOut.close()

            csv_path = os.path.join(output_path, self.t2i_csv_file)
            if not os.path.isfile(csv_path):
                fOut = open(csv_path, mode="w", encoding="utf-8")
                fOut.write(",".join(self.csv_headers))
                fOut.write("\n")

            else:
                fOut = open(csv_path, mode="a", encoding="utf-8")

            output_data = [epoch, steps]
            for k in self.recall_at_k:
                output_data.append(self.t2i_num_hits_at_k[k])

            for k in self.mrr_at_k:
                output_data.append(self.t2i_MRR[k])

            for k in self.ndcg_at_k:
                output_data.append(self.t2i_ndcg[k])

            for k in self.map_at_k:
                output_data.append(self.t2i_AveP_at_k[k])

            output_data.append(self.mR)

            fOut.write(",".join(map(str,output_data)))
            fOut.write("\n")
            fOut.close()

    @staticmethod
    def compute_dcg_at_k(relevances, k):
        dcg = 0
        for i in range(min(len(relevances), k)):
            dcg += relevances[i] / np.log2(i + 2)  #+2 as we start our idx at 0
        return dcg

class EmbeddingImageTextRetrievalEvaluator(ImageTextRetrievalEvaluator):
    """
    Image-Text Retrieval with Embeddings
    It measures Mean Reciprocal Rank (MRR), Recall@k, mean Recall (mR) [Average over all R@k for image and text retrieval]
    and Normalized Discounted Cumulative Gain (NDCG)
    :param batched_sim: If set to a value >0, the cosine similarity will be computed with this as batch size.
    Useful for very large datasets.
    """

    def __init__(self,
                 images: List[str],  #imageids
                 captions: Dict[str, str],  #cid => caption
                 imageid2captions: Dict[str, Set[str]],  #imageid => Set[cid]
                 tags = None,
                 mrr_at_k: List[int] = [10],
                 ndcg_at_k: List[int] = [10],
                 recall_at_k: List[int] = [1, 3, 5, 10],
                 map_at_k: List[int] = [10],
                 mR_ks = [1, 5, 10],
                 show_progress_bar: bool = False,
                 batch_size: int = 32,
                 name: str = '',
                 batched_sim=0):

        super(EmbeddingImageTextRetrievalEvaluator, self).__init__(images, captions, imageid2captions, tags, mrr_at_k,
                                                                   ndcg_at_k, recall_at_k, map_at_k, mR_ks, show_progress_bar, batch_size, name)
        self.batched_sim = batched_sim

    def __call__(self, model, output_path: str = None, epoch: int = -1, steps: int = -1):
        if epoch != -1:
            out_txt = " after epoch {}:".format(epoch) if steps == -1 else " in epoch {} after {} steps:".format(epoch, steps)
        else:
            out_txt = ":"

        logging.info("Image Text Retrieval Evaluation on " + self.name + " dataset" + out_txt)

        self.i2t_num_hits_at_k = {k: [] for k in self.recall_at_k}
        self.i2t_MRR = {k: [] for k in self.mrr_at_k}
        self.i2t_ndcg = {k: [] for k in self.ndcg_at_k}
        self.i2t_AveP_at_k = {k: [] for k in self.map_at_k}

        self.t2i_num_hits_at_k = {k: [] for k in self.recall_at_k}
        self.t2i_MRR = {k: [] for k in self.mrr_at_k}
        self.t2i_ndcg = {k: [] for k in self.ndcg_at_k}
        self.t2i_AveP_at_k = {k: [] for k in self.map_at_k}

        self.mR = 0

        max_k = max(max(self.mrr_at_k), max(self.ndcg_at_k), max(self.recall_at_k), max(self.map_at_k))

        time_start = time.time()

        image_embeddings = model.encode(sentences=self.tags, images=self.images, show_progress_bar=self.show_progress_bar, batch_size=self.batch_size, convert_to_tensor=True)
        caption_embeddings = model.encode(sentences=self.captions, show_progress_bar=self.show_progress_bar, batch_size=self.batch_size, convert_to_tensor=True)

        delta_time = time.time() - time_start

        i2t_result_list = [[] for _ in range(0, len(image_embeddings))]
        t2i_result_list = [[] for _ in range(0, len(caption_embeddings))]

        time_start = time.time()

        if self.batched_sim > 0:
            for start_idx in tqdm(range(0, len(image_embeddings), self.batched_sim), desc="Similarity"):
                cos_scores = pytorch_cos_sim(image_embeddings[start_idx:start_idx+self.batched_sim], caption_embeddings)

                # Image to Text
                cos_scores_top_k_values, cos_scores_top_k_idx = torch.topk(cos_scores, min(max_k, len(cos_scores[0]) - 1), dim=1, largest=True, sorted=False)
                i2t_cos_scores_top_k_values = cos_scores_top_k_values.tolist()
                i2t_cos_scores_top_k_idx = cos_scores_top_k_idx.tolist()

                for query_itr in range(len(i2t_cos_scores_top_k_values)):
                    for captions_id, score in zip(i2t_cos_scores_top_k_idx[query_itr], i2t_cos_scores_top_k_values[query_itr]):
                        captions_id = self.captions_ids[captions_id]
                        i2t_result_list[start_idx+query_itr].append({'captions_id': captions_id, 'score': score})
        else:
            #Compute cosine similarites
            cos_scores = pytorch_cos_sim(image_embeddings, caption_embeddings)

            # Image to Text
            cos_scores_top_k_values, cos_scores_top_k_idx = torch.topk(cos_scores, min(max_k, len(cos_scores[0]) - 1), dim=1, largest=True, sorted=False)
            i2t_cos_scores_top_k_values = cos_scores_top_k_values.tolist()
            i2t_cos_scores_top_k_idx = cos_scores_top_k_idx.tolist()

            for query_itr in range(len(i2t_cos_scores_top_k_values)):
                for captions_id, score in zip(i2t_cos_scores_top_k_idx[query_itr], i2t_cos_scores_top_k_values[query_itr]):
                    captions_id = self.captions_ids[captions_id]
                    i2t_result_list[query_itr].append({'captions_id': captions_id, 'score': score})

        delta_time = time.time() - time_start + delta_time

        self._eval_i2t(i2t_result_list)

        time_start = time.time()

        if self.batched_sim > 0:
            for start_idx in tqdm(range(0, len(caption_embeddings), self.batched_sim), desc="Similarity"):
                cos_scores = pytorch_cos_sim(image_embeddings, caption_embeddings[start_idx:start_idx+self.batched_sim])

                cos_scores_top_k_values, cos_scores_top_k_idx = torch.topk(cos_scores, min(max_k, len(cos_scores[1]) - 1), dim=0, largest=True, sorted=False)
                t2i_cos_scores_top_k_values = cos_scores_top_k_values.permute(1, 0).tolist()
                t2i_cos_scores_top_k_idx = cos_scores_top_k_idx.permute(1, 0).tolist()

                for query_itr in range(len(t2i_cos_scores_top_k_values)):
                    for image_id, score in zip(t2i_cos_scores_top_k_idx[query_itr], t2i_cos_scores_top_k_values[query_itr]):
                        image_id = self.images[image_id]
                        t2i_result_list[start_idx+query_itr].append({'image_id': image_id, 'score': score})
        else:
            cos_scores_top_k_values, cos_scores_top_k_idx = torch.topk(cos_scores, min(max_k, len(cos_scores[1]) - 1), dim=0, largest=True, sorted=False)
            t2i_cos_scores_top_k_values = cos_scores_top_k_values.permute(1, 0).tolist()
            t2i_cos_scores_top_k_idx = cos_scores_top_k_idx.permute(1, 0).tolist()

            for query_itr in range(len(t2i_cos_scores_top_k_values)):
                for image_id, score in zip(t2i_cos_scores_top_k_idx[query_itr], t2i_cos_scores_top_k_values[query_itr]):
                    image_id = self.images[image_id]
                    t2i_result_list[query_itr].append({'image_id': image_id, 'score': score})

        delta_time = time.time() - time_start + delta_time

        self._eval_t2i(t2i_result_list)

        self._average_results()
        self._log_results(delta_time)
        self._write_output(output_path, epoch, steps)

        return self.mR

class CrossEncoderImageTextRetrievalEvaluator(ImageTextRetrievalEvaluator):
    """
    Image-Text Retrieval with Cross Encoder
    It measures Mean Reciprocal Rank (MRR), Recall@k, mean Recall (mR) [Average over all R@k for image and text retrieval]
    and Normalized Discounted Cumulative Gain (NDCG)
    """

    def __init__(self,
                 images: List[str],  #imageids
                 captions: Dict[str, str],  #cid => caption
                 imageid2captions: Dict[str, Set[str]],  #imageid => Set[cid]
                 tags = None,
                 mrr_at_k: List[int] = [10],
                 ndcg_at_k: List[int] = [10],
                 recall_at_k: List[int] = [1, 3, 5, 10],
                 map_at_k: List[int] = [10],
                 mR_ks = [1, 5, 10],
                 show_progress_bar: bool = False,
                 batch_size: int = 32,
                 name: str = ''):

        super(CrossEncoderImageTextRetrievalEvaluator, self).__init__(images, captions, imageid2captions, tags, mrr_at_k,
                                                                   ndcg_at_k, recall_at_k, map_at_k, mR_ks, show_progress_bar, batch_size, name)

    def __call__(self, model, output_path: str = None, epoch: int = -1, steps: int = -1) -> float:
        if epoch != -1:
            out_txt = " after epoch {}:".format(epoch) if steps == -1 else " in epoch {} after {} steps:".format(epoch, steps)
        else:
            out_txt = ":"

        logging.info("Image Text Retrieval Evaluation on " + self.name + " dataset" + out_txt)

        max_k = max(max(self.mrr_at_k), max(self.ndcg_at_k), max(self.recall_at_k), max(self.map_at_k))

        #Init score computation values
        self.i2t_num_hits_at_k = {k: [] for k in self.recall_at_k}
        self.i2t_MRR = {k: [] for k in self.mrr_at_k}
        self.i2t_ndcg = {k: [] for k in self.ndcg_at_k}
        self.i2t_AveP_at_k = {k: [] for k in self.map_at_k}

        self.t2i_num_hits_at_k = {k: [] for k in self.recall_at_k}
        self.t2i_MRR = {k: [] for k in self.mrr_at_k}
        self.t2i_ndcg = {k: [] for k in self.ndcg_at_k}
        self.t2i_AveP_at_k = {k: [] for k in self.map_at_k}

        self.mR = 0

        i2t_result_list = [[] for _ in range(0, len(self.images))]
        t2i_result_list = [[] for _ in range(0, len(self.captions))]

        time_start = time.time()

        #Compute cosine similarites
        cos_scores = torch.zeros((len(self.images), len(self.captions)))
        for i in range(len(self.images)):
            images = []
            captions = []
            for c in range(len(self.captions)):
                images.append(self.images[i])
                if self.tags:
                    captions.append([self.captions[c], self.tags[i]])
                else:
                    captions.append(self.captions[c])
            sims = model.encode(captions, images, batch_size=self.batch_size, output_value="logits", show_progress_bar=True, convert_to_tensor=True)[:,1]
            cos_scores[i] = sims

        # Image to Text
        cos_scores_top_k_values, cos_scores_top_k_idx = torch.topk(cos_scores, min(max_k, len(self.captions) - 1), dim=1, largest=True, sorted=False)
        i2t_cos_scores_top_k_values = cos_scores_top_k_values.tolist()
        i2t_cos_scores_top_k_idx = cos_scores_top_k_idx.tolist()

        delta_time = time.time() - time_start

        for query_itr in range(len(i2t_cos_scores_top_k_values)):
            for captions_id, score in zip(i2t_cos_scores_top_k_idx[query_itr], i2t_cos_scores_top_k_values[query_itr]):
                captions_id = self.captions_ids[captions_id]
                i2t_result_list[query_itr].append({'captions_id': captions_id, 'score': score})

        self._eval_i2t(i2t_result_list)

        time_start = time.time()

        cos_scores_top_k_values, cos_scores_top_k_idx = torch.topk(cos_scores, min(max_k, len(self.images) - 1), dim=0, largest=True, sorted=False)
        t2i_cos_scores_top_k_values = cos_scores_top_k_values.permute(1, 0).tolist()
        t2i_cos_scores_top_k_idx = cos_scores_top_k_idx.permute(1, 0).tolist()

        delta_time = time.time() - time_start + delta_time

        for query_itr in range(len(t2i_cos_scores_top_k_values)):
            for image_id, score in zip(t2i_cos_scores_top_k_idx[query_itr], t2i_cos_scores_top_k_values[query_itr]):
                image_id = self.images[image_id]
                t2i_result_list[query_itr].append({'image_id': image_id, 'score': score})

        self._eval_t2i(t2i_result_list)
        self._average_results()
        self._log_results(delta_time)
        self._write_output(output_path, epoch, steps)

        return self.mR

class RetrieveRerankImageTextRetrievalEvaluator(ImageTextRetrievalEvaluator):
    """
    Image-Text Retrieval with Retrieve & Re-Ranking
    It measures Mean Reciprocal Rank (MRR), Recall@k, mean Recall (mR) [Average over all R@k for image and text retrieval]
    and Normalized Discounted Cumulative Gain (NDCG)
    :param retrieve: Decide how many candidates will be retrieved for re-ranking
    :param scoring: Use another re-ranking approach besides purely using the cross-encoder scores
    :param scoring_factor: weight factor between the embedding score and cross encoder score (if scoring is not 'standard')
    """

    def __init__(self,
                 images: List[str],  #imageids
                 captions: Dict[str, str],  #cid => caption
                 imageid2captions: Dict[str, Set[str]],  #imageid => Set[cid]
                 tags = None,
                 mrr_at_k: List[int] = [10],
                 ndcg_at_k: List[int] = [10],
                 recall_at_k: List[int] = [1, 3, 5, 10],
                 map_at_k: List[int] = [10],
                 mR_ks = [1, 5, 10],
                 show_progress_bar: bool = False,
                 batch_size: int = 32,
                 name: str = '',
                 retrieve=10,
                 scoring="standard",
                 scoring_factor=0.5):
        super(RetrieveRerankImageTextRetrievalEvaluator, self).__init__(images, captions, imageid2captions, tags, mrr_at_k,
         ndcg_at_k, recall_at_k, map_at_k, mR_ks, show_progress_bar, batch_size, name)
        self.retrieve = retrieve
        self.scoring = scoring
        self.scoring_factor = scoring_factor

    def __call__(self, models, output_path: str = None, epoch: int = -1, steps: int = -1) -> float:
        if epoch != -1:
            out_txt = " after epoch {}:".format(epoch) if steps == -1 else " in epoch {} after {} steps:".format(epoch, steps)
        else:
            out_txt = ":"

        if isinstance(models, list):
            biencoder, crossencoder = models
        else:
            biencoder = models
            crossencoder = models
        logging.info("Image Text Retrieval Evaluation on " + self.name + " dataset" + out_txt)

        max_k = max(self.retrieve, max(self.mrr_at_k), max(self.ndcg_at_k), max(self.recall_at_k), max(self.map_at_k))

        time_start = time.time()

        # Compute embedding for the images
        image_embeddings = biencoder.encode(sentences=self.tags, images=self.images, show_progress_bar=self.show_progress_bar, batch_size=self.batch_size, convert_to_tensor=True)
        #Compute embedding for the captions
        caption_embeddings = biencoder.encode(sentences=self.captions, show_progress_bar=self.show_progress_bar, batch_size=self.batch_size, convert_to_tensor=True)

        #Init score computation values
        self.i2t_num_hits_at_k = {k: [] for k in self.recall_at_k}
        self.i2t_MRR = {k: [] for k in self.mrr_at_k}
        self.i2t_ndcg = {k: [] for k in self.ndcg_at_k}
        self.i2t_AveP_at_k = {k: [] for k in self.map_at_k}

        self.t2i_num_hits_at_k = {k: [] for k in self.recall_at_k}
        self.t2i_MRR = {k: [] for k in self.mrr_at_k}
        self.t2i_ndcg = {k: [] for k in self.ndcg_at_k}
        self.t2i_AveP_at_k = {k: [] for k in self.map_at_k}

        self.mR = 0

        i2t_result_list = [[] for _ in range(0, len(image_embeddings))]
        t2i_result_list = [[] for _ in range(0, len(caption_embeddings))]

        #Compute cosine similarites
        cos_scores = pytorch_cos_sim(image_embeddings, caption_embeddings).cpu()

        cos_scores_top_k_values, cos_scores_top_k_idx = torch.topk(cos_scores, max_k, dim=1, largest=True, sorted=False)
        i2t_cos_scores_top_k_values = cos_scores_top_k_values
        i2t_cos_scores_top_k_idx = cos_scores_top_k_idx

        images = []
        captions = []
        assign = []
        for i in range(len(self.images)):
            if self.images[i] in self.imageid2captions:
                for c, c_idx in enumerate(i2t_cos_scores_top_k_idx[i]):
                    images.append(self.images[i])
                    if self.tags:
                        captions.append([self.captions[c_idx], self.tags[i]])
                    else:
                        captions.append(self.captions[c_idx])
                    assign.append((i,c))

        #mb_sims = i2t_cos_scores_top_k_values.detach().clone()
        if self.scoring == "combination":
            sims = torch.softmax(crossencoder.encode(captions, images, batch_size=self.batch_size, output_value="logits", show_progress_bar=True, convert_to_tensor=True), 1)[:,1]
            for s, (i,c) in zip(sims, assign):
                emb_score = i2t_cos_scores_top_k_values[i][c]
                i2t_cos_scores_top_k_values[i][c] = (1-self.scoring_factor)*s + self.scoring_factor*emb_score
        elif self.scoring == "combination_normalized":
            sims = torch.softmax(crossencoder.encode(captions, images, batch_size=self.batch_size, output_value="logits", show_progress_bar=True, convert_to_tensor=True), 1)[:,1]
            sims = sims.reshape(len(self.images), -1)
            for i in range(len(self.images)):
                emb_min_abs = torch.abs(torch.min(i2t_cos_scores_top_k_values[i]))
                emb_score = (i2t_cos_scores_top_k_values[i]-emb_min_abs)/(torch.max(i2t_cos_scores_top_k_values[i])-emb_min_abs)
                i2t_cos_scores_top_k_values[i] = (1-self.scoring_factor)*(sims[i]-torch.min(sims[i]))/(torch.max(sims[i])-torch.min(sims[i])) + \
                                                 self.scoring_factor*emb_score
        elif self.scoring == "rank":
            sims = crossencoder.encode(captions, images, batch_size=self.batch_size, output_value="logits", show_progress_bar=True, convert_to_tensor=True)[:,1]
            sims = sims.reshape(len(self.images), -1)
            for i in range(len(self.images)):
                emb_rank = torch.argsort(i2t_cos_scores_top_k_values[i]).float()
                sim_rank = torch.argsort(sims[i]).float()
                i2t_cos_scores_top_k_values[i] = (1-self.scoring_factor)*emb_rank + self.scoring_factor*sim_rank
        else:
            sims = crossencoder.encode(captions, images, batch_size=self.batch_size, output_value="logits", show_progress_bar=True, convert_to_tensor=True)[:,1]
            for s, (i,c) in zip(sims, assign):
                i2t_cos_scores_top_k_values[i][c] = s

        delta_time = time.time() - time_start

        for query_itr in range(len(i2t_cos_scores_top_k_values)):
            for captions_id, score in zip(i2t_cos_scores_top_k_idx[query_itr], i2t_cos_scores_top_k_values[query_itr]):
                captions_id = self.captions_ids[captions_id]
                i2t_result_list[query_itr].append({'captions_id': captions_id, 'score': score})

        self._eval_i2t(i2t_result_list)

        time_start = time.time()

        cos_scores_top_k_values, cos_scores_top_k_idx = torch.topk(cos_scores, max_k, dim=0, largest=True, sorted=False)
        t2i_cos_scores_top_k_values = cos_scores_top_k_values.permute(1, 0)
        t2i_cos_scores_top_k_idx = cos_scores_top_k_idx.permute(1, 0)

        images = []
        captions = []
        assign = []
        for c in range(len(self.captions)):
            if self.captions_ids[c] in self.imageid2captions:
                for i, i_idx in enumerate(t2i_cos_scores_top_k_idx[c]):
                    images.append(self.images[i_idx])
                    if self.tags:
                        captions.append([self.captions[c], self.tags[i_idx]])
                    else:
                        captions.append(self.captions[c])
                    assign.append((c,i))
        if self.scoring == "combination":
            sims = torch.softmax(crossencoder.encode(captions, images, batch_size=self.batch_size, output_value="logits", show_progress_bar=True, convert_to_tensor=True), 1)[:,1]
            for s, (i,c) in zip(sims, assign):
                emb_score = t2i_cos_scores_top_k_values[i][c]
                t2i_cos_scores_top_k_values[i][c] = (1-self.scoring_factor)*s + self.scoring_factor*emb_score
        elif self.scoring == "combination_normalized":
            sims = torch.softmax(crossencoder.encode(captions, images, batch_size=self.batch_size, output_value="logits", show_progress_bar=True, convert_to_tensor=True), 1)[:,1]
            sims = sims.reshape(len(self.captions), -1)
            for i in range(len(self.images)):
                emb_min_abs = torch.abs(torch.min(t2i_cos_scores_top_k_values[i]))
                emb_score = (t2i_cos_scores_top_k_values[i]-emb_min_abs)/(torch.max(t2i_cos_scores_top_k_values[i])-emb_min_abs)
                sims_score = (sims[i]-torch.min(sims[i]))/(torch.max(sims[i])-torch.min(sims[i]))
                t2i_cos_scores_top_k_values[i] = (1-self.scoring_factor)*sims_score + self.scoring_factor*emb_score
        elif self.scoring == "rank":
            sims = crossencoder.encode(captions, images, batch_size=self.batch_size, output_value="logits", show_progress_bar=True, convert_to_tensor=True)[:,1]
            sims = sims.reshape(len(self.captions), -1)
            for i in range(len(self.images)):
                emb_rank = torch.argsort(t2i_cos_scores_top_k_values[i]).float()
                sim_rank = torch.argsort(sims[i]).float()
                t2i_cos_scores_top_k_values[i] = (1-self.scoring_factor)*emb_rank + self.scoring_factor*sim_rank
        else:
            sims = crossencoder.encode(captions, images, batch_size=self.batch_size, output_value="logits", show_progress_bar=True, convert_to_tensor=True)[:,1]
            for s, (i,c) in zip(sims, assign):
                t2i_cos_scores_top_k_values[i][c] = s

        delta_time = time.time() - time_start + delta_time

        for query_itr in range(len(t2i_cos_scores_top_k_values)):
            for image_id, score in zip(t2i_cos_scores_top_k_idx[query_itr], t2i_cos_scores_top_k_values[query_itr]):
                image_id = self.images[image_id]
                t2i_result_list[query_itr].append({'image_id': image_id, 'score': score})

        self._eval_t2i(t2i_result_list)

        self._average_results()
        self._log_results(delta_time)
        self._write_output(output_path, epoch, steps)

        return self.mR

class ImageQuestionClassification(Evaluator):
    """
    Evaluation for image question datasets like VQA, GQA
    """
    def __init__(self,
                 images,  #imageids
                 questions,  #cid => caption
                 labels,
                 batch_size: int = 32,
                 name: str = ''):

        self.batch_size = batch_size
        self.name = name
        self.questions = questions
        self.labels = labels
        self.images = images
        if name:
            name = "_" + name

        self.file: str = "class" + name + "_results.csv"
        self.csv_headers = ["epoch", "steps", "acc"]

    def __call__(self, model, output_path: str = None, epoch: int = -1, steps: int = -1) -> float:
        if epoch != -1:
            out_txt = " after epoch {}:".format(epoch) if steps == -1 else " in epoch {} after {} steps:".format(epoch, steps)
        else:
            out_txt = ":"

        logging.info("Accuracy on " + self.name + " dataset" + out_txt)

        logits = model.encode(self.questions, self.images, batch_size=self.batch_size, show_progress_bar=True, output_value="logits")

        acc = (self.labels.squeeze() == torch.argmax(logits, dim=1)).squeeze().long().cpu().tolist()

        acc = np.mean(acc)
        logging.info("Questions: {}\n".format(len(self.questions)))
        logging.info(f"Accuracy: {acc:.2f}")
        if output_path is not None:
            csv_path = os.path.join(output_path, self.file)
            if not os.path.isfile(csv_path):
                fOut = open(csv_path, mode="w", encoding="utf-8")
                fOut.write(",".join(self.csv_headers))
                fOut.write("\n")
            else:
                fOut = open(csv_path, mode="a", encoding="utf-8")

            output_data = [epoch, steps, acc]

            fOut.write(",".join(map(str,output_data)))
            fOut.write("\n")
            fOut.close()
        return acc

ext: py

max_stars_repo_name: UKPLab/MMT-Retrieval

max_issues_repo_name: UKPLab/MMT-Retrieval

max_forks_repo_name: UKPLab/MMT-Retrieval

reason: import numpy


text:
#This is Talkbot neural network project
#this project copied form Neuraline code and developed on top of that
#First run this training program to train all chats which system can learn when you run
import json
import random
import pickle
import numpy as np

import nltk
#nltk.download('punkt')
#nltk.download('wordnet')
from nltk.stem import WordNetLemmatizer

from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Activation, Dropout
from tensorflow.keras.optimizers import SGD

lemmatizer = WordNetLemmatizer()

intents = json.loads(open('intents.json').read())

words = []
labels = []
documents = []
ignore_letters = ['?', '!', '.', ',']

for intent in intents['intents']:
    for pattern in intent['patterns']:
        word_list = nltk.word_tokenize(pattern)
        #print(word_list)
        words.extend(word_list)
        #print(words)
        documents.append((word_list, intent['tag']))
        if intent['tag'] not in labels:
            labels.append(intent['tag'])
            #print(labels)
#print(training_data)

words = [lemmatizer.lemmatize(w.lower()) for w in words if w not in ignore_letters] #edited
words = sorted(list(set(words))) #edited
#print(words)
labels = sorted(list(set(labels))) #edited

#print(labels)

pickle.dump(words, open('talkbot_words.pkl', 'wb'))
pickle.dump(labels, open('talkbot_labels.pkl', 'wb'))

training_set = []
output_empty = [0] * len(labels)
#print(output_empty)

for doc in documents:
    bag = []
    word_patterns = doc[0]
    #print(word_patterns)
    word_patterns = [lemmatizer.lemmatize(word.lower()) for word in word_patterns]
    #print(word_patterns)
    for word in words:
        bag.append(1) if word in word_patterns else bag.append(0)

    output_row = list(output_empty)
    output_row[labels.index(doc[1])] = 1
    training_set.append([bag, output_row])

random.shuffle(training_set)
training_set = np.array(training_set)

train_x = list(training_set[:, 0])
train_y = list(training_set[:, 1])

model = Sequential()
model.add(Dense(128, input_shape=(len(train_x[0]),), activation='relu'))
model.add(Dropout(0.5))
model.add(Dense(64, activation='relu'))
model.add(Dropout(0.5))
model.add(Dense(len(train_y[0]), activation='softmax'))

sgd = SGD(lr=0.01, decay=1e-6, momentum=0.9, nesterov=True)
model.compile(loss='categorical_crossentropy', optimizer=sgd, metrics=['accuracy'])

hist = model.fit(np.array(train_x), np.array(train_y), epochs=500, batch_size=5, verbose=1)
model.save('talkbotnlp_model.h5', hist)
print("Done")


text:
# This code is from
# Pareto Multi-Task LearningLin
# Xi Lin, Hui-Ling Zhen, Zhenhua Li, Qingfu Zhang, Sam Kwong
# Neural Information Processing Systems (NeurIPS) 2019
# https://github.com/Xi-L/ParetoMTL

import numpy as np
import torch

from .base import Solver
from .utils import flatten_grad, recover_flattened, flatten_param_grad

from time import time
from datetime import timedelta

@torch.no_grad()
def get_transferences(model, batch, list_of_gradients, initial_task_loss, lr):
    # NOTE: we assume task-specific gradients have been applied.
    initial_params = [v.clone() for k, v in model.model.get_shared_parameters().items()]

    transferences = []
    for gradients in list_of_gradients:
        # take one SGD step
        shared_params = [v for k, v in model.model.get_shared_parameters().items()]
        for index, param in enumerate(shared_params):
            param.data = initial_params[index] - lr * gradients[index]

        # compute transference
        task_loss = model(batch[0], batch[1])
        transference = torch.sum(1 - task_loss / initial_task_loss)
        transferences.append(transference)

    # reset original parameters
    shared_params = [v for k, v in model.model.get_shared_parameters().items()]
    for index, param in enumerate(shared_params):
        param.data = initial_params[index].data

    return torch.stack(transferences)

def get_d_pcgrad(gradients):
    final_grad = 0.0
    for grad_index, grad in enumerate(gradients):
        indices = np.arange(len(gradients))
        indices = np.concatenate([indices[:grad_index], indices[grad_index + 1 :]])
        np.random.shuffle(indices)
        for index in indices:
            other_grad = gradients[index]
            cos_sim = torch.clamp(torch.dot(grad, other_grad), max=0)
            grad = grad - ((cos_sim / torch.linalg.norm(other_grad)) * other_grad)
        final_grad = final_grad + grad
    return final_grad

class ITMTL(Solver):
    @property
    def name(self):
        return "itmtl"

    def update_fn(self, X, ts, model, optimizer):
        # clear the graph
        model.zero_grad()
        # compute shared gradients
        flat_grads = {}
        shared_params = [v for k, v in model.model.get_shared_parameters().items()]

        task_loss = model(X, ts)
        for i in range(self.dataset_config.n_tasks):
            shared_grads = torch.autograd.grad(
                task_loss[i], shared_params, retain_graph=True
            )
            flat_grads[i] = flatten_param_grad(shared_grads)

        # update task parameters
        for i in range(self.dataset_config.n_tasks):
            task_params = [v for k, v in model.model.get_task_parameters(i).items()]
            task_grads = torch.autograd.grad(
                task_loss[i],
                task_params,
                retain_graph=True,
            )
            for index, params in enumerate(task_params):
                params.data = params.data - self.flags.lr * task_grads[index]
                # add a gradient for metrics computation
                params.grad = task_grads[index]

        # compute PCGrad
        pcgrads = [flat_grads[i]["grad"] for i in range(len(flat_grads))]
        pcgrads = get_d_pcgrad(torch.stack(pcgrads))
        pcgrads = recover_flattened(
            pcgrads, flat_grads[0]["indices"], flat_grads[0]["shapes"]
        )

        # compute original gradients
        oggrads = [flat_grads[i]["grad"] for i in range(len(flat_grads))]
        oggrads = torch.mean(torch.stack(oggrads), dim=0)
        oggrads = recover_flattened(
            oggrads, flat_grads[0]["indices"], flat_grads[0]["shapes"]
        )

        # compute transference
        gradient_candidates = [pcgrads, oggrads]
        transferences = get_transferences(
            model, (X, ts), gradient_candidates, task_loss, self.flags.lr
        )
        gradients = gradient_candidates[torch.argmax(transferences).item()]

        # update shared parameters
        shared_params = [v for k, v in model.model.get_shared_parameters().items()]
        for index, params in enumerate(shared_params):
            params.data = params.data - self.flags.lr * gradients[index]
            # add a gradient for metrics computation
            params.grad = gradients[index]

    def run(self):
        """Run itmtl-pcgrad"""
        print(f"**** Now running {self.name} on {self.dataset} ... ")
        start_time = time()
        results = dict()
        for i in range(self.flags.n_preferences):
            self.suffix = f"p{i}"
            model = self.configure_model()
            result, checkpoint = self.train(model, None)
            results[i] = dict(r=None, res=result, checkpoint=checkpoint)
            self.dump(
                results, self.prefix + f"_{self.flags.n_preferences}_from_0-{i}.pkl"
            )

        total_time = timedelta(seconds=round(time() - start_time))
        print(f"**** Time taken for {self.name} on {self.dataset} = {total_time}s.")


text:
from __future__ import division

from shapely.wkt import loads as wkt_loads

import os
import shapely
import shapely.geometry
import shapely.affinity
import h5py
import pandas as pd
import tifffile as tiff
from numba import jit, njit
from tqdm import tqdm
from collections import defaultdict
import csv
import sys

import cv2
from shapely.geometry import MultiPolygon, Polygon
import shapely.wkt
import shapely.affinity
import numpy as np

import skimage.color as color
from skimage.transform import rescale

# dirty hacks from SO to allow loading of big cvs's
# without decrement loop it crashes with C error
# http://stackoverflow.com/questions/15063936/csv-error-field-larger-than-field-limit-131072
maxInt = sys.maxsize
decrement = True

while decrement:
    # decrease the maxInt value by factor 10
    # as long as the OverflowError occurs.

    decrement = False
    try:
        csv.field_size_limit(maxInt)
    except OverflowError:
        maxInt = int(maxInt/10)
        decrement = True

data_path = '../data'
train_wkt = pd.read_csv(os.path.join(data_path, 'train_wkt_v4.csv'))
gs = pd.read_csv(os.path.join(data_path, 'grid_sizes.csv'), names=['ImageId', 'Xmax', 'Ymin'], skiprows=1)
shapes = pd.read_csv(os.path.join(data_path, '3_shapes.csv'))

epsilon = 1e-15

def get_scalers(height, width, x_max, y_min):
    """

    :param height:
    :param width:
    :param x_max:
    :param y_min:
    :return: (xscaler, yscaler)
    """
    w_ = width * (width / (width + 1))
    h_ = height * (height / (height + 1))
    return w_ / x_max, h_ / y_min

def polygons2mask_layer(height, width, polygons, image_id):
    """

    :param height:
    :param width:
    :param polygons:
    :return:
    """

    x_max, y_min = _get_xmax_ymin(image_id)
    x_scaler, y_scaler = get_scalers(height, width, x_max, y_min)

    polygons = shapely.affinity.scale(polygons, xfact=x_scaler, yfact=y_scaler, origin=(0, 0, 0))
    img_mask = np.zeros((height, width), np.uint8)

    if not polygons:
        return img_mask

    int_coords = lambda x: np.array(x).round().astype(np.int32)
    exteriors = [int_coords(poly.exterior.coords) for poly in polygons]
    interiors = [int_coords(pi.coords) for poly in polygons for pi in poly.interiors]

    cv2.fillPoly(img_mask, exteriors, 1)
    cv2.fillPoly(img_mask, interiors, 0)
    return img_mask

def polygons2mask(height, width, polygons, image_id):
    num_channels = len(polygons)
    result = np.zeros((num_channels, height, width))
    for mask_channel in range(num_channels):
        result[mask_channel, :, :] = polygons2mask_layer(height, width, polygons[mask_channel], image_id)
    return result

def generate_mask(image_id, height, width, num_mask_channels=10, train=train_wkt):
    """

    :param image_id:
    :param height:
    :param width:
    :param num_mask_channels: numbers of channels in the desired mask
    :param train: polygons with labels in the polygon format
    :return: mask corresponding to an image_id of the desired height and width with desired number of channels
    """

    mask = np.zeros((num_mask_channels, height, width))

    for mask_channel in range(num_mask_channels):
        poly = train.loc[(train['ImageId'] == image_id)
                         & (train['ClassType'] == mask_channel + 1), 'MultipolygonWKT'].values[0]
        polygons = shapely.wkt.loads(poly)
        mask[mask_channel, :, :] = polygons2mask_layer(height, width, polygons, image_id)
    return mask

def mask2polygons_layer(mask, epsilon=1.0, min_area=10.0):
    # first, find contours with cv2: it's much faster than shapely
    contours, hierarchy = cv2.findContours(((mask == 1) * 255).astype(np.uint8), cv2.RETR_CCOMP, cv2.CHAIN_APPROX_TC89_KCOS)

    # create approximate contours to have reasonable submission size
    if epsilon != 0:
        approx_contours = simplify_contours(contours, epsilon)
    else:
        approx_contours = contours

    if not approx_contours:
        return MultiPolygon()

    all_polygons = find_child_parent(hierarchy, approx_contours, min_area)

    # approximating polygons might have created invalid ones, fix them
    all_polygons = MultiPolygon(all_polygons)

    all_polygons = fix_invalid_polygons(all_polygons)

    return all_polygons

def find_child_parent(hierarchy, approx_contours, min_area):
    # now messy stuff to associate parent and child contours
    cnt_children = defaultdict(list)
    child_contours = set()
    assert hierarchy.shape[0] == 1

    # http://docs.opencv.org/3.1.0/d9/d8b/tutorial_py_contours_hierarchy.html
    for idx, (_, _, _, parent_idx) in enumerate(hierarchy[0]):
        if parent_idx != -1:
            child_contours.add(idx)
            cnt_children[parent_idx].append(approx_contours[idx])

    # create actual polygons filtering by area (removes artifacts)
    all_polygons = []
    for idx, cnt in enumerate(approx_contours):
        if idx not in child_contours and cv2.contourArea(cnt) >= min_area:
            assert cnt.shape[1] == 1
            holes = [c[:, 0, :] for c in cnt_children.get(idx, []) if cv2.contourArea(c) >= min_area]
            contour = cnt[:, 0, :]

            poly = Polygon(shell=contour, holes=holes)

            if poly.area >= min_area:
                all_polygons.append(poly)

    return all_polygons

def simplify_contours(contours, epsilon):
    return [cv2.approxPolyDP(cnt, epsilon, True) for cnt in contours]

def fix_invalid_polygons(all_polygons):
    if not all_polygons.is_valid:
        all_polygons = all_polygons.buffer(0)
        # Sometimes buffer() converts a simple Multipolygon to just a Polygon,
        # need to keep it a Multi throughout
        if all_polygons.type == 'Polygon':
            all_polygons = MultiPolygon([all_polygons])
    return all_polygons

def _get_xmax_ymin(image_id):
    xmax, ymin = gs[gs['ImageId'] == image_id].iloc[0, 1:].astype(float)
    return xmax, ymin

def get_shape(image_id, band=3):
    if band == 3:
        height = shapes.loc[shapes['image_id'] == image_id, 'height'].values[0]
        width = shapes.loc[shapes['image_id'] == image_id, 'width'].values[0]
        return height, width

def read_image_16(image_id):
    img_m = np.transpose(tiff.imread("../data/sixteen_band/{}_M.tif".format(image_id)), (1, 2, 0)) / 2047.0
    img_3 = np.transpose(tiff.imread("../data/three_band/{}.tif".format(image_id)), (1, 2, 0)) / 2047.0
    img_p = tiff.imread("../data/sixteen_band/{}_P.tif".format(image_id)).astype(np.float32) / 2047.0

    height, width, _ = img_3.shape

    rescaled_M = cv2.resize(img_m, (width, height), interpolation=cv2.INTER_CUBIC)
    rescaled_P = cv2.resize(img_p, (width, height), interpolation=cv2.INTER_CUBIC)

    rescaled_M[rescaled_M > 1] = 1
    rescaled_M[rescaled_M < 0] = 0

    rescaled_P[rescaled_P > 1] = 1
    rescaled_P[rescaled_P < 0] = 0

    image_r = img_3[:, :, 0]
    image_g = img_3[:, :, 1]
    image_b = img_3[:, :, 2]
    nir = rescaled_M[:, :, 7]
    re = rescaled_M[:, :, 5]

    L = 1.0
    C1 = 6.0
    C2 = 7.5
    evi = (nir - image_r) / (nir + C1 * image_r - C2 * image_b + L)
    evi = np.expand_dims(evi, 2)

    ndwi = (image_g - nir) / (image_g + nir)
    ndwi = np.expand_dims(ndwi, 2)

    savi = (nir - image_r) / (image_r + nir)
    savi = np.expand_dims(savi, 2)

    ccci = (nir - re) / (nir + re) * (nir - image_r) / (nir + image_r)
    ccci = np.expand_dims(ccci, 2)

    rescaled_P = np.expand_dims(rescaled_P, 2)

    result = np.transpose(np.concatenate([rescaled_M, rescaled_P, ndwi, savi, evi, ccci, img_3], axis=2), (2, 0, 1))
    return result.astype(np.float16)

def make_prediction_cropped(model, X_train, initial_size=(572, 572), final_size=(388, 388), num_channels=19, num_masks=10):
    shift = int((initial_size[0] - final_size[0]) / 2)

    height = X_train.shape[1]
    width = X_train.shape[2]

    if height % final_size[1] == 0:
        num_h_tiles = int(height / final_size[1])
    else:
        num_h_tiles = int(height / final_size[1]) + 1

    if width % final_size[1] == 0:
        num_w_tiles = int(width / final_size[1])
    else:
        num_w_tiles = int(width / final_size[1]) + 1

    rounded_height = num_h_tiles * final_size[0]
    rounded_width = num_w_tiles * final_size[0]

    padded_height = rounded_height + 2 * shift
    padded_width = rounded_width + 2 * shift

    padded = np.zeros((num_channels, padded_height, padded_width))

    padded[:, shift:shift + height, shift: shift + width] = X_train

    # add mirror reflections to the padded areas
    up = padded[:, shift:2 * shift, shift:-shift][:, ::-1]
    padded[:, :shift, shift:-shift] = up

    lag = padded.shape[1] - height - shift
    bottom = padded[:, height + shift - lag:shift + height, shift:-shift][:, ::-1]
    padded[:, height + shift:, shift:-shift] = bottom

    left = padded[:, :, shift:2 * shift][:, :, ::-1]
    padded[:, :, :shift] = left

    lag = padded.shape[2] - width - shift
    right = padded[:, :, width + shift - lag:shift + width][:, :, ::-1]

    padded[:, :, width + shift:] = right

    h_start = range(0, padded_height, final_size[0])[:-1]
    assert len(h_start) == num_h_tiles

    w_start = range(0, padded_width, final_size[0])[:-1]
    assert len(w_start) == num_w_tiles

    temp = []
    for h in h_start:
        for w in w_start:
            temp += [padded[:, h:h + initial_size[0], w:w + initial_size[0]]]

    prediction = model.predict(np.array(temp))

    predicted_mask = np.zeros((num_masks, rounded_height, rounded_width))

    for j_h, h in enumerate(h_start):
         for j_w, w in enumerate(w_start):
             i = len(w_start) * j_h + j_w
             predicted_mask[:, h: h + final_size[0], w: w + final_size[0]] = prediction[i]

    return predicted_mask[:, :height, :width]

max_stars_repo_name: yokoyang/img-class

max_issues_repo_name: yokoyang/img-class

max_forks_repo_name: yokoyang/img-class

reason: import numpy,from numba