#!./venv/bin/python3.7

import argparse
import sys
import os
import cv2
import numpy as np
from src import lib


parser = argparse.ArgumentParser()
parser.add_argument('-i', '--input', help="Input image path", required=True)
parser.add_argument('-o', '--outdir', help="Output directory", required=True)
parser.add_argument(
    '-m', '--model', help="Model to predict", required=True)
parser.add_argument('-it', '--interactive',
                    help='Interactive mode', action='store_true')
parser.add_argument('-a', '--alpha')


args = parser.parse_args()
alpha = args.alpha


img_path = args.input
if not os.path.isfile(img_path):
    raise RuntimeError(f"Could not find image {img_path}")

output_dir = args.outdir
if not os.path.isdir(output_dir):
    raise RuntimeError(f"Could not find dir {output_dir}")

model_path = args.model
if not os.path.isfile(model_path):
    raise RuntimeError(f"Could not find model {model_path}")

original_img = cv2.imread(img_path)
img = cv2.bitwise_not(cv2.cvtColor(original_img, cv2.COLOR_BGR2GRAY))


model_name = lib.extract_name(model_path)
name, extension = lib.split_name(img_path)
out_path = f'{output_dir}/{name}_{model_name}_described'

if not args.interactive:
    img = lib.binarize(img)
    img = lib.morph_open(img, (3, 3))
    img = lib.drop_contacting_objects(img)
else:
    img, alpha = lib.interactive_edit_image(img)

cnts, hierarchy = cv2.findContours(
    img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

h, w = img.shape
model = lib.load_trained_model(model_path)

labels = {}


alpha = 0.0001 if alpha is None else alpha
AR_THRESHOLD = w * h * alpha


nn = len(cnts)
i = 1
for c in cnts:
    print(f'Predicting {i}/{nn}')
    i += 1

    ar = cv2.contourArea(c)
    if ar < AR_THRESHOLD:
        print('Skipped, too small probe')
        continue

    res = lib.create_probe_from_contour(img, c, original_img=original_img)
    if res is None:
        continue

    probe, x, y = res
    probe = lib.reshape_probe(probe)

    label = lib.predict_single_image(model, probe)
    text_label = lib.NUM_TO_LABEL[str(label)]

    if text_label in labels:
        labels[text_label] += 1
    else:
        labels[text_label] = 1

    lib.mark_object(original_img, c, text_label, (x, y))
    print("Predicted and marked")

print(labels)
print(out_path)

cv2.imwrite(out_path + f'.{extension}', original_img)

#!./venv/bin/python3.7

import argparse
import sys
import os
import cv2
import numpy as np
from src import lib


parser = argparse.ArgumentParser()
parser.add_argument('-i', '--input', help="Input image path", required=True)
parser.add_argument('-o', '--outdir', help="Output directory", required=True)
parser.add_argument(
    '-m', '--model', help="Model to predict", required=True)
parser.add_argument('-it', '--interactive',
                    help='Interactive mode', action='store_true')
parser.add_argument('-a', '--alpha')


args = parser.parse_args()
alpha = args.alpha


img_path = args.input
if not os.path.isfile(img_path):
    raise RuntimeError(f"Could not find image {img_path}")

output_dir = args.outdir
if not os.path.isdir(output_dir):
    raise RuntimeError(f"Could not find dir {output_dir}")

model_path = args.model
if not os.path.isfile(model_path):
    raise RuntimeError(f"Could not find model {model_path}")

original_img = cv2.imread(img_path)
img = cv2.bitwise_not(cv2.cvtColor(original_img, cv2.COLOR_BGR2GRAY))


model_name = lib.extract_name(model_path)
name, extension = lib.split_name(img_path)
out_path = f'{output_dir}/{name}_{model_name}_described'

if not args.interactive:
    img = lib.binarize(img)
    img = lib.morph_open(img, (3, 3))
    img = lib.drop_contacting_objects(img)
else:
    img, alpha = lib.interactive_edit_image(img)

cnts, hierarchy = cv2.findContours(
    img, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

h, w = img.shape
model = lib.load_trained_model(model_path)

labels = {}


alpha = 0.0001 if alpha is None else alpha
AR_THRESHOLD = w * h * alpha


nn = len(cnts)
i = 1
for c in cnts:
    print(f'Predicting {i}/{nn}')
    i += 1

    ar = cv2.contourArea(c)
    if ar < AR_THRESHOLD:
        print('Skipped, too small probe')
        continue

    res = lib.create_probe_from_contour(img, c, original_img=original_img)
    if res is None:
        continue

    probe, x, y = res
    probe = lib.reshape_probe(probe)

    label = lib.predict_single_image(model, probe)
    text_label = lib.NUM_TO_LABEL[str(label)]

    if text_label in labels:
        labels[text_label] += 1
    else:
        labels[text_label] = 1

    lib.mark_object(original_img, c, text_label, (x, y))
    print("Predicted and marked")

print(labels)
print(out_path)

cv2.imwrite(out_path + f'.{extension}', original_img)

#!./venv/bin/python3.7