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

#
#
# # Earth Analytics Education - EA  Python Course Spring 2021

# ## Important  - Assignment Guidelines
#
# 1. Before you submit your assignment to GitHub, make sure to run the entire notebook with a fresh kernel. To do this first, **restart the kernel** (in the menubar, select Kernel$\rightarrow$Restart & Run All)
# 2. Always replace the `raise NotImplementedError()` code with your code that addresses the activity challenge. If you don't replace that code, your notebook will not run.
#
# ```
# # YOUR CODE HERE
# raise NotImplementedError()
# ```
#
# 3. Any open ended questions will have a "YOUR ANSWER HERE" within a markdown cell. Replace that text with your answer also formatted using Markdown.
# 4. **DO NOT RENAME THIS NOTEBOOK File!** If the file name changes, the autograder will not grade your assignment properly.
# 6. When you create a figure, comment out `plt.show()` to ensure the autograder can grade your plots. For figure cells, DO NOT DELETE the code that says `DO NOT REMOVE LINE BELOW`.
#
# ```
# ### DO NOT REMOVE LINE BELOW ###
# student_plot1_ax = nb.convert_axes(plt)
# ```
#
# * Only include the package imports, code, and outputs that are required to run your homework assignment.
# * Be sure that your code can be run on any operating system. This means that:
#    1. the data should be downloaded in the notebook to ensure it's reproducible
#    2. all paths should be created dynamically using the `os.path.join`
#
# ## Follow to PEP 8 Syntax Guidelines & Documentation
#
# * Run the `autopep8` tool on all cells prior to submitting (HINT: hit shift + the tool to run it on all cells at once!
# * Use clear and expressive names for variables.
# * Organize your code to support readability.
# * Check for code line length
# * Use comments and white space sparingly where it is needed
# * Make sure all python imports are at the top of your notebook and follow PEP 8 order conventions
# * Spell check your Notebook before submitting it.
#
# For all of the plots below, be sure to do the following:
#
# * Make sure each plot has a clear TITLE and, where appropriate, label the x and y axes. Be sure to include UNITS in your labels.
#

# ### Add Your Name Below
# **Your Name:** Adam Mansur

#

# ---

# # Week 04 and 05 Homework - Automate NDVI Workflow
#
# For this assignment, you will write code to generate a plot of the mean normalized difference vegetation index (NDVI) for two different sites in the United States across one year of data:
#
# * San Joaquin Experimental Range (SJER) in Southern California, United States
# * Harvard Forest (HARV) in the Northeastern United States
#
# The data that you will use for this week is available from **earthpy** using the following download:
#
# `et.data.get_data('ndvi-automation')`
#
# ## Assignment Goals
#
# Your goal in this assignment is to create the most efficient and concise workflow that you can that allows for:
#
# 1. The code to scale if you added new sites or more time periods to the analysis.
# 2. Someone else to understand your workflow.
# 3. The LEAST and most efficient (i.e. runs fast, minimize repetition) amount of code that completes the task.
#
# ### HINTS
#
# * Remove values outside of the landsat valid range of values as specified in the metadata, as needed.
# * Keep any output files SEPARATE FROM input files. Outputs should be created in an outputs directory that is created in the code (if needed) and/or tested for.
# * Use the functions that we demonstrated during class to make your workflow more efficient.
# * BONUS - if you  chose - you can export your data as a csv file. You will get bonus points for doing this.
#
#
# ## Assignment Requirements
#
# Your submission to the GitHub repository should include:
# * This Jupyter Notebook file (.ipynb) with:
#     * The code to create a plot of mean NDVI across a year for  2 NEON Field Sites:
#         * NDVI on the x axis and formatted dates on the y for both NEON sites on one figure/axis object
#     * The **data should be cleaned to remove the influence of clouds**. See the [earthdatascience website for an example of what your plot might look like with and without removal of clouds](https://www.earthdatascience.org/courses/earth-analytics-python/create-efficient-data-workflows/).
# * BONUS: Create one output `.csv` file that has 3 columns - NDVI, Date and Site Name - with values for SJER and HARV.
#
# Your notebook should:
# * Have *at least* 2 well documented and well named functions with docstrings.
# * Include a Markdown cell at the top of the notebook that outlines the overall workflow using pseudocode (i.e. plain language, not code)
# * Include additional Markdown cells throughout the notebook to describe:
#     * the data that you used - and where it is from
#     * how data are being processing
#     * how the code is optimized to run fast and be more concise

# # Replace this cell with your pseudocode  for this workflow
#
# If you happen to be a diagram person a diagram is ok too
#
#

# # Pseudocode
#
# *The cell above is locked*
#
# + Identify and iterate through 2017 Landstat tilesets for HARV and
#   SJER sites. For each site:
#     + Read the site boundary vector file
#     + For each scene:
#         + Read and label red, NIR, and QA bands
#         + Remove extreme values
#         + Crop bands to site boundary
#         + Remove clouds and cloud shadows
#         + Calculate NDVI and mean NDVI
#         + Parse acquisition date from Landsat file names
#         + Create dataframe with site name, date, and mean NDVI
# + Concatenate individual dataframes into a single dataframe
# + Plot NDVI over time for both sites
#
#
# # Approach
#
# The functions for reading, cropping, and masking Landsat data were combined
# into an xarray accessor class called LandsatAccessor. The accessor structure is
# the recommended way to extend xarray objects. I chose this approach because:
#
# 1. I've found it difficult to keep track of what's going on when cleaning and
#    plotting data using xarray. Many of the procedures require remembering to use
#    functions I don't really understand (like plotting_extent) or process bands
#    using list indices that could vary depending on exactly how the data was read.
#    This makes it hard to re-use code or remember how to make various plots. The
#    accessor shifts some of the confusing bits (like deriving and applying the
#    cloud QA mask) behind the scenes.
# 2. The functions I came up with are more readable/reproducible if the Landsat
#    data is read in a specific way (for example, with the band attribute set to
#    match the original band number instead of the read order). That is a good
#    case for bundling methods in a class.
#
# Some other things I like about this approach include that it:
#
# + Allows access to individual bands by name (red, green, blue, nir)
# + Clarifies calculations like NDVI by explicitly naming the bands involved
# + Compiles Landsat-related logic inside a single class (plus a helper function
#   used to read the Landsat images) that are available from any xarray object
#
# Drawbacks to this approach include:
#
# + The ad hoc attributes required to use the accessor methods are not always
#   carried over by array operations, so the accessor is somewhat brittle. This
#   is a limitation of xarray, which prioritizes letting calcualtions run over
#   preserving metadata.
# + The accessor currently only works on Landsat data. Currently Landsat-specific
#   logic is used for mapping bands to colors, reading metadata, and masking
#   using the QA layer. The first thing could be generalized, but not the other
#   two seem to depend heavily on the exact satellite producing the data.
# + Length of code is excessive for the current assignment
#
# The code isn't optimized besides really basic stuff like using built-in array
# functions. The loop used to look for all HARV and SJER data should scale.

# In[1]:

# Autograding imports - do not modify this cell
import matplotcheck.autograde as ag
import matplotcheck.notebook as nb
import matplotcheck.timeseries as ts
from datetime import datetime

# In[2]:

# Import needed packages in PEP 8 order
# and no unused imports listed (10 points total)
from glob import glob
import os
import re

import earthpy as et
import earthpy.mask as em
import earthpy.plot as ep
import geopandas as gpd
import matplotlib.dates as mdates
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
import rioxarray as rxr
import xarray as xr

# Increase default font size on all plots
sns.set(font_scale=1.5)

# Download NDVI automation dataset
et.data.get_data("ndvi-automation")

# Set working directory
os.chdir(os.path.join(et.io.HOME, "earth-analytics", "data"))

# In[3]:

# DO NOT MODIFY THIS CELL
# Tests that the working directory is set to earth-analytics/data

path = os.path.normpath(os.getcwd())
student_wd_parts = path.split(os.sep)

if student_wd_parts[-2:] == ['earth-analytics', 'data']:
    print("\u2705 Great - it looks like your working directory is set correctly to ~/earth-analytics/data")
else:
    print("\u274C Oops, the autograder will not run unless your working directory is set to earth-analytics/data")

# # Figure 1: Plot 1 - Mean NDVI For Each Site Across the Year (50 points)
#
# Create a plot of the mean normalized difference vegetation index (NDVI) for the two different sites in the United States across the year:
#
# * NDVI on the x axis and formatted dates on the y for both NEON sites on one figure/axis object.
# * Each site should be identified with a different color in the plot and legend.
# * The final plot **data should be cleaned to remove the influence of clouds**.
# * Be sure to include appropriate title and axes labels.
#
# Add additional cells as needed for processing data (e.g. defining functions, etc), but be sure to:
# * follow the instructions in the code cells that have been provided to ensure that you are able to use the sanity check tests that are provided.
# * include only the plot code in the cell identified for the final plot code below

# ## Task 1:
#
# In the cell below, create a single dataframe containing MEAN NDVI, the site name,
# and the date of the data for the HARV site
# scene `HARV/landsat-crop/LC080130302017031701T1-SC20181023151837`.  The column names for the  final
# DataFrame should be`mean_ndvi`, and `site`, and the data should be **indexed on the date**.
#
# Use the functions that we reviewed in class (or create your own versions of them) to implement your code
#
# ### In the Cell below Place  All Functions Needed to Run this Notebook (20 points)

# In[4]:

### DO NOT REMOVE THIS LINE OR EDIT / MOVE THIS CELL ###
start_time = datetime.now()

# In[5]:

# In this cell place all of the functions needed to run your notebook
# You will be graded here on function application, docstrings, efficiency so ensure
# All functions are placed here!

