# 첫 번째
# 2023년 1월 30일 mmdetection 오전 수업
# 미니프로젝트에서 wine 라벨 이미지를 사용했던거를 coco데이터를 사용해준다.
# 필요한 패키지를 추가해주었다.
import mmcv
import torch

from mmdet.apis import init_detector, inference_detector
from mmdet.datasets.builder import DATASETS
from mmdet.datasets.coco import CocoDataset
from mmdet.datasets import build_dataset
from mmdet.models import build_detector
from mmdet.apis import set_random_seed
from mmdet.apis import train_detector
from mmdet.utils import collect_env, get_root_logger, setup_multi_processes
from mmcv.runner import get_dist_info, init_dist
from mmcv import Config

# 두 번째
# Dataset regiset 데이터 셋에 등록이 필요하다
# 포맷 자체가 CoCo데이터 셋을 가져온다.
@DATASETS.register_module(force=True)
class WineLabelDataset(CocoDataset): # 상속을 받게 된다.
    # CoCo데이터 셋에는 클래스가 1,000개인데, 우리가 원하는 클래스를 지정해준다
    CLASSES = ('wine-labels', 'AlcoholPercentage', 'Appellation AOC DOC AVARegion', 'Appellation QualityLevel', 'CountryCountry',
            'Distinct Logo', 'Established YearYear', 'Maker-Name', 'Organic', 'Sustainable', 'Sweetness-Brut-SecSweetness-Brut-Sec',
            'TypeWine Type', 'VintageYear')

# 세 번째
# config
# config파일을 불러와야 한다
config_file = './configs/dynamic_rcnn/dynamic_rcnn_r50_fpn_1x_coco.py' # 해당하는 아키텍처를 가져온다.(faster_rcnn을 많이 쓴다. dynamic_rcnn도 많이 쓴다., centernet, mask_rcnn은 세그멘테이션에서 사용한다., cascade_rcnn도 많이 쓴다.(cascade_rcnn_r50_fpn_1x_coco.py많이씀))
cfg = Config.fromfile(config_file) # config를 읽어 들인다.
# 어디에 num_classes가 있는지 확인을 해야 한다.
# 아래 print(cfg.pretty_text)를 사용해서
# print(cfg.pretty_text) # 해당하는 파일들을 가져온다.

# Learning rate setting
# single GPU -> 0.0025
# cfg.optimizer.lr = 0.02/8 <- 보통은 8개 GPU를 사용하기 때문에 이렇게 해준다.
cfg.optimizer.lr = 0.0025 # learning rate를 설정해준다.

# 네 번째
# dataset setting
# 데이터 셋을 설정해준다
cfg.dataset_type = "WineLabelDataset" # coco_detection.py에 있는 것을 바꾸어준다.
cfg.data_root = "./dataset" # 루트를 지정해준다.

# train, val, test dataset >> type data root and file img_prefix setting
# coco_detection.py에 있는 3가지(type, ann_file, img_prefix) 부분을 변경해준다
# 경로 설정을 해준다. coco_detection.py에 있는
# agmentation을 해주려면 train_pipeline쪽에서 수정을 해주어야 한다.
# trian 작성
cfg.data.train.type = "WineLabelDataset"
cfg.data.train.ann_file = "./dataset/train/_annotations.coco.json"
cfg.data.train.img_prefix = "./dataset/train/" # 데이터셋만 지정하면 뒤에 알아서 train이 붙는다.

# val 작성
cfg.data.val.type = "WineLabelDataset"
cfg.data.val.ann_file = "./dataset/valid/_annotations.coco.json"
cfg.data.val.img_prefix = "./dataset/valid/"

# test 작성
cfg.data.test.type = "WineLabelDataset"
cfg.data.test.ann_file = "./dataset/test/_annotations.coco.json"
cfg.data.test.img_prefix = "./dataset/test/"            

# 다섯 번째
# class number
# 항상 이 부분을 주의해야 한다.
# faster_rcnn_r50_fpn.py에서 수정을 해준다.
# 로이 안에서 80인 부분은 13으로 변경해준다.
# 접근 방식을 .으로 접근하면 된다.
cfg.model.roi_head.bbox_head.num_classes = 13

# 8로하게 되면 작은 물체가 더 작아지기 때문에 4로 수정한다.
# small object -> change anchor -> df : size 8 -> size 4로 수정한다
# 8이면 스케일이 켜지는 것인데 작은 이미지는 더 작아지게 되므로 4로 하게 해주었다.(안그러면 작은 데이터를 더 작게 만든다.)
cfg.model.rpn_head.anchor_generator.scales = [4]        

# 여섯 번째
# pretrained call
# pretrained를 가져온다. 모듈을 가져온다. 아키텍처에 해당하는 pth파일을 가져와야 한다.
cfg.load_from = "./dynamic_rcnn_r50_fpn_1x-62a3f276.pth"

# train model save dir
# train 모델을 저장하는 경로를 지정한다. (모델을 저장할 공간을 설정해준다.)
cfg.work_dir = "./work_dirs/0103"

