"""
# YOLOv5 for TACO
## train & detect TACO-2
"""
"""
## Repro [Github](https://github.com/ultralytics/yolov5)
"""
!git clone https://github.com/rkuo2000/yolov5
%cd yolov5
!pip install pycocotools
#from utils.google_utils import *
#attempt_download('weights/yolov5s.pt')
#attempt_download('weights/yolov5m.pt')
#attempt_download('weights/yolov5l.pt')
#attempt_download('weights/yolov5x.pt')
"""
## Prepare Dataset
"""
import os
import shutil
import numpy as np
import tqdm
"""
### read annotations.json
"""
from pycocotools.coco import COCO
data_source = COCO(annotation_file='/kaggle/input/tacotrashdataset/data/annotations.json')
"""
>      'Clear plastic bottle': 5
>      'Plastic bottle cap': 7
>      'Drink can': 12
>      'Other plastic': 29
>      'Plastic film': 36
>      'Other plastic wrapper': 39
>      'Unlabeled litter': 58
>      'Cigarette': 59
"""
# remapping label id to 0~1
label_transfer = {5: 0, 12: 1}
img_ids = data_source.getImgIds()

catIds = data_source.getCatIds()
categories = data_source.loadCats(catIds)
categories.sort(key=lambda x: x['id'])
classes = {}
coco_labels = {}
coco_labels_inverse = {}
for c in categories:
    coco_labels[len(classes)] = c['id']
    coco_labels_inverse[c['id']] = len(classes)
    classes[c['name']] = len(classes)

class_num = {}
"""
### collect images & labels
"""
!mkdir -p tmp/labels tmp/images
save_base_path  = 'tmp/labels/'
save_image_path = 'tmp/images/'
for index, img_id in tqdm.tqdm(enumerate(img_ids), desc='change .json file to .txt file'):
    img_info = data_source.loadImgs(img_id)[0]
    # 將含文件夾的路徑修改為文件名
    save_name = img_info['file_name'].replace('/', '_')
    # 移去文件擴展名
    file_name = save_name.split('.')[0]
    # 獲取單張圖像的寬和高
    height = img_info['height']
    width = img_info['width']
    # 轉換所得txt文件存儲路徑
    save_path = save_base_path + file_name + '.txt'
    is_exist = False  # 記錄圖片是否包含目標垃圾類型對象
    with open(save_path, mode='w') as fp:
        # 根據圖片編號找出垃圾對象的編號集合
        annotation_id = data_source.getAnnIds(img_id)
        boxes = np.zeros((0, 5))
        if len(annotation_id) == 0:  # 集合大小為0
            fp.write('')
            continue
        # 獲取coco格式的標籤
        annotations = data_source.loadAnns(annotation_id)
        lines = ''  # 記錄轉換後yolo格式的標籤
        # 遍歷對象標籤集
        for annotation in annotations:
            # 獲取垃圾對象的標籤
            label = coco_labels_inverse[annotation['category_id']]
            if label in label_transfer.keys():
                # 垃圾類型屬於目標垃圾類型則進行格式轉換
                is_exist = True
                box = annotation['bbox']
                if box[2] < 1 or box[3] < 1:
                    # 如果原標籤中出現無長或寬數據的情況則跳過
                    continue
                # top_x,top_y,width,height==>cen_x,cen_y,width,height
                box[0] = round((box[0] + box[2] / 2) / width, 6)
                box[1] = round((box[1] + box[3] / 2) / height, 6)
                box[2] = round(box[2] / width, 6)
                box[3] = round(box[3] / height, 6)
                label = label_transfer[label]  # 標籤映射
                if label not in class_num.keys():
                    class_num[label] = 0
                class_num[label] += 1
                lines = lines + str(label)  # 先存儲標籤
                for i in box:  # 再存儲位置信息
                    lines += ' ' + str(i)
                lines += '\n'  # 換行
        fp.writelines(lines)
    if is_exist:
        # 存在目標類型對象，則拷貝圖像至指定目錄
        shutil.copy('/kaggle/input/tacotrashdataset/data/{}'.format(img_info['file_name']), os.path.join(save_image_path, save_name))
    else:
        # 不存在則刪除所生成的標籤文件
        os.remove(save_path)
"""
### Split Folder 
#### split into train, val, test (images & labels under each folder)
"""
!pip install split-folders
"""
### Display Test Images
"""
import splitfolders
splitfolders.ratio('tmp', output="taco", seed=1337, ratio=(.8, 0.1,0.1)) 
"""
## Train Model
"""
"""
>     $ python train.py --data coco.yaml --cfg yolov5s.yaml --weights '' --batch-size 64
                                             yolov5m                                40
                                             yolov5l                                24
                                             yolov5x                                16
"""
#!python train.py --img 320 --batch 1 --epochs 100 --data /kaggle/input/yolov5-taco-models/taco2.yaml --cfg models/yolov5s.yaml --weights yolov5s.pt
#!python train.py --img 320 --batch 1 --epochs 100 --data /kaggle/input/yolov5-taco-models/taco2.yaml --cfg models/yolov5s.yaml --weights /kaggle/input/yolov5-taco-models/taco2_epoch100_best.pt
#!python train.py --img 320 --batch 1 --epochs 100 --data /kaggle/input/yolov5-taco-models/taco2.yaml --cfg models/yolov5s.yaml --weights /kaggle/input/yolov5-taco-models/taco2_epoch200_best.pt
!python train.py --img 320 --batch 1 --epochs 100 --data /kaggle/input/yolov5-taco-models/taco2.yaml --cfg models/yolov5s.yaml --weights /kaggle/input/yolov5-taco-models/taco2_epoch300_best.pt
"""
## Detect Trash
"""
## Model weights
#!cp /kaggle/input/yolov5-taco-models/taco2_epoch300_best.pt weights/best.pt
!cp runs/exp0/weights/best.pt weights
!ls taco/test/images
from IPython.display import Image
!python detect.py --weights weights/best.pt --img 320 --conf 0.4 --source taco/test/images/batch_1_000000.jpg
Image('inference/output/batch_1_000000.jpg')
!python detect.py --weights weights/best.pt --img 320 --conf 0.4 --source taco/test/images/batch_5_000017.JPG 
Image('inference/output/batch_5_000017.JPG')
"""
## Detect test images folder
"""
!python detect.py --weights weights/best.pt --img 320 --conf 0.4 --source taco/test/images 
Image('inference/output/batch_7_000114.JPG')
Image('inference/output/batch_8_000009.jpg')
Image('inference/output/batch_8_000020.jpg')
Image('inference/output/batch_8_000027.jpg')
Image('inference/output/batch_8_000069.jpg')
Image('inference/output/batch_8_000085.jpg')
Image('inference/output/batch_9_000001.jpg')
Image('inference/output/batch_9_000074.jpg')