@xr.register_dataarray_accessor("landsat")
class LandsatAccessor:
    """Extends data array to parse and run common calculations on Landsat data

    Attributes
    ----------
    red : xarray DataArray
        red band
    green : xarray DataArray
        green band
    blue : xarray DataArray
        blue band
    nir : xarray DataArray
        near-infrared band
    qa : xarray DataArray
        pixel qa band
    metadata : dict
        Landsat metadata parsed from filename
    """

    def __init__(self, xarr):
        self._obj = xarr
        self._metadata = None
        self._names = {
            2: "blue",
            3: "green",
            4: "red",
            5: "nir",
        }

    @property
    def blue(self):
        return self._obj.sel(band=2)

    @property
    def green(self):
        return self._obj.sel(band=3)

    @property
    def red(self):
        return self._obj.sel(band=4)

    @property
    def nir(self):
        return self._obj.sel(band=5)

    @property
    def qa(self):
        return self._obj.sel(band="qa")

    @property
    def metadata(self):
        """Returns Landsat metadata, deriving it from a filename if needed

        Metadata parsing is deferred so that non-Landsat arrays don't get
        tripped up by the parse function when they are created.
        """
        if self._metadata is None:
            self._metadata = self._parse_product_id()
        return self._metadata

    def clean_all_bands(self, **kwargs):
        """Cleans all bands in the array

        Arguments
        ---------
        kwargs:
            any keyword argument accepted by clean_band()

        Returns
        -------
        xarray DataArray
            Array with out-of-range data set to nan
        """
        # Clean each band and build a new DataArray
        out_xr = []
        for arr in self._obj:
            if arr.band != "qa":
                out_xr.append(self.clean_band(arr, **kwargs))
        out_xr.append(self.qa.copy())

        return xr.concat(out_xr, dim="band")

    def mask_qa_pixels(self, flags):
        """Masks array based on Landsat QA pixel flags

        Arguments
        ---------
        flags : list-like
            list of pixel QA flags to mask

        Returns
        -------
        xarray DataArray
            Array with specified features set to nan
        """

        # Format key for earthy.mask lookup from satellite number
        key = "L{}".format(self.metadata["satellite"].lstrip("0"))

        # Read pixel values for specified flags
        mask = []
        for flag in flags:
            mask.extend(em.pixel_flags["pixel_qa"][key][flag])

        # Mask each data band separately and build a new DataArray
        out_xr = []
        for arr in self._obj:
            if arr.band != "qa":
                out_xr.append(arr.where(~self.qa.isin(mask)))

                # Ad hoc attributes are lost here, so set them manually.
                out_xr[-1]["band"] = arr.band
                out_xr[-1]["path"] = arr.path

        out_xr.append(self.qa.copy())

        return xr.concat(out_xr, dim="band")

    def mask_clouds(self):
        """Masks clouds in an array (Landsat 8 only)

        Returns
        -------
        xarray DataArray
            Array with cloud features set to nan
        """
        return self.mask_qa_pixels(
            ["Cloud", "Cloud Shadow", "High Cloud Confidence"]
        )

    def ndvi(self):
        """Calculates NDVI using the named NIR and red bands

        Returns
        -------
        xarray DataArray
            Array containing NDVI values
        """
        return (self.nir - self.red) / (self.nir + self.red)

    def plot_bands(self, **kwargs):
        """Plots all bands in the array

        Arguments
        ---------
        kwargs:
            any keyword argument accepted by ep.plot_bands(). The title
            arguments is provided if not given.

        Returns
        -------
        None
        """
        # Set default titles for each band
        bands = [a.band.values[()] for a in self._obj]
        kwargs.setdefault("title", [self._names.get(b, b) for b in bands])

        ep.plot_bands(self._obj, **kwargs)

    @staticmethod
    def clean_band(band, valid_range=None):
        """Cleans a single Landsat band

        Parameters
        -----------
        band : str or array
            Either an array or the path to the array to be opened
        valid_range : tuple (optional)
            A tuple of min and max range of values for the data. Default = None

        Returns
        -----------
        xarray DataArray
            Array with out-of-range values set to nan
        """
        # Read band from file if passed as a string
        if isinstance(band, str):
            band = rxr.open_rasterio(band_path, masked=True).squeeze()

        # Remove pixels outside valid range if provided
        if valid_range:
            mask = ((band < min(valid_range)) | (band > max(valid_range)))
            band = band.where(~xr.where(mask, True, False))

        return band

    def _parse_product_id(self):
        """Parses acquisition and processing info from Landsat filename

        Returns
        -------
        dict
            Dictionary containing the parsed metadata
        """
        prod_id = os.path.basename(self._obj[0].path.values[()])
        sat, corr, wrs, acq, proc, coll_num, coll_cat = prod_id.split("_")[:7]

        return {
            "sensor": sat[1],
            "satellite": sat[3:],
            "wrs_path": wrs[:3],
            "wrs_row": wrs[3:],
            "acq_date": datetime.strptime(acq, "%Y%m%d"),
            "proc_date": datetime.strptime(proc, "%Y%m%d"),
            "coll_num": coll_num,
            "coll_cat": coll_cat
        }

def open_landsat(path, min_band=1, max_band=7):
    """Combines Landsat TIFFs in a directory into an xarray

    Also adds two ad hoc attributes, band and path. Some methods, like
    arr.where(), seem to lose track of these attributes in some cases.

    Arguments
    ---------
    path : str
        path to directory containing a set of Landsat TIFFs
    min_band: int
        minimum band to retrieve
    max_band: int
        maximum band to retrieve

    Returns
    -------
    xarray DataArray
         Array merging all matching TIFFs into one object

    Raises
    ------
    IOError
        If wrong number of bands found
    """

    # Read TIFFs along path matching selected bands, then add the QA band
    paths = []
    for fp in glob(os.path.join(path, f"*band[{min_band}-{max_band}]*.tif")):
        paths.append(fp)
    paths.sort()
    paths.extend(glob(os.path.join(path, f"*pixel_qa*.tif")))

    # Raise an error if any of the expected bands are missing
    expected = (max_band - min_band + 2)
    if len(paths) != expected:
        raise IOError(f"Wrong number of TIFFs found on {os.path.abspath(path)}"
                      f" (found {len(paths)}, expected={expected})")

    # Combine ordered list of TIFFs into an xarray
    out_xr =[]
    for fp in paths:
        out_xr.append(rxr.open_rasterio(fp, masked=True).squeeze())

        # Add min_band so that the band id matches original Landsat band
        out_xr[-1]["band"] = len(out_xr) + min_band - 1
        out_xr[-1]["path"] = os.path.realpath(fp)

    # Rename the last band to "qa"
    out_xr[-1]["band"] = "qa"

    return xr.concat(out_xr, dim="band")

def mean_ndvi_from_landsat(path, site_name, site_boundary=None):
    """Calculates mean NDVI for a set of Landsat images

    Arguments
    ---------
    path: str
        path to directory containing a set of Landsat images
    site_name : str
        name of the site imaged by Landsat
    site_boundary : str or geopandas.GeoDatafFrame
        boundary of the site as either a path to a shapefile or GeoDataFrame

    Returns
    -------
    pandas DataFrame
        Date-indexed dataframe with the mean NDVI from the Landsat data
    """
    arr = open_landsat(path, min_band=4, max_band=5)

    # Clip site to boundary
    if site_boundary is not None:

        # Read boundary from file if string given
        if isinstance(site_boundary, str):
            site_boundary = gpd.read_file(site_boundary)

        # Align CRS if not the same
        if arr.rio.crs != site_boundary.crs:
            site_boundary = site_boundary.to_crs(arr.rio.crs)

        arr = arr.rio.clip(site_boundary.geometry)

    # Remove extreme values from all bands
    arr = arr.landsat.clean_all_bands(valid_range=(0, 10000))

    # Mask clouds
    arr = arr.landsat.mask_clouds()

    # Return a data-indexed dataframe with the mean NDVI for the dataset
    return pd.DataFrame({
            "site": site_name,
            "mean_ndvi": float(arr.landsat.ndvi().mean())
        },
        index=[arr.landsat.metadata["acq_date"]]
    )

# ### Data sources
#
# Data for the following calculation of NDVI was collected by the Landsat 8
# satellite in March 2017 at the HARV field site as part of the NEON program.

# In[6]:

# Create dataframe of mean NDVI in this cell using the functions created above
# Important: to use the ungraded tests below as a sanity check,
# name your columns: mean_ndvi and site
# Call the dataframe at the end of the cell so the tests run on it!
# Be sure that the date column is an index of type date
# HINT: the time series lessons may help you remember how to do this!

# Define path to HARV folder
harv_path = os.path.join("ndvi-automation", "sites", "HARV")

# Define path to HARV site boundary
harv_bound_path = os.path.join(harv_path, "vector", "HARV-crop.shp")

# Define path to specific dataset
harv_20170317_path = os.path.join(
    harv_path, "landsat-crop", "LC080130302017031701T1-SC20181023151837"
)

# Calculate mean NDVI
harv_20170317_mean_ndvi = mean_ndvi_from_landsat(
    harv_20170317_path, "HARV", harv_bound_path
)
harv_20170317_mean_ndvi

# In[7]:

# This cell  is testing your data output above

student_ndvi_ts_single_site = _

single_scene_points = 0

# Ensure the data is stored in a dataframe.
if isinstance(student_ndvi_ts_single_site, pd.DataFrame):
    print('\u2705 Your data is stored in a DataFrame!')
    single_scene_points += 1
else:
    print('\u274C It appears your data is not stored in a DataFrame. ',
          'To see what type of object your data is stored in, check its type with type(object)')

# Ensure that the date column is the index
if isinstance(student_ndvi_ts_single_site.index, pd.core.indexes.datetimes.DatetimeIndex):
    print('\u2705 You have the index set to the date column!')
    single_scene_points += 2
else:
    print('\u274C You do not have the index set to the date column.')

# Ensure that the date column is datetime
if isinstance(student_ndvi_ts_single_site.index[0], pd._libs.tslibs.timestamps.Timestamp):
    print('\u2705 The data in your date column is datetime!')
    single_scene_points += 2
else:
    print('\u274C The data in your date column is not datetime.')

# Ensure the site name is correct
if student_ndvi_ts_single_site.site.values[0] == 'HARV':
    print('\u2705 You have the correct site name!')
    single_scene_points += 5
else:
    print('\u274C You do not have the correct site name.')

if np.allclose(0.281131628228094, student_ndvi_ts_single_site.mean_ndvi.values[0]):
    print('\u2705 You have the correct mean NDVI value!')
    single_scene_points += 5
else:
    print('\u274C You do not have the correct mean ndvi value.')

print("\n \u27A1 You received {} out of 15 points for creating a dataframe.".format(
    single_scene_points))
single_scene_points

# ## Task 2:
#
# In the cell below, process all of the landsat scenes. Create a DataFrame that contains the following
# information for each scene
#
#
# |   | index  | site  | mean_ndvi  |
# |---|---|---|---|
# | Date  |   |   |   |
# |  2017-01-07  | 0  | SJER  | .4  |
#
# Be sure to call your dataframe at the end of the cell to ensure autograding works.
# HINT: FOR THIS STEP, leave any rows containing missing values (`NAN`).

# ### Data sources
#
# Data for the following calculation of NDVI was collected by the Landsat 8
# satellite in 2017 for the HARV and SJER field sites as part of the NEON program.

# In[8]:

# Create dataframe of NDVI including the cleaning data to deal with clouds

# Important: to use the ungraded tests below as a sanity check,
# name your columns: mean_ndvi and site
# Don't forget to set date as the index and make the values of type datetime

# Define path to SJER folder (the path to HARV is defined above)
sjer_path = os.path.join("ndvi-automation", "sites", "SJER")

# Read boundaries for each site
boundaries = {
    "HARV": gpd.read_file(os.path.join(harv_path, "vector", "HARV-crop.shp")),
    "SJER": gpd.read_file(os.path.join(sjer_path, "vector", "SJER-crop.shp"))
}

# Calculate NDVI for each scene
ndvi_by_scene = []
for root, dirs, files in os.walk(os.path.join("ndvi-automation", "sites")):
    if os.path.basename(root).startswith("LC08"):
        # Grab site name from root and process scene
        site_name = re.search(r"\b[A-Z]{4}\b", root).group()
        ndvi_by_scene.append(
            mean_ndvi_from_landsat(root, site_name, boundaries[site_name])
        )

mean_ndvi_all = pd.concat(ndvi_by_scene).sort_index()
mean_ndvi_all

# In[9]:

# Last sanity check before creating your plot (10 points)

# Ensure that you call your dataframe at the bottom of the cell above
# and that it has columns called: mean_ndvi and site

# Ensure the data is stored in a dataframe.
student_ndvi_df = _

df_points = 0

if isinstance(student_ndvi_df, pd.DataFrame):
    print('\u2705 Your data is stored in a DataFrame!')
    df_points +=2
else:
    print('\u274C It appears your data is not stored in a DataFrame. ',
          'To see what type of object your data is stored in, check its type with type(object)')

# Check that dataframe contains the appropriate number of NAN values
if student_ndvi_df.mean_ndvi.isna().sum() == 15:
    print('\u2705 Correct number of masked data values!')
    df_points +=2
else:
    print('\u274C The amount of null data in your dataframe is incorrect.')

# Ensure that the date column is the index
if isinstance(student_ndvi_df.index, pd.core.indexes.datetimes.DatetimeIndex):
    print('\u2705 You have the index set to the date column!')
    df_points +=3
else:
    print('\u274C You do not have the index set to the date column.')

# Ensure that the date column is datetime
if isinstance(student_ndvi_df.index[0], pd._libs.tslibs.timestamps.Timestamp):
    print('\u2705 The data in your date column is datetime!')
    df_points +=3
else:
    print('\u274C The data in your date column is not datetime.')

# Output for timer, # DO NOT MODIFY
end_time = datetime.now()
total_time = end_time - start_time
print(
    "Your total run time for processing the data was {0}.".format(total_time))

print("\n \u27A1 You received {} out of 10 points for creating a dataframe.".format(
    df_points))

