Commit 2fe71ac9 authored by jithu.tagore's avatar jithu.tagore

changed classes in init

parent 13b75cd2
......@@ -68,7 +68,7 @@ class CementBagCounter(ModelWrapper):
# with open("/home/ilens/container_weights/classes.json", 'r') as f:
# self.classes = json.loads(f.read())
# self.classes = {int(k): v for k, v in self.classes.items()}
self.classes = {0: 'ambuja_plus', 1: 'acc_gold', 2: 'acc_suraksha_power_plus', 3: 'ambuja_buildcem', 4: 'mrp', 5: 'acc_suraksha_power', 6: 'acc_nfr', 7: 'acc_concrete_plus'}
self.classes = {0: 'acc_gold', 1: 'mrp'}
self.yolo_v5_wrapper = YoloV5TRT(engine_file_path, model_config.get('conf_thresh', 0.5),
model_config.get('iou_thresh', 0.4))
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment