[Feature]: cvmart competition baseline

cvmart/Helmet_identification_10163/convert_to_coco.py

@@ -0,0 +1,153 @@
+# coding:utf-8
+
+# pip install lxml
+
+import os
+import glob
+import json
+import shutil
+import numpy as np
+
+import xml.etree.ElementTree as ET
+
+path2 = "."
+
+START_BOUNDING_BOX_ID = 1
+
+
+def get(root, name):
+    return root.findall(name)
+
+
+def get_and_check(root, name, length):
+    vars = root.findall(name)
+    if len(vars) == 0:
+        raise NotImplementedError('Can not find %s in %s.' % (name, root.tag))
+    if length > 0 and len(vars) != length:
+        raise NotImplementedError('The size of %s is supposed to be %d, but is %d.' % (name, length, len(vars)))
+    if length == 1:
+        vars = vars[0]
+    return vars
+
+
+def convert(xml_list, json_file):
+    json_dict = {"images": [], "type": "instances", "annotations": [], "categories": []}
+    categories = pre_define_categories.copy()
+    bnd_id = START_BOUNDING_BOX_ID
+    all_categories = {}
+    for index, line in enumerate(xml_list):
+        # print("Processing %s"%(line))
+        xml_f = line
+        tree = ET.parse(xml_f)
+        root = tree.getroot()
+
+        filename = os.path.basename(xml_f)[:-4] + ".jpg"
+        image_id = 20190000001 + index
+        size = get_and_check(root, 'size', 1)
+        width = int(get_and_check(size, 'width', 1).text)
+        height = int(get_and_check(size, 'height', 1).text)
+        image = {'file_name': filename, 'height': height, 'width': width, 'id': image_id}
+        json_dict['images'].append(image)
+        ## Cruuently we do not support segmentation
+        #  segmented = get_and_check(root, 'segmented', 1).text
+        #  assert segmented == '0'
+        for obj in get(root, 'object'):
+            category = get_and_check(obj, 'name', 1).text
+            if category in all_categories:
+                all_categories[category] += 1
+            else:
+                all_categories[category] = 1
+            if category not in categories:
+                if only_care_pre_define_categories:
+                    continue
+                new_id = len(categories) + 1
+                print(
+                    "[warning] category '{}' not in 'pre_define_categories'({}), create new id: {} automatically".format(
+                        category, pre_define_categories, new_id))
+                categories[category] = new_id
+            category_id = categories[category]
+            bndbox = get_and_check(obj, 'bndbox', 1)
+            xmin = int(float(get_and_check(bndbox, 'xmin', 1).text))
+            ymin = int(float(get_and_check(bndbox, 'ymin', 1).text))
+            xmax = int(float(get_and_check(bndbox, 'xmax', 1).text))
+            ymax = int(float(get_and_check(bndbox, 'ymax', 1).text))
+            assert (xmax > xmin), "xmax <= xmin, {}".format(line)
+            assert (ymax > ymin), "ymax <= ymin, {}".format(line)
+            o_width = abs(xmax - xmin)
+            o_height = abs(ymax - ymin)
+            ann = {'area': o_width * o_height, 'iscrowd': 0, 'image_id':
+                image_id, 'bbox': [xmin, ymin, o_width, o_height],
+                   'category_id': category_id, 'id': bnd_id, 'ignore': 0,
+                   'segmentation': []}
+            json_dict['annotations'].append(ann)
+            bnd_id = bnd_id + 1
+
+    for cate, cid in categories.items():
+        cat = {'supercategory': 'none', 'id': cid, 'name': cate}
+        json_dict['categories'].append(cat)
+    json_fp = open(json_file, 'w')
+    json_str = json.dumps(json_dict)
+    json_fp.write(json_str)
+    json_fp.close()
+    print("------------create {} done--------------".format(json_file))
+    print("find {} categories: {} -->>> your pre_define_categories {}: {}".format(len(all_categories),
+                                                                                  all_categories.keys(),
+                                                                                  len(pre_define_categories),
+                                                                                  pre_define_categories.keys()))
+    print("category: id --> {}".format(categories))
+    print(categories.keys())
+    print(categories.values())
+
+
+if __name__ == '__main__':
+    classes = ['person', 'head', 'hat']
+    pre_define_categories = {}
+    for i, cls in enumerate(classes):
+        pre_define_categories[cls] = i + 1
+    # pre_define_categories = {'a1': 1, 'a3': 2, 'a6': 3, 'a9': 4, "a10": 5}
+    only_care_pre_define_categories = True
+    # only_care_pre_define_categories = False
+
+    train_ratio = 0.9
+    save_json_train = '/project/train/src_repo/coco_annotations/instances_train2014.json'
+    save_json_val = 'instances_val2014.json'
+    xml_dir = "/home/data/831/"
+
+    xml_list = glob.glob(xml_dir + "/*.xml")
+    xml_list = np.sort(xml_list)
+    np.random.seed(100)
+    np.random.shuffle(xml_list)
+
+    train_num = int(len(xml_list) * train_ratio)
+    xml_list_train = xml_list[:train_num]
+    xml_list_val = xml_list[train_num:]
+
+    convert(xml_list_train, save_json_train)
+    convert(xml_list_val, save_json_val)
+
+    if os.path.exists(path2 + "/annotations"):
+        shutil.rmtree(path2 + "/annotations")
+    os.makedirs(path2 + "/annotations")
+    if os.path.exists(path2 + "/images/train2014"):
+        shutil.rmtree(path2 + "/images/train2014")
+    os.makedirs(path2 + "/images/train2014")
+    if os.path.exists(path2 + "/images/val2014"):
+        shutil.rmtree(path2 + "/images/val2014")
+    os.makedirs(path2 + "/images/val2014")
+
+    f1 = open("train.txt", "w")
+    for xml in xml_list_train:
+        img = xml[:-4] + ".jpg"
+        f1.write(os.path.basename(xml)[:-4] + "\n")
+        shutil.copyfile(img, path2 + "/images/train2014/" + os.path.basename(img))
+
+    f2 = open("test.txt", "w")
+    for xml in xml_list_val:
+        img = xml[:-4] + ".jpg"
+        f2.write(os.path.basename(xml)[:-4] + "\n")
+        shutil.copyfile(img, path2 + "/images/val2014/" + os.path.basename(img))
+    f1.close()
+    f2.close()
+    print("-------------------------------")
+    print("train number:", len(xml_list_train))
+    print("val number:", len(xml_list_val))

