"""Create recognition crops from train.json (segmentation model).

--data_dir C:/temp/data --output_dir C:/temp/out -c 1.2 -p 0.95 -a True -f True -r True
чтоб сделать датасет для ландмарков
--data_dir C:/temp/data --output_dir C:/temp/recogn -c 1.2 --make_square True  --randomize=True
"""

import json
import os
import glob
import torch
from argparse import ArgumentParser

import cv2
import numpy as np
import tqdm

from segmentation.models import get_segmentation_model
from utils.image_transforms import get_crop

RESIZE_DIMS = (122, 28)
from utils.image_transforms import get_crop, apply_mask, get_sorted_polyline, \
    parallelaze_polygon, plot_circles, affine_transform


def parse_arguments():
    parser = ArgumentParser()
    parser.add_argument("--data_dir", help="Path to dir containing images.")
    parser.add_argument("-o", "--output_dir", dest="output_dir", default="seg_out", help="Path to output dir'.")
    parser.add_argument("-l", "--load", dest="load", help="segmantation model name")
    parser.add_argument("-c", "--crop_scale", dest="crop_scale", default=1.2, type=float, help="scale bounding box")
    parser.add_argument("-p", "--polygon_scale", dest="polygon_scale", default=1, type=float, help="scale polygon ")


    return parser.parse_args()


def main(args):
    os.makedirs(args.output_dir, exist_ok=True)
    test_images = glob.glob(os.path.join(args.data_dir, '*.jpg')) + \
                  glob.glob(os.path.join(args.data_dir, '*.bmp'))

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    segmentation_model = get_segmentation_model()
    with open(args.load, "rb") as fp:
        print(f"load {args.load}")
        state_dict = torch.load(fp, map_location="cpu")

    segmentation_model.load_state_dict(state_dict)
    segmentation_model.to(device)
    segmentation_model.eval()
    with torch.no_grad():
        for img_file_name in tqdm.tqdm(test_images, total=len(test_images)):
            image = cv2.imread(os.path.join(args.data_dir, img_file_name))
            if image is None:
                print(f"wrong file {img_file_name}")
                continue
            image_n = image.astype(np.float32) / 255
            t_image = torch.as_tensor(image_n.transpose(2, 0, 1)).to(device)
            predict = segmentation_model([t_image])
            path_list = os.path.normpath(img_file_name).lstrip(os.path.sep).split(os.path.sep)

            # path_list = image_filename.split(os.sep)
            image_base, ext = os.path.splitext(path_list[-1])
            mask_result = np.zeros(image_n.shape[0:2])

            for i, (class_num, score, mask) in enumerate(zip(predict[0]["labels"], predict[0]["scores"],
                                                                   predict[0]["masks"])):
                if class_num == 1 and score > 0.4:
                    mask = mask.cpu().squeeze().numpy()
                    mask_result += mask


            mask_result[mask_result > 0.3] = 1
            mask_result = mask_result.astype(np.uint8)

            contours, _ = cv2.findContours(mask_result.astype(np.uint8), 1, 1)
            for i, contour in enumerate(contours):
                x, y, w, h = cv2.boundingRect(contour)
                crop_image, _ = get_crop(image,np.array([[x,y],[x + w, y+  h]], dtype=np.int32),
                                         args.crop_scale, make_square=True)

                fout = os.path.join(args.output_dir, f"{image_base}_{i}{ext}")
                cv2.imwrite(fout, crop_image)





if __name__ == "__main__":
    main(parse_arguments())

"""Create recognition crops from train.json (segmentation model).

--data_dir C:/temp/data --output_dir C:/temp/out -c 1.2 -p 0.95 -a True -f True -r True
чтоб сделать датасет для ландмарков
--data_dir C:/temp/data --output_dir C:/temp/recogn -c 1.2 --make_square True  --randomize=True
"""

import json
import os
import glob
import torch
from argparse import ArgumentParser

import cv2
import numpy as np
import tqdm

from segmentation.models import get_segmentation_model
from utils.image_transforms import get_crop

RESIZE_DIMS = (122, 28)
from utils.image_transforms import get_crop, apply_mask, get_sorted_polyline, \
    parallelaze_polygon, plot_circles, affine_transform


def parse_arguments():
    parser = ArgumentParser()
    parser.add_argument("--data_dir", help="Path to dir containing images.")
    parser.add_argument("-o", "--output_dir", dest="output_dir", default="seg_out", help="Path to output dir'.")
    parser.add_argument("-l", "--load", dest="load", help="segmantation model name")
    parser.add_argument("-c", "--crop_scale", dest="crop_scale", default=1.2, type=float, help="scale bounding box")
    parser.add_argument("-p", "--polygon_scale", dest="polygon_scale", default=1, type=float, help="scale polygon ")


    return parser.parse_args()


def main(args):
    os.makedirs(args.output_dir, exist_ok=True)
    test_images = glob.glob(os.path.join(args.data_dir, '*.jpg')) + \
                  glob.glob(os.path.join(args.data_dir, '*.bmp'))

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    segmentation_model = get_segmentation_model()
    with open(args.load, "rb") as fp:
        print(f"load {args.load}")
        state_dict = torch.load(fp, map_location="cpu")

    segmentation_model.load_state_dict(state_dict)
    segmentation_model.to(device)
    segmentation_model.eval()
    with torch.no_grad():
        for img_file_name in tqdm.tqdm(test_images, total=len(test_images)):
            image = cv2.imread(os.path.join(args.data_dir, img_file_name))
            if image is None:
                print(f"wrong file {img_file_name}")
                continue
            image_n = image.astype(np.float32) / 255
            t_image = torch.as_tensor(image_n.transpose(2, 0, 1)).to(device)
            predict = segmentation_model([t_image])
            path_list = os.path.normpath(img_file_name).lstrip(os.path.sep).split(os.path.sep)

            # path_list = image_filename.split(os.sep)
            image_base, ext = os.path.splitext(path_list[-1])
            mask_result = np.zeros(image_n.shape[0:2])

            for i, (class_num, score, mask) in enumerate(zip(predict[0]["labels"], predict[0]["scores"],
                                                                   predict[0]["masks"])):
                if class_num == 1 and score > 0.4:
                    mask = mask.cpu().squeeze().numpy()
                    mask_result += mask


            mask_result[mask_result > 0.3] = 1
            mask_result = mask_result.astype(np.uint8)

            contours, _ = cv2.findContours(mask_result.astype(np.uint8), 1, 1)
            for i, contour in enumerate(contours):
                x, y, w, h = cv2.boundingRect(contour)
                crop_image, _ = get_crop(image,np.array([[x,y],[x + w, y+  h]], dtype=np.int32),
                                         args.crop_scale, make_square=True)

                fout = os.path.join(args.output_dir, f"{image_base}_{i}{ext}")
                cv2.imwrite(fout, crop_image)





if __name__ == "__main__":
    main(parse_arguments())

Create recognition crops from train.json (segmentation model).  --data_dir C:/temp/data --output_dir C:/temp/out -c 1.2 -p 0.95 -a True -f True -r True чтоб сделать датасет для ландмарков --data_dir C:/temp/data --output_dir C:/temp/recogn -c 1.2 --make_square True  --randomize=True # path_list = image_filename.split(os.sep)