df_points

# In[10]:

# Add only the plot code to this cell

# This is the final figure of mean NDVI
# for both sites across the year
# with data cleaned to deal with clouds

# Remove null NDVIs so lines connect all points
mean_ndvi_sub = mean_ndvi_all.dropna(how="any", axis=0)

fig, ax = plt.subplots(1, 1, figsize=(12, 6))
ax.set(
    title="Mean NDVI at two NEON sites in 2017",
    xlabel="Date",
    ylabel="Mean NDVI",
    ylim=(0, 1)
)

# Plot scenes for each site as a line
for site, rows in mean_ndvi_sub.groupby("site"):
    ax.plot(rows.index, rows.mean_ndvi, marker="o", label=site)

# Use short month names to label the x axis
ax.xaxis.set_major_formatter(mdates.DateFormatter('%b\n%Y'))

ax.legend()

fig.text(0.5, -0.1, f"Source: NEON (Landsat 8)", fontsize=14, ha="center")

### DO NOT REMOVE LINES BELOW ###
final_masked_solution = nb.convert_axes(plt, which_axes="current")

# In[11]:

# Ignore this cell for the autograding tests

# In[12]:

# Ignore this cell for the autograding tests

# # Question 1 (10 points)
#
# Imagine that you are planning NEON’s upcoming flight season to capture remote sensing data in these locations and want to ensure that you fly the area when the vegetation is the most green.
#
# When would you recommend the flights take place for each site?
#
# Answer the question in 2-3 sentences in the Markdown cell below.

# Healthy, green vegetation will show a moderate to high NDVI. The highest NDVI (and
# therefore greenest vegeation) occurred at the SJER site between March and mid-April
# and at the HARV site between mid-May and late September, so those would be the best
# windows for scheduling flights.

# # Question 2 (10 points)
#
# How could you modify your workflow to look at vegetation changes over time in each site?
#
# Answer the question in 2-3 sentences in the Markdown cell below.

# One approach would be to compare results from periods of high and low NDVI. For SJER,
# that would mean scheduling a flight for sometime between July and December in addition
# to the high NDVI flight above. For HARV, the low NDVI flight would be scheduled for
# March to April (and possibily earlier since data from January and February are not
# available but NDVI at this site is trending down in December).

# # Do not edit this cell! (10 points)
#
# The notebook includes:
# * additional Markdown cells throughout the notebook to describe:
#     * the data that you used - and where it is from
#     * how data are being processing
#     * how the code is optimized to run fast and be more concise

# # Do not edit this cell! (20 points)
#
# The notebook will also be checked for overall clean code requirements as specified at the **top** of this notebook. Some of these requirements include (review the top cells for more specifics):
#
# * Notebook begins at cell [1] and runs on any machine in its entirety.
# * PEP 8 format is applied throughout (including lengths of comment and code lines).
# * No additional code or imports in the notebook that is not needed for the workflow.
# * Notebook is fully reproducible. This means:
#    * reproducible paths using the os module.
#    * data downloaded using code in the notebook.
#    * all imports at top of notebook.

# ## BONUS - Export a  .CSV File to Share (10 points possible)
#
# This is optional - if you export a **.csv** file with the columns specified above: Site, Date and NDVI Value you can get an additional 10 points.
#
# * FULL CREDIT: File exists in csv format and contains the columns specified.
# We will check your github repo for this file!
#

# In[13]:

# Export CSV with mean NDVI to outputs folder
mean_ndvi_all.to_csv(os.path.join("ndvi-automation", "outputs", "mean_ndvi.csv"),
                     header=["Site", "NDVI Value"],
                     index_label="Date")

# In[ ]:


#
#   Lightnet related postprocessing
#   Copyright EAVISE
#

import logging
import numpy as np
import torch
from ..util import BaseTransform
from ..._imports import bb

__all__ = ['NMS', 'NMSFast', 'NMSSoft', 'NMSSoftFast', 'NonMaxSuppression']
log = logging.getLogger(__name__)

class NMS(BaseTransform):
    """ Performs non-maximal suppression on the bounding boxes, filtering boxes with a high overlap.
    This function can either work on a pytorch bounding box tensor or a brambox dataframe.

    Args:
        nms_thresh (Number [0-1]): Overlapping threshold to filter detections with non-maxima suppresion
        class_nms (Boolean, optional): Whether to perform nms per class; Default **True**
        force_cpu (Boolean, optional): Whether to force a part of the computation on CPU (tensor only, see Note); Default **True**
        reset_index (Boolean, optional): Whether to reset the index of the returned dataframe (dataframe only); Default **True**

    Input:
        boxes (Tensor [Boxes x 7] or pandas.Dataframe): bounding boxes

    Returns:
        boxes (Tensor [Boxes x 7] or pandas.Dataframe): filtered bounding boxes

    Note:
        This post-processing function expects the input bounding boxes to be either a PyTorch tensor or a brambox dataframe.

        The PyTorch tensor needs to be formatted as follows: **[batch_num, x_tl, y_tl, x_br, y_br, confidence, class_id]** for every bounding box.
        This corresponds to the output from the Get***Boxes classes available in lightnet.

        The brambox dataframe should be a detection dataframe,
        as it needs the x_top_left, y_top_left, width, height, confidence and class_label columns.

    Note:
        A part of the computation of NMS involve a sequential loop through the boxes.
        This is quite difficult to implement efficiently on the GPU, using the PyTorch API.
        We thus concluded empirically that it is usually more efficient to move GPU tensors to the CPU for this part of the computations. |br|
        By passing false to `force_cpu`, you can disable this behavior and perform all the computations on the original device of the input tensor.
    """
    def __init__(self, nms_thresh, class_nms=True, force_cpu=True, reset_index=True):
        super().__init__()
        self.nms_thresh = nms_thresh
        self.class_nms = class_nms
        self.force_cpu = force_cpu
        self.reset_index = reset_index

    def forward(self, boxes):
        if isinstance(boxes, torch.Tensor):
            return self._torch(boxes)
        else:
            return self._pandas(boxes)

    def _torch(self, boxes):
        if boxes.numel() == 0:
            return boxes

        batches = boxes[:, 0]
        keep = torch.empty(boxes.shape[0], dtype=torch.bool, device=boxes.device)
        for batch in torch.unique(batches, sorted=False):
            mask = batches == batch
            keep[mask] = self._torch_nms(boxes[mask])

        return boxes[keep]

    def _torch_nms(self, boxes):
        if boxes.numel() == 0:
            return boxes

        bboxes = boxes[:, 1:5]
        scores = boxes[:, 5]
        classes = boxes[:, 6]

        # Sort coordinates by descending score
        scores, order = scores.sort(0, descending=True)
        x1, y1, x2, y2 = bboxes[order].split(1, 1)

        # Compute dx and dy between each pair of boxes (these mat contain every pair twice...)
        dx = (x2.min(x2.t()) - x1.max(x1.t())).clamp(min=0)
        dy = (y2.min(y2.t()) - y1.max(y1.t())).clamp(min=0)

        # Compute iou
        intersections = dx * dy
        areas = (x2 - x1) * (y2 - y1)
        unions = (areas + areas.t()) - intersections
        ious = intersections / unions

        # Filter based on iou (and class)
        conflicting = (ious > self.nms_thresh).triu(1)
        if self.class_nms:
            classes = classes[order]
            same_class = (classes.unsqueeze(0) == classes.unsqueeze(1))
            conflicting = (conflicting & same_class)

        if self.force_cpu:
            conflicting = conflicting.cpu()
        keep = torch.zeros(conflicting.shape[0], dtype=torch.bool, device=conflicting.device)
        supress = torch.zeros(conflicting.shape[0], dtype=torch.bool, device=conflicting.device)
        for i, row in enumerate(conflicting):
            if not supress[i]:
                keep[i] = True
                supress[row] = True

        keep = keep.to(boxes.device)
        return keep.scatter(0, order, keep)

    @torch.jit.ignore
    def _pandas(self, boxes):
        if len(boxes.index) == 0:
            return boxes
        boxes = boxes.groupby('image', group_keys=False, observed=True).apply(self._pandas_nms)
        if self.reset_index:
            return boxes.reset_index(drop=True)
        return boxes

    def _pandas_nms(self, boxes):
        boxes = boxes.sort_values('confidence', ascending=False)
        ious = bb.stat.coordinates.iou(boxes, boxes, bias=0)
        scores = boxes['confidence'].values

        # Filter based on iou (and class)
        conflicting = np.triu(ious > self.nms_thresh, 1)
        if self.class_nms:
            classes = boxes['class_label'].values
            same_class = (classes[None, ...] == classes[..., None])
            conflicting = (conflicting & same_class)

        keep = np.zeros(conflicting.shape[0], dtype=np.bool)
        supress = np.zeros(conflicting.shape[0], dtype=np.bool)
        for i, row in enumerate(conflicting):
            if not supress[i]:
                keep[i] = True
                supress[row] = True

        return boxes[keep]

class NMSFast(NMS):
    """ Performs fast NMS on the bounding boxes, filtering boxes with a high overlap.
    This function can either work on a pytorch bounding box tensor or a brambox dataframe.

    This faster alternative makes a small "mistake" during NMS computation,
    in order to remove a necessary loop in the code, allowing it to run faster.
    The speed increase should be mostly notable when performing NMS with PyTorch tensors on the GPU.

    The difference is explained in the image below, where the boxes A and B overlap enough to be filtered out
    and the boxes B and C as well (but A and C do not). |br|
    Regular NMS will keep box C in this situation, because box B gets filtered out and is thus not there to remove C.
    Fast NMS will not do this and will only keep box A in this situation. |br|
    Depending on the use-case (closely clustered and overlapping objects), this might be a problem or not.

    .. figure:: /.static/api/nms-fast.*
       :width: 100%
       :alt: Fast NMS problem

       Regular NMS will keep both boxes A and C, but Fast NMS will only keep A in this example.

    Args:
        nms_thresh (Number [0-1]): Overlapping threshold to filter detections with non-maxima suppresion
        class_nms (Boolean, optional): Whether to perform nms per class; Default **True**
        reset_index (Boolean, optional): Whether to reset the index of the returned dataframe (dataframe only); Default **True**

    Input:
        boxes (Tensor [Boxes x 7] or pandas.Dataframe): bounding boxes

    Returns:
        boxes (Tensor [Boxes x 7] or pandas.Dataframe): filtered bounding boxes

    Note:
        This post-processing function expects the input bounding boxes to be either a PyTorch tensor or a brambox dataframe.

        The PyTorch tensor needs to be formatted as follows: **[batch_num, x_tl, y_tl, x_br, y_br, confidence, class_id]** for every bounding box.
        This corresponds to the output from the Get***Boxes classes available in lightnet.

        The brambox dataframe should be a detection dataframe,
        as it needs the x_top_left, y_top_left, width, height, confidence and class_label columns.
    """
    def __init__(self, nms_thresh, class_nms=True, reset_index=True):
        super().__init__(nms_thresh, class_nms, False, reset_index)

    def _torch_nms(self, boxes):
        if boxes.numel() == 0:
            return boxes

        bboxes = boxes[:, 1:5]
        scores = boxes[:, 5]
        classes = boxes[:, 6]

        # Sort coordinates by descending score
        scores, order = scores.sort(0, descending=True)
        x1, y1, x2, y2 = bboxes[order].split(1, 1)

        # Compute dx and dy between each pair of boxes (these mat contain every pair twice...)
        dx = (x2.min(x2.t()) - x1.max(x1.t())).clamp(min=0)
        dy = (y2.min(y2.t()) - y1.max(y1.t())).clamp(min=0)

        # Compute iou
        intersections = dx * dy
        areas = (x2 - x1) * (y2 - y1)
        unions = (areas + areas.t()) - intersections
        ious = intersections / unions

        # Filter based on iou (and class)
        conflicting = (ious > self.nms_thresh).triu(1)
        if self.class_nms:
            classes = classes[order]
            same_class = (classes.unsqueeze(0) == classes.unsqueeze(1))
            conflicting = (conflicting & same_class)

        keep = conflicting.sum(0) == 0
        return keep.scatter(0, order, keep)

    @torch.jit.ignore
    def _pandas_nms(self, boxes):
        boxes = boxes.sort_values('confidence', ascending=False)
        ious = bb.stat.coordinates.iou(boxes, boxes, bias=0)
        scores = boxes['confidence'].values

        conflicting = np.triu(ious > self.nms_thresh, 1)
        if self.class_nms:
            classes = boxes['class_label'].values
            same_class = (classes[None, ...] == classes[..., None])
            conflicting = (conflicting & same_class)

        # Return filtered boxes
        keep = conflicting.sum(0) == 0
        return boxes[keep]