# lr, hyp setting
# 러닝 레이트랑 하이퍼파라메타를 세팅한다.
cfg.lr_config.warmup = None
cfg.log_config.interval = 10 # 10번째 마나 띄워준다. # 어느정도에서 볼지 정해준다.

# coco dataset evaluation type = bbox
# coco 데이터 셋 평가 타입을 bbox로 잡는다. (그러면 0.5 ~ 0.95사이의 값을 갖는다.)
# mAP iou threshold 0.5 ~ 0.95
cfg.evaluation.metric = 'bbox'
cfg.evaluation.interval = 10 # 언제쯤 볼지 정해준다.
cfg.checkpoint_config.interval = 10 # 언제쯤 저장을 해줄지 해주는 것이다.
# 투스테이션에서는 best를 뽑을 수는 없다.

# epoch setting
# 보통 8 x 12정도로 하게 된다. -> 96(8개 GPU를 사용할 때, 12번 돌린다)
# num work랑 gpu를 설정해주는 방법이다.
cfg.runner.max_epochs = 88 # 90정도의 에포크를 돌린다. 데이터가 충분히 있다면 88정도 해주어도 괜찮다.
cfg.seed = 777
cfg.data.samples_per_gpu = 6 # 싱글로 할 때는 6하고 2는 고정이다.
cfg.data.workers_per_gpu = 2
print("cfg.data >>", cfg.data)
cfg.gpu_ids = range(1) # gpu가 몇개인지 개수를 파악하는 것이다.
cfg.device = "cuda"
set_random_seed(777, deterministic=False)
print("cfg info show", cfg.pretty_text)

datasets = [build_dataset((cfg.data.train))] # 트레인용 데이터 셋을 생성하는 것이다.
print("dataset[0]", datasets[0])

# datasets[0].__dict__ variables key val
datasets[0].__dict__.keys()

# 모델을 불러주는 것이다.
model = build_detector(cfg.model, train_cfg=cfg.get("train_cfg"),
                    test_cfg=cfg.get('test_cfg'))
model.CLASSES = datasets[0].CLASSES
print(model.CLASSES)

if __name__ == "__main__":
    train_detector(model, datasets, cfg, distributed=False, validate=True)

import os
import json
from tqdm import tqdm
import shutil


def convert_bbox_coco2yolo(img_width, img_height, bbox):
"""
Convert bounding box from COCO  format to YOLO format

Parameters
----------
img_width : int
        width of image
img_height : int
        height of image
bbox : list[int]
        bounding box annotation in COCO format:
        [top left x position, top left y position, width, height]

Returns
-------
list[float]
        bounding box annotation in YOLO format:
        [x_center_rel, y_center_rel, width_rel, height_rel]
"""

# YOLO bounding box format: [x_center, y_center, width, height]
# (float values relative to width and height of image)
x_tl, y_tl, w, h = bbox

dw = 1.0 / img_width
dh = 1.0 / img_height

x_center = x_tl + w / 2.0
y_center = y_tl + h / 2.0

x = x_center * dw
y = y_center * dh
w = w * dw
h = h * dh

return [x, y, w, h]

def make_folders(path="output"):
if os.path.exists(path):
        shutil.rmtree(path)
os.makedirs(path)
return path

def convert_coco_json_to_yolo_txt(output_path, json_file):

path = make_folders(output_path)

with open(json_file) as f:
        json_data = json.load(f)

label_file = os.path.join(output_path, "_darknet.labels")
with open(label_file, 'w') as f:
        for category in tqdm(json_data["categories"], desc="Categores"):
        category_name = category["name"]
        f.write(f"{category_name}\n")

for image in tqdm(json_data["images"], desc="Annotation txt for each image"):
        img_id = image["id"]
        img_name = image["file_name"]
        img_width = image["width"]
        img_height = image["height"]

        anno_in_image = [anno for anno in json_data["annotations"] if anno["image_id"] == img_id]
        anno_txt = os.path.join(output_path, img_name[:-4] + ".txt")
        with open(anno_txt, "w") as f:
        for anno in anno_in_image:
                category = anno["category_id"]
                bbox_COCO = anno["bbox"]
                x, y, w, h = convert_bbox_coco2yolo(img_width, img_height, bbox_COCO)
                f.write(f"{category} {x:.6f} {y:.6f} {w:.6f} {h:.6f}\n")

print("Converting COCO Json to YOLO txt finished!")

# # test파일에 있는 json파일을 yolo파일로 변환해주었다.
# convert_coco_json_to_yolo_txt("./output/test", "./wine labels_coco/test/_annotations.coco.json")

# # train파일에 있는 json파일을 yolo파일로 변환해주었다.
# convert_coco_json_to_yolo_txt("./output/train", "./wine labels_coco/train/_annotations.coco.json")

# valid파일에 있는 json파일을 yolo파일로 변환해주었다.
convert_coco_json_to_yolo_txt("./output/valid", "./wine labels_coco/valid/_annotations.coco.json")