cvmart/Helmet_identification_10163/is_there_tensorboard.py

@@ -0,0 +1,8 @@
+import os
+
+if os.path.exists('/project/train/tensorboard'):
+    # os.system('cp -r /project/train/models/train/exp/weights/* /project/train/tensorboard/')
+    os.system('rm -rf /project/train/tensorboard')
+else:
+    os.system('mkdir /project/train/tensorboard')
+    # os.system('cp -r /project/train/tensorboard/* /project/train/models/train/exp/weights/')

cvmart/Helmet_identification_10163/ji.py

@@ -0,0 +1,102 @@
+import json    # json模块是Python标准库中的一个模块，用于处理JSON格式的数据。在该代码中，使用json模块将检测结果转换为JSON格式。
+from mmdet.apis import init_detector, inference_detector
+
+import mmcv
+
+# from mmdet_custom.datasets import D10007Dataset
+
+def init():
+
+    config_file = '/project/train/src_repo/mmyolo/tools/rtmdet_tiny_syncbn_fast_8xb32-300e_coco.py'
+    checkpoint_file = '/project/train/models/train/exp/weights/epoch_320.pth'
+    model = init_detector(config_file, checkpoint_file)
+    # init_detector函数用于初始化MMDetection模型，输入参数包括模型配置文件路径config_file和模型权重文件路径checkpoint_file。该函数返回一个MMDetection模型对象model，用于后续的目标检测操作。
+    return model
+
+def process_image(handle=None, input_image=None, args=None, **kwargs):
+
+    # CLASSES = D10007Dataset.CLASSES  # D10007Dataset是一个自定义的数据集类，用于加载数据集的类别信息。在该代码中，通过D10007Dataset.CLASSES获取数据集的类别信息，用于后续的处理。
+    CLASSES = ['person', 'hat', 'head']
+    # CLASSES 是一个 Python列表  例如CLASSES的值为["person", "car", "dog"]
+
+    # result = inference_detector(handle, input_image)
+    # # inference_detector函数用于对输入图像进行目标检测，输入参数包括MMDetection模型对象handle和输入图像input_image。
+    # # 该函数返回一个包含检测结果的列表，其中每个元素表示一个检测框的信息，包括坐标、置信度、类别等。
+
+    # labels = result.pred_instances.labels
+    # objects = []
+    # for i, class_name in enumerate(CLASSES):   # 遍历数据集的类别信息CLASSES，对每个类别的检测结果进行处理。
+    #     # 具体来说，对于每个类别的检测结果，遍历其中的每个检测框，将检测框的 坐标、置信度、类别 等信息转换为字典格式，并添加到objects列表中。
+
+    #     fires = result[i]
+    #     for fire in fires:     # 具体来说，对于每个类别的检测结果，从result列表中获取该类别的检测结果fires，然后遍历其中的每个检测框fire。
+    #         obj = dict(
+    #             xmin = int(fire[0].item()),
+    #             ymin= int(fire[1].item()),
+    #             xmax=int(fire[2].item()),
+    #             ymax=int(fire[3].item()),
+    #             confidence=fire[4].item(),
+    #             name = CLASSES[i]
+    #         )
+    #     # 对于每个检测框，将其坐标、置信度、类别等信息转换为字典格式，并添加到objects列表中。
+    #     # 具体来说，使用dict函数创建一个字典对象obj，其中包含检测框的左上角和右下角坐标、置信度、类别等信息。然后，将obj添加到objects列表中。
+
+    #         if obj['confidence' ] >0.5:
+    #             objects.append(obj)
+
+    result = inference_detector(handle, input_image)
+    bboxes = result.pred_instances.bboxes
+    scores = result.pred_instances.scores
+    labels = result.pred_instances.labels
+    objects = []
+    fan = len(bboxes)
+    for i in range(len(bboxes)):
+
+            obj = dict(
+            xmin = int(bboxes[i][0]),
+            ymin = int(bboxes[i][1]),
+            xmax = int(bboxes[i][2]),
+            ymax = int(bboxes[i][3]),
+            confidence = float(scores[i]),
+            name = CLASSES[labels[i]])
+            if obj['confidence'] > 0.5:
+                            objects.append(obj)
+
+
+
+
+    # model.show_result(img, result)
+    # model.show_result(img, result, out_file='result.jpg', score_thr = 0.3)
+    r_json = dict()
+    r_json['algorithm_data'] = dict(target_info=objects, is_alert=False, target_count=0)
+    r_json['model_data'] = dict(objects=objects)
+
+    '''
+    这段代码的作用是将MMDetection模型的检测结果转换为字典格式，并添加到objects列表中。具体分析如下：
+
+    objects是一个空列表，用于存储检测结果的信息。
+
+    在循环体中，使用enumerate函数遍历数据集的类别信息CLASSES，对每个类别的检测结果进行处理。具体来说，对于每个类别的检测结果，从result列表中获取该类别的检测结果fires，然后遍历其中的每个检测框fire。
+
+    对于每个检测框，将其坐标、置信度、类别等信息转换为字典格式，并添加到objects列表中。具体来说，使用dict函数创建一个字典对象obj，
+    其中包含检测框的左上角和右下角坐标、置信度、类别等信息。然后，将obj添加到objects列表中。
+
+    在添加obj到objects列表之前，使用if语句判断检测框的置信度是否大于0.5。如果是，则将obj添加到objects列表中；否则，忽略该检测框。
+
+    最后，将objects列表转换为字典格式，并添加到r_json字典中。其中，r_json字典包含两个键值对，分别为algorithm_data和model_data。
+    algorithm_data表示算法的输出结果，包括目标信息、是否报警、目标数量等；model_data表示模型的输出结果，包括检测框的信息。
+    '''
+
+    if objects.__len__( ) >0:
+        r_json['algorithm_data']['is_alert'] = True
+        r_json['algorithm_data']['target_count'] = objects.__len__()
+
+    # return json.dumps(objects, indent=4)
+    return json.dumps(r_json, indent=4)
+
+if __name__ == "__main__":
+    handle = init()
+    # 或者 img = mmcv.imread(img), 这将只加载图像一次．
+    img = "/home/data/831/helmet_10809.jpg"
+
+    process_image(handle, img, '{"mask_output_path": "result.png"}')

cvmart/Helmet_identification_10163/move_log.py

@@ -0,0 +1,7 @@
+import os
+
+if os.path.exists('/project/train/tensorboard'):
+    os.system('cp -r /project/train/models/train/exp/weights/* /project/train/tensorboard/')
+else:
+    os.system('mkdir /project/train/tensorboard')
+    os.system('cp -r /project/train/models/train/exp/weights/* /project/train/tensorboard/')

cvmart/Helmet_identification_10163/move_model.py

@@ -0,0 +1,8 @@
+import os
+
+if os.path.exists('/project/train/tensorboard'):
+    # os.system('cp -r /project/train/models/train/exp/weights/* /project/train/tensorboard/')
+    os.system('cp -r /project/train/tensorboard/* /project/train/models/train/exp/weights/')
+else:
+    os.system('mkdir /project/train/tensorboard')
+    os.system('cp -r /project/train/tensorboard/* /project/train/models/train/exp/weights/')