class NMSSoft(BaseTransform):
    """ Performs soft NMS with exponential decaying on the bounding boxes, as explained in :cite:`soft_nms`.
    This function can either work on a pytorch bounding box tensor or a brambox dataframe.

    Args:
        sigma (Number): Sensitivity value for the confidence rescaling (exponential decay)
        conf_thresh (Number [0-1], optional): Confidence threshold to filter the bounding boxes after decaying them; Default **0**
        class_nms (Boolean, optional): Whether to perform nms per class; Default **True**
        force_cpu (Boolean, optional): Whether to force a part of the computation on CPU (tensor only, see Note); Default **True**
        reset_index (Boolean, optional): Whether to reset the index of the returned dataframe (dataframe only); Default **True**

    Input:
        boxes (Tensor [Boxes x 7] or pandas.Dataframe): bounding boxes

    Returns:
        boxes (Tensor [Boxes x 7] or pandas.Dataframe): filtered bounding boxes

    Note:
        This post-processing function expects the input bounding boxes to be either a PyTorch tensor or a brambox dataframe.

        The PyTorch tensor needs to be formatted as follows: **[batch_num, x_tl, y_tl, x_br, y_br, confidence, class_id]** for every bounding box.
        This corresponds to the output from the Get***Boxes classes available in lightnet.

        The brambox dataframe should be a detection dataframe,
        as it needs the x_top_left, y_top_left, width, height, confidence and class_label columns.
    """
    def __init__(self, sigma, conf_thresh=0, class_nms=True, force_cpu=True, reset_index=True):
        super().__init__()
        self.sigma = sigma
        self.conf_thresh = conf_thresh
        self.class_nms = class_nms
        self.force_cpu = force_cpu
        self.reset_index = reset_index

    def forward(self, boxes):
        if isinstance(boxes, torch.Tensor):
            return self._torch(boxes)
        else:
            return self._pandas(boxes)

    def _torch(self, boxes):
        if boxes.numel() == 0:
            return boxes

        batches = boxes[:, 0]
        for batch in torch.unique(batches, sorted=False):
            mask = batches == batch
            boxes[mask, 5] = self._torch_nms(boxes[mask])

        if self.conf_thresh > 0:
            keep = boxes[:, 5] > self.conf_thresh
            return boxes[keep]

        return boxes

    def _torch_nms(self, boxes):
        if boxes.numel() == 0:
            return boxes

        bboxes = boxes[:, 1:5]
        scores = boxes[:, 5]
        classes = boxes[:, 6]

        # Sort coordinates by descending score
        scores, order = scores.sort(0, descending=True)
        x1, y1, x2, y2 = bboxes[order].split(1, 1)

        # Compute dx and dy between each pair of boxes (these mat contain every pair twice...)
        dx = (x2.min(x2.t()) - x1.max(x1.t())).clamp(min=0)
        dy = (y2.min(y2.t()) - y1.max(y1.t())).clamp(min=0)

        # Compute iou
        intersections = dx * dy
        areas = (x2 - x1) * (y2 - y1)
        unions = (areas + areas.t()) - intersections
        ious = intersections / unions

        # Filter class
        if self.class_nms:
            classes = classes[order]
            same_class = (classes.unsqueeze(0) == classes.unsqueeze(1))
            ious *= same_class

        # Decay scores
        decay = torch.exp(-(ious ** 2) / self.sigma)
        if self.force_cpu:
            scores = scores.cpu()
            order = order.cpu()
            decay = decay.cpu()

        tempscores = scores.clone()
        for _ in range(scores.shape[0]):
            maxidx = tempscores.argmax()
            maxscore = tempscores[maxidx]
            if maxscore <= self.conf_thresh:
                break

            tempscores[maxidx] = -1
            mask = tempscores != -1
            tempscores[mask] *= decay[maxidx, mask]
            scores[mask] = tempscores[mask]

        scores = scores.to(boxes.device)
        order = order.to(boxes.device)
        return scores.scatter(0, order, scores)

    @torch.jit.ignore
    def _pandas(self, boxes):
        if len(boxes.index) == 0:
            return boxes

        boxes = boxes.groupby('image', group_keys=False, observed=True).apply(self._pandas_nms)
        if self.conf_thresh > 0:
            boxes = boxes[boxes.confidence > self.conf_thresh].copy()
        if self.reset_index:
            return boxes.reset_index(drop=True)

        return boxes

    def _pandas_nms(self, boxes):
        boxes = boxes.sort_values('confidence', ascending=False)
        scores = boxes['confidence'].values
        ious = bb.stat.coordinates.iou(boxes, boxes, bias=0)

        # Filter class
        if self.class_nms:
            classes = boxes['class_label'].values
            same_class = (classes[None, ...] == classes[..., None])
            ious *= same_class

        # Decay scores
        decay = np.exp(-(ious ** 2) / self.sigma)
        tempscores = scores.copy()
        for _ in range(scores.shape[0]):
            maxidx = tempscores.argmax()
            maxscore = tempscores[maxidx]
            if maxscore <= self.conf_thresh:
                break

            tempscores[maxidx] = -1
            mask = tempscores != -1
            tempscores[mask] *= decay[maxidx, mask]
            scores[mask] = tempscores[mask]

        # Set scores back
        boxes['confidence'] = scores
        return boxes

class NMSSoftFast(BaseTransform):
    """ Faster version of SoftNMS which filters boxes with a high overlap, using exponential decay.
    This function can either work on a pytorch bounding box tensor or a brambox dataframe.

    This version of NMS makes the same "mistake" as :class:`~lightnet.data.transform.NMSFast`,
    which in turn allows it to be faster than the regular :class:`~lightnet.data.transform.NMSSoft` algorithm.
    The speed increase should be mostly notable when performing NMS with PyTorch tensors on the GPU.

    Args:
        sigma (Number): Sensitivity value for the confidence rescaling (exponential decay)
        conf_thresh (Number [0-1], optional): Confidence threshold to filter the bounding boxes after decaying them; Default **0**
        class_nms (Boolean, optional): Whether to perform nms per class; Default **True**
        reset_index (Boolean, optional): Whether to reset the index of the returned dataframe (dataframe only); Default **True**

    Input:
        boxes (Tensor [Boxes x 7] or pandas.Dataframe): bounding boxes

    Returns:
        boxes (Tensor [Boxes x 7] or pandas.Dataframe): filtered bounding boxes

    Note:
        This post-processing function expects the input bounding boxes to be either a PyTorch tensor or a brambox dataframe.

        The PyTorch tensor needs to be formatted as follows: **[batch_num, x_tl, y_tl, x_br, y_br, confidence, class_id]** for every bounding box.
        This corresponds to the output from the Get***Boxes classes available in lightnet.

        The brambox dataframe should be a detection dataframe,
        as it needs the x_top_left, y_top_left, width, height, confidence and class_label columns.
    """
    def __init__(self, sigma, conf_thresh=0, class_nms=True, reset_index=True):
        super().__init__()
        self.sigma = sigma
        self.conf_thresh = conf_thresh
        self.class_nms = class_nms
        self.reset_index = reset_index

    def forward(self, boxes):
        if isinstance(boxes, torch.Tensor):
            return self._torch(boxes)
        else:
            return self._pandas(boxes)

    def _torch(self, boxes):
        if boxes.numel() == 0:
            return boxes

        batches = boxes[:, 0]
        for batch in torch.unique(batches, sorted=False):
            mask = batches == batch
            boxes[mask, 5] = self._torch_nms(boxes[mask])

        if self.conf_thresh > 0:
            keep = boxes[:, 5] > self.conf_thresh
            return boxes[keep]

        return boxes

    def _torch_nms(self, boxes):
        if boxes.numel() == 0:
            return boxes

        bboxes = boxes[:, 1:5]
        scores = boxes[:, 5]
        classes = boxes[:, 6]

        # Sort coordinates by descending score
        scores, order = scores.sort(0, descending=True)
        x1, y1, x2, y2 = bboxes[order].split(1, 1)

        # Compute dx and dy between each pair of boxes (these mat contain every pair twice...)
        dx = (x2.min(x2.t()) - x1.max(x1.t())).clamp(min=0)
        dy = (y2.min(y2.t()) - y1.max(y1.t())).clamp(min=0)

        # Compute iou
        intersections = dx * dy
        areas = (x2 - x1) * (y2 - y1)
        unions = (areas + areas.t()) - intersections
        ious = intersections / unions

        # Filter class
        if self.class_nms:
            classes = classes[order]
            same_class = (classes.unsqueeze(0) == classes.unsqueeze(1))
            ious *= same_class

        # Decay scores
        decay = ious.triu(1)
        decay = torch.exp(-(decay ** 2) / self.sigma).prod(0)
        scores *= decay
        return scores.scatter(0, order, scores)

    @torch.jit.ignore
    def _pandas(self, boxes):
        if len(boxes.index) == 0:
            return boxes

        boxes = boxes.groupby('image', group_keys=False, observed=True).apply(self._pandas_nms)
        if self.conf_thresh > 0:
            boxes = boxes[boxes.confidence > self.conf_thresh].copy()
        if self.reset_index:
            return boxes.reset_index(drop=True)

        return boxes

    def _pandas_nms(self, boxes):
        boxes = boxes.sort_values('confidence', ascending=False)
        scores = boxes['confidence'].values
        ious = bb.stat.coordinates.iou(boxes, boxes, bias=0)

        # Filter class
        if self.class_nms:
            classes = boxes['class_label'].values
            same_class = (classes[None, ...] == classes[..., None])
            ious *= same_class

        # Decay scores
        decay = np.triu(ious, 1)
        decay = np.prod(np.exp(-(decay ** 2) / self.sigma), 0)
        scores *= decay

        # Set scores back
        boxes['confidence'] = scores
        return boxes

def NonMaxSuppression(*args, **kwargs):
    log.deprecated('NonMaxSuppression is deprecated, please use "NMS"')
    return NMS(*args, **kwargs)


import os
import sys
import random
import math
import numpy as np
import skimage.io
import matplotlib
import matplotlib.pyplot as plt
from mrcnn import utils
import mrcnn.model as modellib
from mrcnn import visualize

import cv2
import time

