import sagemaker
input_s3_uri = sagemaker.session.Session().upload_data(path='./conveyor_image/', bucket=sagemaker.session.Session().default_bucket(), key_prefix='conveyor_image')

import sagemaker
import os, json
TARGET_DIR = 'annotation_result'
!mkdir -p {TARGET_DIR}
sagemaker.session.Session().download_data(TARGET_DIR,bucket=sagemaker.session.Session().default_bucket(),key_prefix='conveyor_image/conveyor-image/')
ITE1_WORKER_RESPONSE = './annotation_result/annotations/worker-response/iteration-1/'
max_time = 0
userid = ""
for directory in sorted(os.listdir(ITE1_WORKER_RESPONSE)):
    tmp_dir = os.path.join(ITE1_WORKER_RESPONSE,directory)
    tmp_json_file = os.path.join(tmp_dir,os.listdir(tmp_dir)[0])
    with open(tmp_json_file) as f:
        tmp_json = json.loads(f.read())['answers'][0]
        if max_time < float(tmp_json['timeSpentInSeconds']):
            max_time = float(tmp_json['timeSpentInSeconds'])
            userid = tmp_json['workerMetadata']['identityData']['sub']
print(max_time,userid)

sagemaker.session.Session().download_data('./manifest/',key_prefix=f'conveyor_image/conveyor-image/manifests/output/output.manifest',bucket=bucket)

# 追記 : ベルトコンベアの画像を格納したディレクトリに定数を変更
TRAIN_RAWIMAGE_DIR = 'conveyor_image'

# ラベリング結果をテキストとして読み込む
with open('./manifest/output.manifest','r') as f:
    manifest_line_list = f.readlines()

# クロップした結果のきのこの山やたけのこの里の位置情報を格納する辞書 
annotation_dict = {
    'images':[],
    'annotations':[]
}

# クロップサイズの定数
IMAGE_SIZE_TUPLE=(640,480) # 変更 : 画像サイズ変更

# クロップした画像のファイル名に使う一意なシーケンス番号
IMAGE_ID = 0

# クロップした画像の保存先
OUTPUT_DIR = './train_random_crop_images/'

# (re-run用の削除コマンド)
!rm -rf {OUTPUT_DIR}*.png

# ラベリング結果の行数分ループする
# ラベリング結果は 1 行につき 1 画像格納される
for manifest_line in manifest_line_list:
    # 画像のラベリング結果の読み込み
    manifest_dict = json.loads(manifest_line)
    # 画像のファイル名取得(ラベリング結果に格納されている)
    filename = manifest_dict['source-ref'].split('/')[-1]
    # 元画像のサイズを取得(ラベリング結果に格納されている)
    image_size_tuple=(manifest_dict['conveyor-image']['image_size'][0]['width'],manifest_dict['conveyor-image']['image_size'][0]['height'])
    # PIL で画像を開く
    raw_img = Image.open(os.path.join(TRAIN_RAWIMAGE_DIR,filename))
    # 変更 : クロップ回数は1回(=ランダムクロップしない)
    for i in range(1):
#       ループさせる必要がないので無視
#         # ループするかどうかのフラグ(画像にきのこの山やたけのこの里が 2 枚未満だったらクロップをやりなおし)
#         loop = True
#         while loop:
        # クロップを行う左上の座標を設定
        rand_x = 0
        rand_y = 0
        # クロップする
        crop_img = raw_img.crop((
            rand_x,
            rand_y,
            rand_x + IMAGE_SIZE_TUPLE[0],
            rand_y + IMAGE_SIZE_TUPLE[1]
        ))
        # クロップ後のきのこの山やたけのこの里の位置を格納するリスト
        annotation_list = []
        # 元画像のラベリング結果をループ
        for annotation in manifest_dict['conveyor-image']['annotations']:
            # クロップした後のきのこの山やたけのこの里の座標に補正
            left = annotation['left'] - rand_x
            top = annotation['top'] - rand_y
            right = annotation['left'] + annotation['width'] - rand_x
            bottom = annotation['top'] + annotation['height'] - rand_y
                # 変更 : きのこの山やたけのこの里があるかどうかを判定する必要がないのでコメントアウト
#                 judge,(left,top,right,bottom) = fix_bbox(left,top,right,bottom,IMAGE_SIZE_TUPLE[0],IMAGE_SIZE_TUPLE[1])
#                 if judge:
#                     # きのこの山やたけのこの里があったら位置とラベルを追加
            # 変更 : 判定がなくなったのでインデントを前へ
            annotation_list.append(
                {
                    'bbox':[left,top,right,bottom],
                    'category_id':annotation['class_id']
                }
            )
            # 変更 : クロップやり直しは不要なのでコメントアウト
#             # きのこの山やたけのこの里と数が2未満だったらクロップやり直し
#             if len(annotation_list):
#                 loop = False
        
        # クロップしたら画像を保存する
        save_file_name = f'{str(IMAGE_ID).zfill(5)}_{str(i).zfill(5)}_{filename}'.replace('jpg','png')
        crop_img.save(os.path.join(OUTPUT_DIR,save_file_name)) 
        # 補正済ラベリング結果を出力用辞書に格納
        annotation_dict['images'].append(
            {
                'file_name' : save_file_name,
                'height' : IMAGE_SIZE_TUPLE[1],
                'width' : IMAGE_SIZE_TUPLE[0],
                'id' : IMAGE_ID
            }
        )
        for annotation in annotation_list:                  
            annotation_dict['annotations'].append(
                {
                    'image_id': IMAGE_ID,
                    'bbox':annotation['bbox'],
                    'category_id':annotation['category_id']
                }
            )
        IMAGE_ID += 1

