import json
import mmcv
import numpy as np
import os.path as osp
from mmdet.apis import train_detector, set_random_seed
from mmdet.datasets import build_dataset
from mmdet.models import build_detector

class ObjectDetector:
    def __init__(self, config_path):
        with open(config_path, 'r') as file:
            self.config = json.load(file)
        
        self.cfg = mmcv.Config.fromfile('./configs/faster_rcnn/faster_rcnn_r50_caffe_fpn_mstrain_1x_coco.py')  # Update this path to your actual config file
        self._update_config()

    def _update_config(self):
        # Dataset configuration
        self.cfg.dataset_type = 'CocoDataset'
        self.cfg.data_root = 'C:/Users/User/Desktop/Ronit-Projects/Detr_Train/dataset'

        self.cfg.data.train.type = self.cfg.dataset_type
        self.cfg.data.train.data_root = self.config['train_folder']
        self.cfg.data.train.ann_file = self.config['train_json']
        self.cfg.data.train.img_prefix = ''

        self.cfg.data.val.type = self.cfg.dataset_type
        self.cfg.data.val.data_root = self.config['val_folder']
        self.cfg.data.val.ann_file = self.config['val_json']
        self.cfg.data.val.img_prefix = ''

        self.cfg.data.test.type = self.cfg.dataset_type
        self.cfg.data.test.data_root = self.config['test_folder']
        self.cfg.data.test.ann_file = self.config['test_json']
        self.cfg.data.test.img_prefix = ''

        # Model settings
        self.cfg.model.roi_head.bbox_head.num_classes = self.config['num_classes']

        # Pre-trained model
        self.cfg.load_from = 'checkpoints/mask_rcnn_r50_caffe_fpn_mstrain-poly_3x_coco_bbox_mAP-0.408__segm_mAP-0.37_20200504_163245-42aa3d00.pth' 


        # Training settings
        self.cfg.runner.max_epochs = self.config['epochs']
        self.cfg.work_dir = './your_training_output_directory'
        self.cfg.optimizer.lr = 0.02 / 8  # Adjust as necessary
        self.cfg.lr_config.warmup = None
        self.cfg.log_config.interval = 10

        self.cfg.evaluation.metric = 'mAP'
        self.cfg.evaluation.interval = 12
        self.cfg.checkpoint_config.interval = 12

        # Seed for reproducibility
        self.cfg.seed = 0
        set_random_seed(0, deterministic=False)
        self.cfg.gpu_ids = range(1)  # Adjust the number of GPUs

    def train(self):
        # Build dataset and model
        datasets = [build_dataset(self.cfg.data.train)]
        model = build_detector(self.cfg.model, train_cfg=self.cfg.train_cfg, test_cfg=self.cfg.test_cfg)
        model.CLASSES = datasets[0].CLASSES

        # Create directory for work
        mmcv.mkdir_or_exist(osp.abspath(self.cfg.work_dir))
        train_detector(model, datasets, self.cfg, distributed=False, validate=True)

# Usage
config_path = 'config.json'
detector = ObjectDetector(config_path)
detector.train()