class MaskRCNN():

    def __init__(self,mask_rcnn_root):
        # Root directory of the project
        ROOT_DIR = mask_rcnn_root

        # Import Mask RCNN
        sys.path.append(ROOT_DIR)  # To find local version of the library
        sys.path.append(os.path.join(ROOT_DIR, "samples/coco/"))  # To find local version
        import coco
        # Directory to save logs and trained model
        MODEL_DIR = os.path.join(ROOT_DIR, "logs")

        # Local path to trained weights file
        COCO_MODEL_PATH = os.path.join(ROOT_DIR, "mask_rcnn_coco.h5")
        # Download COCO trained weights from Releases if needed
        if not os.path.exists(COCO_MODEL_PATH):
            utils.download_trained_weights(COCO_MODEL_PATH)

        class InferenceConfig(coco.CocoConfig):
            # Set batch size to 1 since we'll be running inference on
            # one image at a time. Batch size = GPU_COUNT * IMAGES_PER_GPU
            GPU_COUNT = 1
            IMAGES_PER_GPU = 1

        config = InferenceConfig()
        config.display()

        # Create model object in inference mode.
        self.model = modellib.MaskRCNN(mode="inference", model_dir=MODEL_DIR, config=config)

        # Load weights trained on MS-COCO
        self.model.load_weights(COCO_MODEL_PATH, by_name=True)

        self.class_names = ['BG', 'person', 'bicycle', 'car', 'motorcycle', 'airplane',
                       'bus', 'train', 'truck', 'boat', 'traffic light',
                       'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird',
                       'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant', 'bear',
                       'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie',
                       'suitcase', 'frisbee', 'skis', 'snowboard', 'sports ball',
                       'kite', 'baseball bat', 'baseball glove', 'skateboard',
                       'surfboard', 'tennis racket', 'bottle', 'wine glass', 'cup',
                       'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple',
                       'sandwich', 'orange', 'broccoli', 'carrot', 'hot dog', 'pizza',
                       'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed',
                       'dining table', 'toilet', 'tv', 'laptop', 'mouse', 'remote',
                       'keyboard', 'cell phone', 'microwave', 'oven', 'toaster',
                       'sink', 'refrigerator', 'book', 'clock', 'vase', 'scissors',
                       'teddy bear', 'hair drier', 'toothbrush']
        self.ax = []
        self.plt = []

    def devide_images(self,image_path_list):
        image_save_index = 0
        for image_path in image_path_list:
            image = skimage.io.imread(image_path)
            results = self.model.detect([image],verbose=1)
            r = results[0]

            rois = r["rois"]
            N = rois.shape[0]
            class_ids = r["class_ids"]
            masks = r["masks"]
            scores = r["scores"]
            if not os.path.exists("output"):
                os.mkdir("output")
            for i in range(N):
                if not np.any(rois[i]):continue
                y1,x1,y2,x2 = rois[i]
                #print(x1,x2,y1,y2)
                #print(image.shape)
                new_image = image[y1:y2,x1:x2,:]

                #print(new_image.shape)
                mask = masks[:,:,i]

                label = self.class_names[class_ids[i]]

                plt.imshow(new_image)

                try:
                    plt.savefig("output\%s\%s%s.png"%(label,label,image_save_index))
                    image_save_index += 1
                except FileNotFoundError:

                     os.mkdir("output\%s" % label)
                     plt.savefig("output\%s\%s%s.png" % (label, label, image_save_index))
                     image_save_index += 1
                except:
                    continue

def tst():
    # set MaskRCNN root: https://github.com/matterport/Mask_RCNN
    temp = MaskRCNN("G:\wbc\GitHub\Mask_RCNN")
    # images
    temp.devide_images(["C:\\Users\Administrator\OneDrive - shanghaitech.edu.cn\图片\\"+ i for i in
                        [
                        "9247489789_132c0d534a_z.jpg",
                          "3132016470_c27baa00e8_z.jpg",
                          "3651581213_f81963d1dd_z.jpg"

                          ]

                        ]

                       )

if __name__ == '__main__':
    tst()


import os
import argparse
import cv2
import numpy as np
import sys
import time
from threading import Thread
import importlib.util
import re

def loadLabels(labelPath):
    p = re.compile(r'\s*(\d+)(.+)')
    with open(labelPath, 'r', encoding='utf-8') as labelFile:
        lines = (p.match(line).groups() for line in labelFile.readlines())
        return {int(num): text.strip() for num, text in lines}

pkg = importlib.util.find_spec('tflite_runtime')

if pkg:
    print('tflite')
    from tflite_runtime.interpreter import Interpreter
else:
    print('tf.lite')
    from tensorflow.lite.python.interpreter import Interpreter

interpreter = Interpreter(model_path='tflite_model/model.tflite')
interpreter.allocate_tensors()

labels = loadLabels('tflite_model/labels.txt')

input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
height = input_details[0]['shape'][1]
width = input_details[0]['shape'][2]
floating_model = (input_details[0]['dtype'] == np.float32)

input_mean = 127.5
input_std = 127.5

frame_rate_calc = 0
freq = cv2.getTickFrequency()

cap = cv2.VideoCapture(0)

while cap.isOpened():

    t1 = cv2.getTickCount()

    ret, frame = cap.read()

    scr = frame.copy()

    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    frame_resized = cv2.resize(frame_rgb, (width, height))

    # 加一個維度 補 batch size
    input_data = np.expand_dims(frame_resized, axis=0)

    if floating_model:
        input_data = (np.float32(input_data) - input_mean) / input_std

    interpreter.set_tensor( (input_details[0]['index']) , input_data)

    interpreter.invoke()

    results = interpreter.get_tensor( output_details[0]['index'])[0]

    print_res = f'FPS: {frame_rate_calc:.2f} , Label: {labels[np.argmax(results)]}'

    t2 = cv2.getTickCount()
    time1 = (t2-t1)/freq
    frame_rate_calc= 1/time1

    cv2.putText(frame, print_res, (30,50),cv2.FONT_HERSHEY_SIMPLEX,1,(255,255,0),2,cv2.LINE_AA)

    cv2.imshow('Object detector', frame)

    if cv2.waitKey(1) == ord('q'):
        break

# Clean up
cv2.destroyAllWindows()
cap.release()


#!/usr/bin/env python
import cv2
import numpy as np
import time
import sys
import signal
from time import sleep
import urllib.request
net = cv2.dnn.readNet("C:/Users/wenchih/Desktop/Robot_V3/yolov4-custom_best.weights","C:/Users/wenchih/Desktop/Robot_V3/yolov3.cfg")
classes = []
with open('MyObject.names','r') as f:
    classes = f.read().splitlines()
# print(classes)
count=0

cap = cv2.VideoCapture(0)

while True:
    _, img = cap.read()
    height, width, _ = img.shape # 前兩個維度是圖片的高度與寬度
    # 第三個維度則是圖片的 channel,彩色=3
    blob = cv2.dnn.blobFromImage(img, 1/255, (416, 416), (0, 0, 0), swapRB = True, crop = False)
    # Normalize the pixel by value 255
    # conver BGR to RGB
    '''
    for b in blob:
        for n, img_blob in enumerate(b):
            cv2.imshow(str(n), img_blob)
    '''
    net.setInput(blob) # set the input from the blob into the network
    ouput_layers_names = net.getUnconnectedOutLayersNames() # To get the output layers names
    layerOutputs = net.forward(ouput_layers_names) # passing names to forward function
    # run the forward pass and obtain their output layer
    boxes = []
    confidences = []
    class_ids = []
    for output in layerOutputs: # extract all the information
        for detection in output: # extract information from each of the output
            scores = detection[5:]  # store all the acting classes prediction
            class_id = np.argmax(scores)  # find out the location contains the higher scores
            confidence = scores[class_id]  # class_ids?
            if confidence > 0.5:
                center_x = int(detection[0]*width)
                center_y = int(detection[1]*height)
                w = int(detection[2]*width)
                h = int(detection[3]*height)
                # the value in detection is normalized
                # get the position for the upper left corner
                x = int(center_x - w/2)
                y = int(center_y - h/2)
                # bounding boxes information
                boxes.append([x, y, w, h])
                confidences.append((float(confidence)))
                class_ids.append(class_id)

                if (center_x < (0.3*width)):  # 0.3 / 0.4 / 0.3
                    left = 1
                    right = 0
                    print('left')
                    #time.sleep(0.3)
                elif (center_x > (0.7*width)):
                    left = 0
                    right = 1
                    print('right')
                    #time.sleep(0.3)
                else:
                    left = 0
                    right = 0
                    print('forward')
                    #time.sleep(0.3)

    indexes = cv2.dnn.NMSBoxes(boxes, confidences, 0.5, 0.4) # ??
    font = cv2.FONT_HERSHEY_PLAIN # choose font
    colors = np.random.uniform(0, 255, size=(len(boxes), 3)) # 3channels

    if len(indexes) > 0:
        for i in indexes.flatten():
            x, y, w, h = boxes[i]
            label = str(classes[class_ids[i]])
            confidence = str(round(confidences[i], 2))
            color = colors[i]
            cv2.rectangle(img, (x, y), (x+w, y+h), color, 2)
            cv2.putText(img, label + " " + confidence, (x, y+20), font, 2, (255, 255, 255), 2)
    cv2.namedWindow('Image', cv2.WINDOW_NORMAL) # windows can adjust the size automatically
    cv2.imshow('Image', img)
    key = cv2.waitKey(1)  # Wait for user's instruct, the parameter is the time to wait(minisecond)
    if key==27:
        break

cap.release()
cv2.destroyAllWindows #  shut down all the windows of OpenCV


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

import time

import wiringpi as wi
import numpy as np

# from scipy.optimize import least_squares