# ランダムクロップ補正後のラベリング結果を出力
with open('annotations.json','wt') as f:
    f.write(json.dumps(annotation_dict))        

# 出力したディレクトリを prefix として使う
prefix = OUTPUT_DIR[2:-1]
# re-run 用の削除コマンド
!aws s3 rm s3://{sagemaker.session.Session().default_bucket()}/{prefix} --recursive
# ランダムクロップした画像をアップロード    
image_s3_uri = sagemaker.session.Session().upload_data(OUTPUT_DIR,key_prefix=f'{prefix}/images')
# ラベリング結果をアップロード
annotatione_s3_uri = sagemaker.session.Session().upload_data('./annotations.json',key_prefix=prefix,bucket=sagemaker.session.Session().default_bucket())
# Fine-Tune で使う URI を出力
paste_str = image_s3_uri.replace('/images','')
print(f"paste string to S3 bucket address：{paste_str}")

from sagemaker import image_uris, model_uris, script_uris

model_id, model_version = "mxnet-od-ssd-512-mobilenet1-0-coco", "1.1.0"
training_instance_type = "ml.g4dn.xlarge"
instance_count = 1

# Retrieve the JumpStart base model S3 URI
base_model_uri = model_uris.retrieve(
    model_id=model_id, model_version=model_version, model_scope="training"
)

# Retrieve the training script and Docker image
training_script_uri = script_uris.retrieve(
    model_id=model_id, model_version=model_version, script_scope="training"
)
training_image_uri = image_uris.retrieve(
    region=None,
    framework=None,
    image_scope="training",
    model_id=model_id,
    model_version=model_version,
    instance_type=training_instance_type,
)

from sagemaker.estimator import Estimator
from sagemaker.session import Session
from sagemaker import hyperparameters

# URI of your training dataset
training_dataset_s3_path = paste_str

# Get the default JumpStart hyperparameters
hps = hyperparameters.retrieve_default(
    model_id=model_id,
    model_version=model_version,
)
# [Optional] Override default hyperparameters with custom values
hps['epochs'] = '40'
hps['adam-learning-rate'] = '0.0001'


# Create your SageMaker Estimator instance
estimator = Estimator(
    image_uri=training_image_uri,
    source_dir=training_script_uri,
    model_uri=base_model_uri,
    entry_point="transfer_learning.py",
    role=Session().get_caller_identity_arn(),
    hyperparameters=hps,
    instance_count=instance_count,
    instance_type=training_instance_type,
    enable_network_isolation=True,
)

# Specify the S3 location of training data for the training channel
estimator.fit({"training": training_dataset_s3_path})

import tarfile
prefix = estimator.latest_training_job.describe()['ModelArtifacts']['S3ModelArtifacts'].replace(f's3://{sagemaker.session.Session().default_bucket()}/','')
sagemaker.session.Session().download_data('./', sagemaker.session.Session().default_bucket(), key_prefix=prefix)
with tarfile.open('./model.tar.gz', 'r') as f:
    f.extractall()

git clone https://github.com/JetsonHacksNano/installSwapfile
cd installSwapfile
./installSwapfile.sh

$ uname -a
Linux nvidia 4.9.140-tegra #1 SMP PREEMPT Fri Oct 16 12:32:46 PDT 2020 aarch64 aarch64 aarch64 GNU/Linux

$ python3 --version
Python 3.7.5

$ python3
>>> import mxnet, gluoncv, cv2
>>> mxnet.__version__
'1.8.0'
>>> gluoncv.__version__
'0.11.0'
>>> cv2.__version__
'4.5.5'

$ export AWS_ACCESS_KEY_ID=
$ export AWS_SECRET_ACCESS_KEY=
$ export AWS_DEFAULT_REGION=ap-northeast-1

$ sudo apt install -y curl default-jdk

$ sudo -E java -Droot="/greengrass/v2" -Dlog.store=FILE \
 -jar ./GreengrassInstaller/lib/Greengrass.jar \
 --aws-region $AWS_DEFAULT_REGION \
 --thing-name bf202205_Nano \
 --thing-group-name bf202205_Nano_Group \
 --thing-policy-name bf202205_Nano_IoTPolicy \
 --tes-role-name bf202205_Nano_TESRole \
 --tes-role-alias-name bf202205_Nano_TESRoleAlias \
 --component-default-user ggc_user:ggc_group \
 --provision true \
 --setup-system-service true

# AWS IoT Greengrass のコンポーネントを保存するバケットを作成
component_bucket_name = f"bf202205-ggcomponent-{account_id}"
component_bucket_name

# AWS IoT Greengrass を Jetson Nano 上にインストールした時の、
# Thing 名(--thing-name)、グループ名( --thing-group-name)、 Role 名(--tes-role-name)を変数に指定します
GG_THING_NAME="bf202205_Nano"
GG_GROUP_NAME="bf202205_Nano_Group"
GG_TES_ROLE="bf202205_Nano_TESRole"