# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/03_brand_classification.ipynb (unless otherwise specified).

__all__ = ['Classification', 'train_cli']

# Cell

from fractalai.engine.trainer import train
from fractalai.engine.utils import get_config
import fractalai.engine.infer as infer
import cv2
from glob import glob
import pandas as pd
import random
import os
import numpy as np

from fastscript import call_parse
from fastscript import *
from fastcore.transform import store_attr

# Cell

class Classification():
    def __init__(self, config_path):
        self.cfg = config_path
    def train(self):
        config = get_config(self.cfg)
        train(config)
    def infer_img(self, img_loc: str, weights: str):
        label = infer.infer_img(self.cfg, img_loc, weights, False)
        return label

# Cell
@call_parse
def train_cli(cfg: Param("config file", type=str)="/mnt/nfshome1/FRACTAL/sindhura.k/dlg-home-content/data/classify_tv_brands.yaml"):
    cls = Classification(cfg)
    cls.train()