class MPU9250(object):
    wi.wiringPiSetup()
    i2c = wi.I2C()

    _address = 0x68  # addresses of gyroscope and accelerometer
    _addr_AK8963 = 0x0C  # a address of magnetometer (self.AK8963)
    # sensor constant
    _REG_PWR_MGMT_1 = 0x6B
    _REG_INT_PIN_CFG = 0x37
    _REG_ACCEL_CONFIG1 = 0x1C
    _REG_ACCEL_CONFIG2 = 0x1D
    _REG_GYRO_CONFIG = 0x1B

    _MAG_MODE_POWERDOWN = 0  # 磁気センサpower down
    _MAG_MODE_SERIAL_1 = 1  # 磁気センサ8Hz連続測定モード
    _MAG_MODE_SERIAL_2 = 2  # 磁気センサ100Hz連続測定モード
    _MAG_MODE_SINGLE = 3  # 磁気センサ単発測定モード
    _MAG_MODE_EX_TRIGER = 4  # 磁気センサ外部トリガ測定モード
    _MAG_MODE_SELF_TEST = 5  # 磁気センサセルフテストモード
    _MAG_ACCESS = False  # 磁気センサへのアクセス可否
    _MAG_MODE = 0  # 磁気センサモード
    _MAG_BIT = 16  # 磁気センサが出力するbit数
    _gyro_range = 1000  # 250, 500, 1000, 2000　'dps'から選択
    _accel_range = 8  # +-2, +-4, +-8, +-16 'g'から選択
    _mag_range = 4912  # 'μT'

    mpu9250 = i2c.setup(_address)  # i2cアドレス0x68番地をmpu9250として設定(アドレスは$sudo i2cdetect 1で見られる)
    AK8963 = i2c.setup(_addr_AK8963)

    # 加速度の測定レンジを設定
    _val = 8  # val = 16, 8, 4, 2(default)

    def __init__(self):
        # オフセット用変数
        self._offset_gyro_x = 0
        self._offset_gyro_y = 0
        self._offset_gyro_z = 0
        self._offset_accel_x = 0
        self._offset_accel_y = 0
        self._offset_accel_z = 0
        self._offset_mag_x = -7.497336140
        self._offset_mag_y = -57.461323250
        self._offset_mag_z = 63.096101850

        self._reset_register()
        self._power_wakeup()
        self._gyro_coefficient = self._gyro_range / float(0x8000)  # coefficient : sensed decimal val to dps val.
        self._accel_coefficient = self._accel_range / float(0x8000)  # coefficient : sensed decimal val to g val
        self._mag_coefficient_16 = self._mag_range / 32760.0  # coefficient : sensed decimal val to μT val (16bit)
        self._mag_coefficient_14 = self._mag_range / 8190.0  # coefficient : sensed decimal val to μT val (14bit)
        self._set_accel_range(val=self._accel_range, _calibration=True)
        self._set_gyro_range(val=self._gyro_range, _calibration=True)
        # self._set_accel_range(val=self._accel_range, _calibration=False)
        # setself._gyro_range(val=self._gyro_range, _calibration=False)
        self._set_mag_register('100Hz', '16bit')

    # レジスタを初期設定に戻す。
    def _reset_register(self):
        if self._MAG_ACCESS is True:
            self.i2c.writeReg8(self.AK8963, 0x0B, 0x01)
        self.i2c.writeReg8(self.mpu9250, 0x6B, 0x80)
        self._MAG_ACCESS = False
        time.sleep(0.1)

    # センシング可能な状態にする。
    def _power_wakeup(self):
        # PWR_MGMT_1をクリア
        self.i2c.writeReg8(self.mpu9250, self._REG_PWR_MGMT_1, 0x00)
        time.sleep(0.1)
        # I2Cで磁気センサ機能(self.AK8963)へアクセスできるようにする(BYPASS_EN=1)
        self.i2c.writeReg8(self.mpu9250, self._REG_INT_PIN_CFG, 0x02)
        self._MAG_ACCESS = True
        time.sleep(0.1)

    def _set_accel_range(self, val, _calibration=False):
        # +-2g (00), +-4g (01), +-8g (10), +-16g (11)
        if val == 16:
            _accel_range = 16
            _data = 0x18
        elif val == 8:
            _accel_range = 8
            _data = 0x10
        elif val == 4:
            _accel_range = 4
            _data = 0x08
        else:
            _accel_range = 2
            _data = 0x00
        print("set _accel_range=%d [g]" % _accel_range)
        self.i2c.writeReg8(self.mpu9250, self._REG_ACCEL_CONFIG1, _data)
        self._accel_coefficient = _accel_range / float(0x8000)
        time.sleep(0.1)

        # Calibration
        if _calibration is True:
            self._calib_accel(1000)
        return

    # ジャイロの測定レンジを設定します。
    # val= 2000, 1000, 500, 250(default)
    def _set_gyro_range(self, val, _calibration=False):
        if val == 2000:
            _gyro_range = 2000
            _data = 0x18
        elif val == 1000:
            _gyro_range = 1000
            _data = 0x10
        elif val == 500:
            _gyro_range = 500
            _data = 0x08
        else:
            _gyro_range = 250
            _data = 0x00
        print("set _gyro_range=%d [dps]" % _gyro_range)
        self.i2c.writeReg8(self.mpu9250, self._REG_GYRO_CONFIG, _data)
        self._gyro_coefficient = _gyro_range / float(0x8000)
        time.sleep(0.1)

        # Calibration
        if _calibration is True:
            self._calib_gyro(1000)
        return

    # 磁気センサのレジスタを設定する
    def _set_mag_register(self, _mode, _bit, _calibration=False):
        if self._MAG_ACCESS is False:
            # 磁気センサへのアクセスが有効になっていない場合は例外
            raise Exception('001 Access to a sensor is invalid.')

        _writeData = 0x00
        # 測定モードの設定
        if _mode == '8Hz':  # Continuous measurement mode 1
            _writeData = 0x02
            self._MAG_MODE = self._MAG_MODE_SERIAL_1
        elif _mode == '100Hz':  # Continuous measurement mode 2
            _writeData = 0x06
            self._MAG_MODE = self._MAG_MODE_SERIAL_2
        elif _mode == 'POWER_DOWN':  # Power down mode
            _writeData = 0x00
            self._MAG_MODE = self._MAG_MODE_POWERDOWN
        elif _mode == 'EX_TRIGER':  # Trigger measurement mode
            _writeData = 0x04
            self._MAG_MODE = self._MAG_MODE_EX_TRIGER
        elif _mode == 'SELF_TEST':  # self test mode
            _writeData = 0x08
            self._MAG_MODE = self._MAG_MODE_SELF_TEST
        else:  # _mode='SINGLE'    # single measurment mode
            _writeData = 0x01
            self._MAG_MODE = self._MAG_MODE_SINGLE

        # 出力するbit数
        if _bit == '14bit':  # output 14bit
            _writeData = _writeData | 0x00
            self._MAG_BIT = 14
        else:  # _bit='16bit'      # output 16bit
            _writeData = _writeData | 0x10
            self._MAG_BIT = 16
        print("set self._MAG_MODE=%s, %d bit" % (_mode, self._MAG_BIT))
        self.i2c.writeReg8(self.AK8963, 0x0A, _writeData)
        time.sleep(0.1)

        # Calibration
        if _calibration is True:
            self._calib_mag(3000)
        return

    # センサからのデータはそのまま使おうとするとunsignedとして扱われるため、signedに変換(16ビット限定）
    def _u2s(self, unsigneddata):
        if unsigneddata & (0x01 << 15):
            return -1 * ((unsigneddata ^ 0xffff) + 1)
        return unsigneddata

    # 加速度値を取得
    def get_accel(self):
        mpu9250 = self.mpu9250
        i2c = self.i2c
        ACCEL_XOUT_H = i2c.readReg8(mpu9250, 0x3B)
        ACCEL_XOUT_L = i2c.readReg8(mpu9250, 0x3C)
        ACCEL_YOUT_H = i2c.readReg8(mpu9250, 0x3D)
        ACCEL_YOUT_L = i2c.readReg8(mpu9250, 0x3E)
        ACCEL_ZOUT_H = i2c.readReg8(mpu9250, 0x3F)
        ACCEL_ZOUT_L = i2c.readReg8(mpu9250, 0x40)
        raw_x = self._accel_coefficient * self._u2s(ACCEL_XOUT_H << 8 | ACCEL_XOUT_L) + self._offset_accel_x
        raw_y = self._accel_coefficient * self._u2s(ACCEL_YOUT_H << 8 | ACCEL_YOUT_L) + self._offset_accel_y
        raw_z = self._accel_coefficient * self._u2s(ACCEL_ZOUT_H << 8 | ACCEL_ZOUT_L) + self._offset_accel_z
        return raw_x, raw_y, raw_z

    # ジャイロ値を取得
    def get_gyro(self):
        mpu9250 = self.mpu9250
        i2c = self.i2c
        GYRO_XOUT_H = i2c.readReg8(mpu9250, 0x43)
        GYRO_XOUT_L = i2c.readReg8(mpu9250, 0x44)
        GYRO_YOUT_H = i2c.readReg8(mpu9250, 0x45)
        GYRO_YOUT_L = i2c.readReg8(mpu9250, 0x46)
        GYRO_ZOUT_H = i2c.readReg8(mpu9250, 0x47)
        GYRO_ZOUT_L = i2c.readReg8(mpu9250, 0x48)
        raw_x = self._gyro_coefficient * self._u2s(GYRO_XOUT_H << 8 | GYRO_XOUT_L) + self._offset_gyro_x
        raw_y = self._gyro_coefficient * self._u2s(GYRO_YOUT_H << 8 | GYRO_YOUT_L) + self._offset_gyro_y
        raw_z = self._gyro_coefficient * self._u2s(GYRO_ZOUT_H << 8 | GYRO_ZOUT_L) + self._offset_gyro_z
        return raw_x, raw_y, raw_z

    # 磁気値を取得
    def get_mag(self):
        AK8963 = self.AK8963
        i2c = self.i2c

        if self._MAG_ACCESS is False:
            # 磁気センサへのアクセスが有効になっていない場合は例外
            raise Exception('002 Access to a sensor is invalid.')

        # 事前処理
        if self._MAG_MODE == self._MAG_MODE_SINGLE:
            # 単発測定モードは測定終了と同時にPower Downになるので、もう一度モードを変更する
            if self._MAG_BIT == 14:  # output 14bit
                _writeData = 0x01
            else:  # output 16bit
                _writeData = 0x11
            self.i2c.writeReg8(self.AK8963, 0x0A, _writeData)
            time.sleep(0.01)

        elif self._MAG_MODE == self._MAG_MODE_SERIAL_1 or self._MAG_MODE == self._MAG_MODE_SERIAL_2:
            status = self.i2c.readReg8(self.AK8963, 0x02)
            if (status & 0x02) == 0x02:
                # if (status[0] & 0x02) == 0x02:
                # データオーバーランがあるので再度センシング
                self.i2c.readReg8(self.AK8963, 0x09)

        elif self._MAG_MODE == self._MAG_MODE_EX_TRIGER:
            # 未実装
            return

        elif self._MAG_MODE == self._MAG_MODE_POWERDOWN:
            raise Exception('003 Mag sensor power down')

        # ST1レジスタを確認してデータ読み出しが可能か確認する
        status = i2c.readReg8(AK8963, 0x02)
        while (status & 0x01) != 0x01:
            # while (status[0] & 0x01) != 0x01:
            # Wait until data ready state.
            # time.sleep(0.01)
            status = self.i2c.readReg8(self.AK8963, 0x02)

        # データ読み出し
        MAG_XOUT_L = i2c.readReg8(AK8963, 0x03)
        MAG_XOUT_H = i2c.readReg8(AK8963, 0x04)
        MAG_YOUT_L = i2c.readReg8(AK8963, 0x05)
        MAG_YOUT_H = i2c.readReg8(AK8963, 0x06)
        MAG_ZOUT_L = i2c.readReg8(AK8963, 0x07)
        MAG_ZOUT_H = i2c.readReg8(AK8963, 0x08)
        MAG_OF = i2c.readReg8(AK8963, 0x09)
        raw_x = self._u2s(MAG_XOUT_H << 8 | MAG_XOUT_L)
        raw_y = self._u2s(MAG_YOUT_H << 8 | MAG_YOUT_L)
        raw_z = self._u2s(MAG_ZOUT_H << 8 | MAG_ZOUT_L)
        st2 = MAG_OF
        # data    = self.i2c.readReg8(addrAK8963, 0x03 ,7)
        # raw_x    = u2s(data[1] << 8 | data[0])  # Lower bit is ahead.
        # raw_y    = u2s(data[3] << 8 | data[2])  # Lower bit is ahead.
        # raw_z    = u2s(data[5] << 8 | data[4])  # Lower bit is ahead.
        # st2     = data[6]

        # オーバーフローチェック
        if (st2 & 0x08) == 0x08:
            # オーバーフローのため正しい値が得られていない
            raise Exception('004 Mag sensor over flow')

        # μTへの変換
        if self._MAG_BIT == 16:  # output 16bit
            raw_x = raw_x * self._mag_coefficient_16
            raw_y = raw_y * self._mag_coefficient_16
            raw_z = raw_z * self._mag_coefficient_16
        else:  # output 14bit
            raw_x = raw_x * self._mag_coefficient_14
            raw_y = raw_y * self._mag_coefficient_14
            raw_z = raw_z * self._mag_coefficient_14

        # ---- offset by myself ----
        # raw_x -= 7.49733614
        # raw_y -= 57.46132325
        # raw_z -= -63.09610185
        # ---- offset by myself ----

        raw_x += self._offset_mag_x
        raw_y += self._offset_mag_y
        raw_z += self._offset_mag_z

        return raw_x, raw_y, raw_z

    # 加速度センサを較正する
    # 本当は緯度、高度、地形なども考慮する必要があるとは思うが、簡略で。
    # z軸方向に正しく重力がかかっており、重力以外の加速度が発生していない前提
    def _calib_accel(self, _count=1000):
        print("Accel calibration start")
        _sum = [0, 0, 0]

        # データのサンプルを取る
        for _ in range(_count):
            _data = self.get_accel()
            _sum[0] += _data[0]
            _sum[1] += _data[1]
            _sum[2] += _data[2]

        # 平均値をオフセットにする
        self._offset_accel_x = -1.0 * _sum[0] / _count
        self._offset_accel_y = -1.0 * _sum[1] / _count
        # self._offset_accel_z = -1.0 * _sum[2] / _count
        self._offset_accel_z = -1.0 * ((_sum[2] / _count) - 1.0)  # 重力分を差し引く

        # I want to register an offset value in a register. But I do not know the behavior, so I will put it on hold.
        print("Accel calibration complete")
        print(
            f'Accel error X: {self._offset_accel_x:.2f}  Y: {self._offset_accel_y:.2f}  Z: {self._offset_accel_z:.2f}')

    # ジャイロセンサを較正する
    # 各軸に回転が発生していない前提
    def _calib_gyro(self, _count=1000):
        print("Gyro calibration start")
        _sum = [0, 0, 0]

        # データのサンプルを取る
        for _i in range(_count):
            _data = self.get_gyro()
            _sum[0] += _data[0]
            _sum[1] += _data[1]
            _sum[2] += _data[2]

        # 平均値をオフセットにする
        self._offset_gyro_x = -1.0 * _sum[0] / _count
        self._offset_gyro_y = -1.0 * _sum[1] / _count
        self._offset_gyro_z = -1.0 * _sum[2] / _count

        # I want to register an offset value in a register. But I do not know the behavior, so I will put it on hold.
        print("Gyro calibration complete")
        print(f'Gyro error X: {self._offset_gyro_x:.2f}  Y: {self._offset_gyro_y:.2f}  Z: {self._offset_gyro_z:.2f}')

#    # 地磁気センサを較正する
#    # def _calib_mag(self, _count=3000):
#
#        def model(param, x):
#            return np.array([((xt[0] - param[0]) / param[3]) ** 2 + ((xt[1] - param[1]) / param[4]) ** 2 + (
#                    (xt[2] - param[2]) / param[5]) ** 2 for xt in x])
#
#        # 誤差関数
#        def residuals(param, x, y):
#            return y - model(param, x)
#
#        print("Mag calibration start")
#        print("Please keep the sensor turning around")
#        log = np.array([])
#
#        # データのサンプルを取る
#        for _ in range(_count):
#            _data = self.get_gyro()
#            np.append(log, [_data])
#
#        x_s = log.T
#        y_s = np.array([1.0] * len(x_s))
#        predicted_offset = np.array([5.0, 55, -65, 40, 40, 40])  # パラメータ初期値
#        res = least_squares(residuals, predicted_offset, args=(x_s, y_s))
#        offset_values = res.x
#
#        # 平均値をオフセットにする
#        self._offset_mag_x = -1.0 * offset_values[0]
#        self._offset_mag_y = -1.0 * offset_values[1]
#        self._offset_mag_z = -1.0 * offset_values[2]
#
#        # I want to register an offset value in a register. But I do not know the behavior, so I will put it on hold.
#        print("Mag calibration complete")
#        print(f'Mag error X: {self._offset_mag_x:.2f}  Y: {self._offset_mag_y:.2f}  Z: {self._offset_mag_z:.2f}')


import cv2
import numpy as np
import tensorflow as tf

from detect import find_face

def load_graph(file_name):
    with open(file_name, 'rb') as f:
        content = f.read()
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(content)
    with tf.Graph().as_default() as graph:
        tf.import_graph_def(graph_def, name='')
    return graph

def predict(result_tensor):
    prediction = tf.argmax(result_tensor, 1)
    return prediction

labels = dict()
with open('data/output_labels.txt', 'r') as f:
    temp = f.readlines()
    for i in range(len(temp)):
        labels.update({i: temp[i]})

graph = load_graph('data/output_graph.pb')
image_buffer_input = graph.get_tensor_by_name('Mul:0')
final_tensor = graph.get_tensor_by_name('final_result:0')
with tf.Session(graph=graph) as sess:
    video_camera = cv2.VideoCapture(0)
    while video_camera.isOpened():
        if (cv2.waitKey(1) & 0xFF) == ord('q'):
            break
        ret, frame = video_camera.read()
        cropped, bb = find_face(frame, 299)
        if bb:
            cropped = np.reshape(cropped, newshape=(1, 299, 299, 3))
            prediction = predict(final_tensor)
            predicted_label = sess.run(prediction, feed_dict={image_buffer_input: cropped})
            print labels[predicted_label[0]]
            x1, x2, y1, y2 = None, None, None, None
            for box in bb:
                x1 = box.left()
                y1 = box.top()
                x2 = box.right()
                y2 = box.bottom()
                cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                cv2.putText(frame, labels[predicted_label[0]],
                            org=(int(x2) - 50, int(y2) + 30), fontFace=cv2.FONT_HERSHEY_COMPLEX, color=0.2, thickness=1,
                            fontScale=0.4)
        cv2.imshow('Camera', frame)


from sentence_transformers import SentenceTransformer
from transformers import logging
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np

logging.set_verbosity_error()

class TextFilter:
    def __init__(self, threshold=0.45):
        self.model = SentenceTransformer('bert-base-nli-mean-tokens')
        self.threshold = threshold

    def run(self, path, text):
        storedEmbeds = np.load(path)
        textEmbeds = self.model.encode([text])
        probs = cosine_similarity(storedEmbeds, textEmbeds)
        return True if probs.item(0) > self.threshold else False

    def generateFilter(self, textList, path):
        textEmbeds = self.model.encode(textList)
        meanEmbeds = np.mean(textEmbeds, axis=0)[np.newaxis, :]
        np.save(path, meanEmbeds)
        return meanEmbeds

if __name__ == '__main__':
    sentences = [
        "I always wrote this series off as being a complete stink-fest because Jim Belushi was involved in it, and heavily. But then one day a tragic happenstance occurred. After a White Sox game ended I realized that the remote was all the way on the other side of the room somehow. Now I could have just gotten up and walked across the room to get the remote, or even to the TV to turn the channel. But then why not just get up and walk across the country to watch TV in another state? ""Nuts to that"", I said. So I decided to just hang tight on the couch and take whatever Fate had in store for me. What Fate had in store was an episode of this show, an episode about which I remember very little except that I had once again made a very broad, general sweeping blanket judgment based on zero objective or experiential evidence with nothing whatsoever to back my opinions up with, and once again I was completely right! This show is a total crud-pie! Belushi has all the comedic delivery of a hairy lighthouse foghorn. The women are physically attractive but too Stepford-is to elicit any real feeling from the viewer. There is absolutely no reason to stop yourself from running down to the local TV station with a can of gasoline and a flamethrower and sending every copy of this mutt howling back to hell.

Except..

Except for the wonderful comic sty lings of Larry Joe Campbell, America's Greatest Comic Character Actor. This guy plays Belushi's brother-in-law, Andy, and he is gold. How good is he really? Well, aside from being funny, his job is to make Belushi look good. That's like trying to make butt warts look good. But Campbell pulls it off with style. Someone should invent a Nobel Prize in Comic Buffoonery so he can win it every year. Without Larry Joe this show would consist of a slightly vacant looking Courtney Thorne-Smith smacking Belushi over the head with a frying pan while he alternately beats his chest and plays with the straw on the floor of his cage. 5 stars for Larry Joe Campbell designated Comedic Bacon because he improves the flavor of everything he's in!",
        "It's been about 14 years since Sharon Stone awarded viewers a leg-crossing that twisted many people's minds. And now, God knows why, she's in the game again. ""Basic Instinct 2"" is the sequel to the smash-hit erotica ""Basic Instinct"" featuring a sexy Stone and a vulnerable Michael Douglas. However, fans of the original might not even get close to this one, since ""Instinct 2"" is painful film-making, as the mediocre director Michael Caton-Jones assassinates the legacy of the first film.

The plot of the movie starts when a car explosion breaks in right at the beginning. Catherine Tramell (Sharon Stone, trying to look forcefully sexy) is a suspect and appears to be involved in the murder. A psychiatrist (a horrible David Morrisey) is appointed to examine her, but eventually falls for an intimate game of seduction.

And there it is, without no further explanations, the basic force that moves this ""Instinct"". Nothing much is explained and we have to sit through a sleazy, C-class erotic film. Sharon Stone stars in her first role where she is most of the time a turn-off. Part of it because of the amateurish writing, the careless direction, and terrifyingly low chemistry. The movie is full of vulgar dialogues and even more sexuality (a menage a trois scene was cut off so that this wouldn't be rated NC-17) than the first entrance in the series. ""Instinct"" is a compelling torture.

To top it off, everything that made the original film a guilty pleasure is not found anywhere in the film. The acting here is really bad. Sharon Stone has some highlights, but here, she gets extremely obnoxious. David Morrisey stars in the worst role of his life, and seems to never make more than two expressions in the movie- confused and aroused. ""Instinct 2"" is a horrible way to continue an otherwise original series, that managed to put in thriller with erotica extremely well. Paul Verhoeven, how I miss you....

""Basic Instinct 2"" never sounded like a good movie, and, indeed, it isn't. Some films should never get out of paper, and that is the feeling you get after watching this. Now, it is much easier to understand why Douglas and David Cronenberg dropped out, and why Sharon Stone was expecting a huge paycheck for this......-----3/10",
        "someone needed to make a car payment... this is truly awful... makes jean Claude's cyborg look like gone with the wind... this is an hour I wish I could sue to get back... luckily it produced severe somnolence... from which I fell asleep. how can actors of this caliber create this dog? I would rather spend the time watching algae grow on the side of a fish tank than partake of this wholly awful concoction of several genre. I now use the DVD as a coaster on my coffee table. $5.99 at walmart is far too much to spend on this movie... if you really have to have it, wait till they throw them out after they have carried them on the inventory for several years and are frustrated that they would not sell.

please for the love of god let this movie die of obscurity.",
        "The Guidelines state that a comment must contain a minimum of four lines. That is the only reason I am saying anything more about Tomcats. Because after all, my one line summary really says everything there is to say. There is absolutely NOTHING remotely entertaining in this film.",
        "This movie is a muddled mish-mash of clichés from recent cinema. There are some promising ideas in there, but while the director was clearly aiming to wind up with a hauntingly ambiguous film, what he ended up with was a confusing mess. Lead actor Daniel Wu does a fair job but with no central theme it seems as though he doesn't have much to work with. Furthermore, the movie is largely devoid of scares (although, in fairness, there are some creepy moments amid the drudgery).

*MILD SPOILERS*

We have the mysterious death of an estranged twin, diabolical librarians, ghostly love interests, identity confusion, death by savage monkeys, oedipal conflict, abusive stepfathers, sublimated homosexuality, and crime gang connections. The only real commonality these elements share seems to be that they cause the protagonist to express a vague sense of confusion and discontent.

Perhaps the most disappointing aspect to this film is that despite the brother's death by monkeys being strongly featured on the DVD cover, the act itself is never directly portrayed. Instead, director Julian Lee uses what appears to be stock footage of monkeys - not very scary.

*END SPOILERS*

Avoid this one. For an excellent psychological, ambiguous horror tale, check out the Korean film A Tale of Two Sisters (2003)."
        "Why did I waste 1.5 hours of my life watching this? Why was this film even made? Why am I even commenting on this film?

One reviewer said this film took patience to watch and it was n't for everybody. I cannot figure out who this movie is for. maybe after dropping a hit of acid, SOMEBODY, SOMEWHERE could watch this and make some sense out of it. It is incoherent, it isn't experimental, it's plain and simple garbage. The film follows no plot line whatsoever, just when you think you have something, well.....you don't.

I think the ending brought some finality to the film (no pun intended), the viewer gets a glimpse of what might have been going on. I don't think I put a spoiler in here, not that it would matter. This film is another must miss in the world of filmdom.",
        "This film takes you on one family's impossible journey, and makes you feel every step of their odyssey. Beautifully acted and photographed, heartbreakingly real. Its last line, with its wistful hope, is one of the more powerful in memory.",
        "The Russian space station 'Avna' with a crew of four Russians and two Americans is threatening to re-enter the Earth's atmosphere in a matter of days. Russia asks for NASA's help in rescuing the stranded crew and NASA scrambles the space shuttle Atlantis. The NSA also have an interest in the 'Prometheus', a prototype microwave power source being tested aboard 'Avna' and organise for one of their men to be placed on the mission.

That's the plot. Onto less important things. The space station and the shuttle are the same, blatantly obvious models used in 'Fallout', 'Memorial Day' and 'Dark Breed' (and a handful of other films, I suspect). The model effects are so obvious throughout the entire movie and make the film look very 1960s. The sets are a little better but are far too '80s for what is supposedly a brand new station built by an American company (which later comes in as part of a conspiracy to destroy 'Avna' and the 'Prometheus' and claim the insurance. The script has a few good moments (including Yuri's farewell and the little spiel at the end) but is otherwise fairly bland and sub-standard. The acting is okay; the only real standout performance comes from Alex Veadov who offers up some of the film's better dialogue. Michael Dudikoff is, surprisingly, one of the best parts about this film. Ice-T is Ice-T. 'Nuff said. The film offers a few surprises, though, that I don't wish to spoil.

Certainly one of the better low-grade, contemporary-set sci-fi films of the last six years, but not the best. The film is watchable but the special effects and plot will probably put a lot of viewers off. Rent the other 'Stranded' sci-fi film instead.",
        "the more i think about it, there was nothing redeeming about this

movie. i saw it 9 months ago, so my memory might have made it

worse than it was, but i do know it was at least as bad as a 4 out of

10.

after seeing the movie, i met the director. he seemed so clueless

as to what he was doing or what he had done, and as far as i

could tell, he didn't care for the film either. even he agreed that he

didn't really know what he was doing, and he was forced to do

certain things because it was filmed digitally.

i felt that the movie was trying to hard to fit in to the formula that it

built for itself: ""9 people all have to be connected in some way. how

can we get from point 'A' to point 'B'"" so in order get from the

prostitute we see in the start and back to her at the end they 10

minutes on each character's relationship to another person. it

makes one feel choked by the 2 demensional, badly drawn

characters.

I just remembered the one redeeming part of the movie... Steve

Bouchemi there is one scene where he is amazing. that's it. as i

say... 4 out of 10.",
        "This is very dated, but that's part of the charm with this 1933 movie. You can say the same for most Pre-Code films; they're just different, and usually in an interesting way.

It was the short running time, the great acting of Spencer Tracy and the beautiful face and sweetness of Loretta Young's character which kept me watching and enjoying this stagy-but-intriguing film.

You'd be hard-pressed to find a nicer girl than ""Trinna,"" played by the 20-year-old Young who was already into making her 50th movie! (She started acting as a small child. That, and the fact they made movies quickly back in the old days.) The camera, although in soft focus throughout much of the film, zoomed in on Loretta's face and eyes many times and I was mesmerized by her beauty.

Playing a crotchety man with a cynical outlook on life, Tracy's ""Bill"" slowly transformed into a loving man, thanks to Trinna. Spencer delivered his lines here with such naturalness that you hardly knew he was acting.

Although they have small roles, supporting actors Walter Connolly, Marjorie Rambeau, Arthur Hohl and Glenda Farrell leave lasting impressions long after viewing this 75-minute film. I was particularly fascinated with Connolly's role as the minister/father figure of the camp.

The story is a little far-fetched but - hey - that's the movies. This story is about two lonely Great Depression victims trying to survive in a ""Hooverville""-type camp and it winds up to be a very touching tale.",
        "A powerful adaptation of the best-selling book and the smash Broadway play about the lives of Bessie and Sadie Delany, two ""colored"" sisters who lived past the age of 100. Wonderfully played in their old age by Ruby Dee and Diahann Carroll, respectively, they tell their story in flashbacks to Amy Hill Hearth (played by Amy Madigan), a white New York Times reporter. The flashback and present-day scenes don't have as much inspirational value in them as in the book, but really are powerful. However, certain aspects of the sisters' lives, such as the inter-racial background of their mother and the reasons behind their father's stern personality are not presented clearly. You need to read the book to fully understand these things. Which is just as well, because the book's just as great! Aside from those flaws, it's wonderfully done and performed, especially by Dee and Carroll, and a very powerful and educational movie.",
        "This movie's origins are a mystery to me, as I only know as much as IMDB did before I rented it. I assume that before ""Starship Troopers"", ""Killshot"" was one of the countless unaired pilots that never made it to network, cable, or otherwise. The new title of ""Kill Shot"" is comically thrown into the opening sequence, the first of many quick clues that this was not ever intended for the cinema. The quick cuts, cheesy ""Melrose Place"" music, and short 2-second close-up candid shots of the main actors let you know what you're in for.

And I don't mind at all. I rented this movie seeing the repackaging that puts Casper Van Dien and Denise Richards on the cover in front of a volleyball net thinking it would be funny to see them in a movie besides the SciFi travesty of Starship Troopers (an excellent book, in my opinion, not so hot a movie - but that's another review). After looking it up on IMDB, my roommate and I surmised that the pilot was dragged up after the apparent success of Troopers and Richards own career (see Bond-Girl and Wild Things references here). They threw in a sex scene involving a minor character to reach the coveted R-rated status - coveted in suspense Video Rental sections, that is. In any event, they should have left it unrated if you're trying to sell it in the suspense/softcore porn section.

All in all, it's entertaining. I hate to spoil the fun of telling you it's a TV pilot, though. That was the biggest pull while watching it - when you expect a cinematic movie and get a TV show, the differences between them make themselves more clear than usual.

Would I rent it again? No. Would I watch this TV show? Well, why not - it's better than Baywatch. And their meager attempts at hitting all demographics would have done well back in the mid 90s. Token black guy (who's gay to avoid the TV taboo of inter-racial dating), token Asian (Japanese, I assume from the name Koji) more adept at science and computers than talking to women, beautiful, intelligent Latina pre-med student who has everything going for her except her family's bank account - this show probably would have done ok.

But as a movie it just cracks me.

I gave it a 7 out of 10, considering what it was and what it was forced to become. It made for a very enjoyable evening, and that's all I ask of rentals.",
        "I really enjoyed the first film and when it turned up again, without thinking, or checking, I took a family of friends to see it. I was ashamed that I had enthused so much about it to them.

Disney processed the original film just like the human body processes a delicious meal - takes in something good and turns out ... well, you know. And by having a dark-skinned person as the FBI man, the results of fingerprinting the informant were subdued.

Taken as an isolated film, I suppose it is not too bad if one likes that weird sort of thing, but when one has read the book or seen the first film - horrible!"
    ]

    text = "I had seen 'Kalifornia' before (must be about 10 years ago) and I still remember to be very impressed by it. That's why I wanted to see it again and all I can say is that it still hasn't lost its power, even though I'm used to a lot more when it comes to movies than that I was ten years ago.

'Kalifornia' tells the tale of the writer Brian Kessler and his girlfriend Carrie Laughlin, a photographer, who want to move to California. But instead of stepping on a plain and flying right to the state where they say it never rains, they choose to make a trip by car. He wants to write a book about America's most famous serial killers and she will make the matching pictures. But because their car uses an enormous amount of petrol, they decide to take another couple with them, so they can spread the costs of the trip. Only one couple has answered the add, so they will automatically be the lucky ones. But they haven't met each other yet and when seeing the other couple for the first time, when their trip has already started, Carrie is shocked. Without wanting to be prejudiced, she can only conclude that Early Grayce and Adele Corners are poor white trailer park trash. She definitely doesn't want them in her car, but Brian doesn't really mind to take them with them and decides to stop and pick them up anyway. At first the couple doesn't seem to be that bad after all, but gradually Early Grayce changes from a trashy hillbilly into a remorseless murderer...

Not only is the story very impressive, so is the acting from our four leads. Brad Pitt is incredible as Early Grayce. His performance in this movie may well be his best ever. The same for Juliette Lewis. She plays the childish and naive girlfriend that doesn't want to hear a bad word about her Early and does that really very well. But David Duchovny and Michelle Forbes are a surprise as well. They both did a very good job and I really wonder why we never heard anything from Forbes again since this movie, because she really proves to have a lot of talent.

Overall this is a very good and impressive psychological thriller with a very powerful story, but because of the graphic violence, I can imagine that it may not be to everybody's taste (although I don't really see another way how to portray a serial killer in a believable way). Personally I really liked this movie a lot and the violence never bothered me (it's a part of the story that's too important to be left out). I reward this movie with an 8/10."
    tf = TextFilter()
    tf.generateFilter(sentences, './filter.npy')
    print(tf.run('./filter.npy', text))


from enum import IntEnum
from pathlib import Path
from typing import List, Union

import cv2
import numpy as np
from xlib.image import ImageProcessor
from xlib.net import ThreadFileDownloader

from xlib import tf

class DFMModelInfo:
    """
    Model for face swapping.

    """
    def __init__(self, celeb_type, device):
        celeb_name = str(celeb_type).split('.')[-1]
        model_path = self._model_path = Path(__file__).parent / 'CELEB_MODEL' / f'{celeb_name}.onnx'

        self._dl_error = None
        self._dl = None
        if not model_path.exists():
            self._dl = ThreadFileDownloader(rf'https://github.com/iperov/DeepFaceLive/releases/download/CELEB_MODEL/{celeb_name}.onnx', savepath=model_path)

        self._device = device
        self._sess = None

    @staticmethod
    def get_available_devices(celeb_type):
        return tf.get_available_devices_info()

    def get_download_error(self) -> Union[str,None]:
        """
        returns download error or None if no error.
        """
        if self._dl_error is not None:
            return 'DFMModelInfo download error: ' + self._dl_error
        return None

    def update_download_progress(self) -> float:
        """
        returns [0.0..100.0] where 100.0 mean the model is ready to convert.
        """
        if self._dl_error is not None:
            progress = 0.0
        elif self._dl is None:
            progress = 100.0
        else:
            err = self._dl_error = self._dl.get_error()
            if err is not None:
                progress = 0.0
            else:
                progress = self._dl.get_progress()

        if progress == 100.0 and self._sess is None:
            self._dl = None

            model_path = r'D:\DevelopPPP\projects\DeepFaceLive\github_project\xlib\model_hub\tf\TOM_CRUISE.pb'
            self._sess = tf.TFInferenceSession(model_path,
                                 in_tensor_names=['in_face:0', 'morph_value:0'],
                                 out_tensor_names=['out_face_mask:0','out_celeb_face:0','out_celeb_face_mask:0'],
                                 device_info=self._device)

        return progress

    def convert(self, img, morph_factor=0.75):
        """
         img    np.ndarray  HW,HWC,NHWC uint8,float32

        returns

         img        NHW3  same dtype as img
         celeb_mask NHW1  same dtype as img
         face_mask  NHW1  same dtype as img

        convert the face
        if the model is not ready, it will return black images
        """

        ip = ImageProcessor(img)

        N,H,W,C = ip.get_dims()
        dtype = ip.get_dtype()

        progress = self.update_download_progress()
        if progress == 100.0:

            img = ip.resize( (224,224) ).ch(3).to_ufloat32().get_image('NHWC')

            #out_face_mask, out_celeb, out_celeb_mask = self._sess.run(None, {self._input_name: img})

            out_face_mask, out_celeb, out_celeb_mask = self._sess.run([img, [morph_factor] ])

            out_celeb      = ImageProcessor(out_celeb).resize((W,H)).ch(3).to_dtype(dtype).get_image('NHWC')
            out_celeb_mask = ImageProcessor(out_celeb_mask).resize((W,H)).ch(1).to_dtype(dtype).get_image('NHWC')
            out_face_mask  = ImageProcessor(out_face_mask).resize((W,H)).ch(1).to_dtype(dtype).get_image('NHWC')
        else:
            out_celeb = np.zeros( (N,H,W,3), dtype=dtype )
            out_celeb_mask = ImageProcessor(np.full( (N,H,W,1), 255, dtype=np.uint8 )).to_dtype(dtype).get_image('NHWC')
            out_face_mask = ImageProcessor(np.full( (N,H,W,1), 255, dtype=np.uint8 )).to_dtype(dtype).get_image('NHWC')

        return out_celeb, out_celeb_mask, out_face_mask

    class CelebType(IntEnum):
        TOM_CRUISE = 0
        VLADIMIR_PUTIN = 1

    CelebTypeNames = ['Tom Cruise',
                      'Vladimir Putin']