Upload 7 files

run/json2png.py

@@ -0,0 +1,45 @@
+import base64
+import json
+import os
+import os.path as osp
+
+import numpy as np
+import PIL.Image
+from labelme import utils
+
+
+def json2mask(root, classes)
+    before_path = os.path.join(root, 'json') # labelme files
+    jpgs_path = os.path.join(root,   'jpgs') # images
+    pngs_path = os.path.join(root, 'masks')  # annotations
+    if not os.path.exists(jpgs_path):
+        os.makedirs(jpgs_path)
+    if not os.path.exists(pngs_path):
+        os.makedirs(pngs_path)
+
+    path = before_path
+    for file in os.listdir(path):
+        if file.endswith('json'):
+            data = json.load(open(os.path.join(path, file)))
+
+            if data['imageData']:
+                imageData = data['imageData']
+            else:
+                imagePath = os.path.join(os.path.dirname(path), data['imagePath'])
+                with open(imagePath, 'rb') as f:
+                    imageData = f.read()
+                    imageData = base64.b64encode(imageData).decode('utf-8')
+
+            img = utils.img_b64_to_arr(imageData)
+            label_name_to_value = {classes[0]: 0, classes[2]:0, classes[1]: 1 }
+
+            lbl = utils.shapes_to_label(img.shape, data['shapes'], label_name_to_value)
+
+            PIL.Image.fromarray(img).save(osp.join(jpgs_path, file.split(".")[0] + '.jpg'))
+
+            utils.lblsave(osp.join(pngs_path, file.split(".")[0] + '.png'), lbl)
+
+if __name__ == '__main__':
+    root = 'J:/dataset_panicle/2023/only_plant/images'
+    classes = ["_background_",  "panicle", "other"]
+    json2mask(root,classes)

run/labelme2mask.py

@@ -0,0 +1,93 @@
+import os
+import json
+import numpy as np
+import cv2
+import shutil
+from tqdm import tqdm
+
+
+def labelme2mask_single_img(img_path, labelme_json_path, class_info):
+    '''
+    Convert a single image's LabelMe annotation to a mask.
+    '''
+    img_bgr = cv2.imread(img_path)
+    img_mask = np.zeros(img_bgr.shape[:2], dtype=np.uint8)  # Create an empty mask image (0 - background)
+
+    with open(labelme_json_path, 'r', encoding='utf-8') as f:
+        labelme = json.load(f)
+
+    for one_class in class_info:  # Iterate over each class in class_info
+        for each in labelme['shapes']:  # Iterate over all shapes in the annotation
+            if each['label'] == one_class['label']:
+                if one_class['type'] == 'polygon':  # Handle polygon annotation
+                    points = [np.array(each['points'], dtype=np.int32).reshape((-1, 1, 2))]  # Ensure correct shape
+                    img_mask = cv2.fillPoly(img_mask, points, color=one_class['color'])
+                elif one_class['type'] == 'line' or one_class['type'] == 'linestrip':  # Handle line annotation
+                    points = [np.array(each['points'], dtype=np.int32).reshape((-1, 1, 2))]
+                    img_mask = cv2.polylines(img_mask, points, isClosed=False, color=one_class['color'],
+                                             thickness=one_class.get('thickness', 1))
+                elif one_class['type'] == 'circle':  # Handle circle annotation
+                    points = np.array(each['points'], dtype=np.int32)
+                    center_x, center_y = points[0][0], points[0][1]
+                    edge_x, edge_y = points[1][0], points[1][1]
+                    radius = int(np.linalg.norm([center_x - edge_x, center_y - edge_y]))
+                    img_mask = cv2.circle(img_mask, (center_x, center_y), radius, one_class['color'], -1)
+                else:
+                    print('Unknown annotation type:', one_class['type'])
+
+    return img_mask
+
+
+def convert_labelme_to_mask(Dataset_Path):
+    '''
+    Convert all LabelMe annotations in the dataset to mask images.
+    '''
+    # Dataset directories
+    img_dir = os.path.join(Dataset_Path, 'images')
+    ann_dir = os.path.join(Dataset_Path, 'labelme_jsons')
+
+    # Class information for mask generation
+    class_info = [
+        {'label': 'panicle', 'type': 'polygon', 'color': 1}
+    ]
+
+    # Create target directories
+    images_target_dir = os.path.join(Dataset_Path, 'img_dir')
+    ann_target_dir = os.path.join(Dataset_Path, 'ann_dir')
+
+    # Create target directories if they do not exist
+    os.makedirs(images_target_dir, exist_ok=True)
+    os.makedirs(ann_target_dir, exist_ok=True)
+
+    # Process each image in the images directory
+    for img_name in tqdm(os.listdir(img_dir), desc="Converting images to masks"):
+        try:
+            img_path = os.path.join(img_dir, img_name)
+            labelme_json_path = os.path.join(ann_dir, f'{os.path.splitext(img_name)[0]}.json')
+
+            if os.path.exists(labelme_json_path):
+                # Convert LabelMe annotations to mask
+                img_mask = labelme2mask_single_img(img_path, labelme_json_path, class_info)
+
+                # Save the mask to the target directory
+                mask_path = os.path.join(ann_target_dir, f'{os.path.splitext(img_name)[0]}.png')
+                cv2.imwrite(mask_path, img_mask)
+
+                # Move the image to the target directory
+                shutil.move(img_path, os.path.join(images_target_dir, img_name))
+            else:
+                print(f"Annotation file missing for {img_name}")
+
+        except Exception as e:
+            print(f"Failed to convert {img_name}: {e}")
+
+    # Optionally remove the original directories if they are empty
+    shutil.rmtree(img_dir, ignore_errors=True)
+    shutil.rmtree(ann_dir, ignore_errors=True)
+
+    print("Conversion completed.")
+
+
+if __name__ == '__main__':
+    Dataset_Path = 'CVRP'  # Update this to the path of your dataset
+    convert_labelme_to_mask(Dataset_Path)

run/mask2json.py

@@ -0,0 +1,74 @@
+import os
+import io
+import json
+import numpy as np
+from pycococreatortools import pycococreatortools
+from PIL import Image
+import base64
+import cv2
+from tqdm import tqdm
+def img_tobyte(img_pil):
+    
+    ENCODING = 'utf-8'
+    img_byte = io.BytesIO()
+    img_pil.save(img_byte, format='PNG')
+    binary_str2 = img_byte.getvalue()
+    imageData = base64.b64encode(binary_str2)
+    base64_string = imageData.decode(ENCODING)
+    return base64_string
+
+def mask2json(ROOT_DIR):
+    Image_DIR = os.path.join(ROOT_DIR, "pngs")
+    Label_DIR = os.path.join(ROOT_DIR, "pre")
+
+    Label_files = os.listdir(Label_DIR)
+    # classs_names
+    class_names = ['_background_', 'panicle']
+    for Label_filename in tqdm(Label_files):
+
+        Json_output = {
+            "version": "3.16.7",
+            "flags": {},
+            "fillColor": [255, 0, 0, 128],
+            "lineColor": [0, 255, 0, 128],
+            "imagePath": {},
+            "shapes": [],
+            "imageData": {}}
+        name = Label_filename.split('.', 3)[0]
+        name1 = name + '.png'
+        Json_output["imagePath"] = name1
+
+        image = Image.open(Image_DIR + '/' + name1)
+        imageData = img_tobyte(image)
+        Json_output["imageData"] = imageData
+        binary_mask = np.asarray(np.array(Image.open(Label_DIR + '/' + Label_filename))).astype(np.uint8)
+        mask_image = cv2.imread(Label_DIR + '/' + Label_filename, cv2.IMREAD_GRAYSCALE)
+        temp_mask = np.asarray((mask_image != 0), dtype=np.uint8)
+
+        segmentation = pycococreatortools.binary_mask_to_polygon(temp_mask, tolerance=2)
+        for item in segmentation:
+            if len(item) > 10:
+                list1 = []
+                for j in range(0, len(item), 2):
+                    list1.append([item[j], item[j + 1]])
+                # There is only one non-background class, so just use class_names[1]
+                label = class_names[1]
+                seg_info = {'points': list1, "fill_color": None, "line_color": None, "label": label,
+                            "shape_type": "polygon", "flags": {}}
+                Json_output["shapes"].append(seg_info)
+
+        Json_output["imageHeight"] = binary_mask.shape[0]
+        Json_output["imageWidth"] = binary_mask.shape[1]
+        json_path = os.path.join(ROOT_DIR,'json')
+        if not os.path.exists(json_path):
+            os.makedirs(json_path)
+        full_path = os.path.join(json_path, '{}.json'.format(name))
+        with open(full_path, 'w') as output_json_file:
+            json.dump(Json_output, output_json_file)
+
+
+if __name__ == '__main__':
+
+    ROOT_DIR = ''
+    mask2json(ROOT_DIR)
+

run/requirements.txt

@@ -0,0 +1,9 @@
+opencv-python
+pillow
+matplotlib
+seaborn
+tqdm
+ftfy
+regex
+pytorch-lightning
+mmdet>=3.1.0

run/run_configs.py

@@ -0,0 +1,206 @@
+import os
+import argparse
+from mmengine import Config
+
+def create_deeplabv3plus_config(model_config_path, dataset_config_path, num_class, work_dir, save_dir, batch_size, max_iters, val_interval):
+    cfg = Config.fromfile(model_config_path)
+    dataset_cfg = Config.fromfile(dataset_config_path)
+    cfg.merge_from_dict(dataset_cfg)
+    
+    # Set crop size
+    cfg.crop_size = (512, 512)
+    cfg.model.data_preprocessor.size = cfg.crop_size
+
+      # Configure normalization
+    cfg.norm_cfg = dict(type='BN', requires_grad=True)
+    cfg.model.backbone.norm_cfg = cfg.norm_cfg
+    cfg.model.decode_head.norm_cfg = cfg.norm_cfg
+    cfg.model.auxiliary_head.norm_cfg = cfg.norm_cfg
+
+    cfg.model.decode_head.num_classes = num_class
+    cfg.model.auxiliary_head.num_classes = num_class
+
+    cfg.train_dataloader.batch_size = batch_size
+    
+    # Set training configurations
+    cfg.train_cfg.max_iters = max_iters
+    cfg.train_cfg.val_interval = val_interval
+    cfg.default_hooks.logger.interval = 100
+    cfg.default_hooks.checkpoint.interval = 2500
+    cfg.default_hooks.checkpoint.max_keep_ckpts = 1
+    cfg.default_hooks.checkpoint.save_best = 'mIoU'
+
+    cfg['randomness'] = dict(seed=0)
+    # Set work directory
+    cfg.save_dir = save_dir
+    name = os.path.basename(dataset_config_path).split('_')[0] + "_" + os.path.dirname(model_config_path).split(os.sep)[1]
+    cfg.work_dir = os.path.join(work_dir,name)
+    os.makedirs(cfg.work_dir, exist_ok=True)
+    save_config_file = os.path.join(save_dir, f"{name}.py")
+    cfg.dump(save_config_file)
+    print(f"Configuration saved to: {save_config_file}")
+
+def create_knet_config(model_config_path, dataset_config_path, num_class, work_dir, save_dir, batch_size, max_iters, val_interval):
+   
+    cfg = Config.fromfile(model_config_path)
+    dataset_cfg = Config.fromfile(dataset_config_path)
+    
+    cfg.merge_from_dict(dataset_cfg)
+    
+    cfg.norm_cfg = dict(type='BN', requires_grad=True)
+    cfg.model.data_preprocessor.size = cfg.crop_size
+
+    cfg.model.decode_head.kernel_generate_head.num_classes = num_class
+    cfg.model.auxiliary_head.num_classes = num_class
+
+    cfg.train_dataloader.batch_size = batch_size
+    cfg.work_dir = work_dir
+
+    cfg.train_cfg.max_iters = max_iters
+    cfg.train_cfg.val_interval = val_interval
+    cfg.default_hooks.logger.interval = 100
+    cfg.default_hooks.checkpoint.interval = 2500
+    cfg.default_hooks.checkpoint.max_keep_ckpts = 1
+    cfg.default_hooks.checkpoint.save_best = 'mIoU'
+
+    cfg['randomness'] = dict(seed=0)
+
+    cfg.save_dir = save_dir
+    name = os.path.basename(dataset_config_path).split('_')[0] + "_" + os.path.dirname(model_config_path).split(os.sep)[1]
+    cfg.work_dir = os.path.join(work_dir, name)
+    os.makedirs(cfg.work_dir, exist_ok=True)
+    save_config_file = os.path.join(save_dir, f"{name}.py")
+    cfg.dump(save_config_file)
+    print(f"Configuration saved to: {save_config_file}")
+
+def create_mask2former_config(model_config_path, dataset_config_path, num_class, work_dir, save_dir, batch_size, max_iters, val_interval):
+    cfg = Config.fromfile(model_config_path)
+    dataset_cfg = Config.fromfile(dataset_config_path)
+    cfg.merge_from_dict(dataset_cfg)
+    
+    # Set crop size
+    cfg.crop_size = (512, 512)
+    cfg.model.data_preprocessor.size = cfg.crop_size
+
+      # Configure normalization
+    cfg.norm_cfg = dict(type='BN', requires_grad=True)
+
+    cfg.model.decode_head.num_classes = num_class
+    cfg.model.decode_head.loss_cls.class_weight = [1.0] * num_class + [0.1]
+
+    cfg.train_dataloader.batch_size = batch_size
+    
+    # Set training configurations
+    cfg.train_cfg.max_iters = max_iters
+    cfg.train_cfg.val_interval = val_interval
+    cfg.default_hooks.logger.interval = 100
+    cfg.default_hooks.checkpoint.interval = 2500
+    cfg.default_hooks.checkpoint.max_keep_ckpts = 1
+    cfg.default_hooks.checkpoint.save_best = 'mIoU'
+
+    cfg['randomness'] = dict(seed=0)
+    # Set work directory
+    cfg.save_dir = save_dir
+    name = os.path.basename(dataset_config_path).split('_')[0] + "_" + os.path.dirname(model_config_path).split(os.sep)[1]
+    cfg.work_dir = os.path.join(work_dir,name)
+    os.makedirs(cfg.work_dir, exist_ok=True)
+    save_config_file = os.path.join(save_dir, f"{name}.py")
+    cfg.dump(save_config_file)
+    print(f"Configuration saved to: {save_config_file}")
+
+def create_segformer_config(model_config_path, dataset_config_path, num_class, work_dir, save_dir, batch_size, max_iters, val_interval):
+    cfg = Config.fromfile(model_config_path)
+    dataset_cfg = Config.fromfile(dataset_config_path)
+    cfg.merge_from_dict(dataset_cfg)
+
+      # Configure normalization
+    cfg.norm_cfg = dict(type='BN', requires_grad=True)
+    cfg.model.data_preprocessor.size = cfg.crop_size
+    cfg.model.decode_head.norm_cfg = cfg.norm_cfg
+
+    cfg.model.decode_head.num_classes = num_class
+
+    cfg.train_dataloader.batch_size = batch_size
+    
+    # Set training configurations
+    cfg.train_cfg.max_iters = max_iters
+    cfg.train_cfg.val_interval = val_interval
+    cfg.default_hooks.logger.interval = 100
+    cfg.default_hooks.checkpoint.interval = 2500
+    cfg.default_hooks.checkpoint.max_keep_ckpts = 1
+    cfg.default_hooks.checkpoint.save_best = 'mIoU'
+
+    cfg['randomness'] = dict(seed=0)
+    # Set work directory
+    cfg.save_dir = save_dir
+    name = os.path.basename(dataset_config_path).split('_')[0] + "_" + os.path.dirname(model_config_path).split(os.sep)[1]
+    cfg.work_dir = os.path.join(work_dir,name)
+    os.makedirs(cfg.work_dir, exist_ok=True)
+    save_config_file = os.path.join(save_dir, f"{name}.py")
+    cfg.dump(save_config_file)
+    print(f"Configuration saved to: {save_config_file}")
+
+def main():
+    parser = argparse.ArgumentParser(description='Train configuration setup for different models.')
+    
+    parser.add_argument('--model_name', type=str, required=True, choices=['deeplabv3plus', 'knet', 'mask2former', 'segformer'],
+                        help='Model name to generate the config for.')
+    parser.add_argument('-m', '--model_config', type=str, required=True, help="Path to the model config file")
+    parser.add_argument('-d', '--dataset_config', type=str, required=True, help='Path to the dataset config file.')
+    parser.add_argument('-c', '--num_class', type=int, required=True, help="Number of classes in the dataset")
+    parser.add_argument('-w','--work_dir', type=str, required=True, help='Directory to save the train result.')
+    parser.add_argument('-s', '--save_dir', type=str, required=True, help="Directory to save the generated config file")
+    parser.add_argument('--batch_size', type=int, default=2, help='Batch size for training.')
+    parser.add_argument('--max_iters', type=int, default=20000, help='Number of training iterations.')
+    parser.add_argument('--val_interval', type=int, default=500, help='Interval for validation during training.')
+
+
+    args = parser.parse_args()
+
+    if args.model_name == 'deeplabv3plus':
+        create_deeplabv3plus_config(
+            model_config_path=args.model_config,
+            dataset_config_path=args.dataset_config,
+            num_class=args.num_class,
+            work_dir=args.work_dir,
+            save_dir =args.save_dir,
+            batch_size=args.batch_size,
+            max_iters=args.max_iters,
+            val_interval=args.val_interval
+        )
+    if args.model_name == 'knet':
+        create_knet_config(
+            model_config_path=args.model_config,
+            dataset_config_path=args.dataset_config,
+            num_class=args.num_class,
+            work_dir=args.work_dir,
+            save_dir =args.save_dir,
+            batch_size=args.batch_size,
+            max_iters=args.max_iters,
+            val_interval=args.val_interval
+        )
+    if args.model_name == 'mask2former':
+            create_mask2former_config(
+            model_config_path=args.model_config,
+            dataset_config_path=args.dataset_config,
+            num_class=args.num_class,
+            work_dir=args.work_dir,
+            save_dir =args.save_dir,
+            batch_size=args.batch_size,
+            max_iters=args.max_iters,
+            val_interval=args.val_interval
+        )
+    elif args.model_name == 'segformer':
+            create_segformer_config(
+            model_config_path=args.model_config,
+            dataset_config_path=args.dataset_config,
+            num_class=args.num_class,
+            work_dir=args.work_dir,
+            save_dir =args.save_dir,
+            batch_size=args.batch_size,
+            max_iters=args.max_iters,
+            val_interval=args.val_interval
+        )
+
+if __name__ == '__main__':
+    main()

run/split_dataset.py

@@ -0,0 +1,70 @@
+import os
+import shutil
+import random
+from tqdm import tqdm
+
+def main():
+    # Set the dataset path
+    Dataset_Path = 'CVRPDataset'
+    img_dir = os.path.join(Dataset_Path, 'img_dir')
+    ann_dir = os.path.join(Dataset_Path, 'ann_dir')
+
+    # Set the ratio of training set to test set
+    test_frac = 0.1 
+    random.seed(123) 
+
+    # Create directories if not exist
+    os.makedirs(os.path.join(Dataset_Path, 'train'), exist_ok=True)
+    os.makedirs(os.path.join(Dataset_Path, 'val'), exist_ok=True)
+
+    # Get image file paths
+    img_paths = os.listdir(img_dir)
+    random.shuffle(img_paths) 
+
+    val_number = int(len(img_paths) * test_frac)  
+    train_files = img_paths[val_number:]          
+    val_files = img_paths[:val_number]            
+
+    print(f"Total images: {len(img_paths)}")
+    print(f"Training set images: {len(train_files)}")
+    print(f"Test set images: {len(val_files)}")
+
+    # Move the training set images to the train directory
+    for each in tqdm(train_files, desc="Move the training set images"):
+        src_path = os.path.join(img_dir, each)
+        dst_path = os.path.join(Dataset_Path, 'train', each)
+        shutil.move(src_path, dst_path)
+        
+    # Move the test set images to the test directory
+    for each in tqdm(val_files, desc="Move the test set images"):
+        src_path = os.path.join(img_dir, each)
+        dst_path = os.path.join(Dataset_Path, 'val', each)
+        shutil.move(src_path, dst_path)
+
+    # Move the train and val directories into img_dir
+    shutil.move(os.path.join(Dataset_Path, 'train'), os.path.join(img_dir, 'train'))
+    shutil.move(os.path.join(Dataset_Path, 'val'), os.path.join(img_dir, 'val'))
+
+    # Process annotation files
+    # Ensure the annotation directories exist
+    os.makedirs(os.path.join(Dataset_Path, 'train'), exist_ok=True)
+    os.makedirs(os.path.join(Dataset_Path, 'val'), exist_ok=True)
+
+    # Move the training set annotation files to the train directory
+    for each in tqdm(train_files, desc="Move the training set annotations"):
+        src_path = os.path.join(ann_dir, each.split('.')[0] + '.png')
+        dst_path = os.path.join(Dataset_Path, 'train', each.split('.')[0] + '.png')
+        shutil.move(src_path, dst_path)
+
+    # Move the test set annotation files to the test directory
+    for each in tqdm(val_files, desc="Move the test set annotations"):
+        src_path = os.path.join(ann_dir, each.split('.')[0] + '.png')
+        dst_path = os.path.join(Dataset_Path, 'val', each.split('.')[0] + '.png')
+        shutil.move(src_path, dst_path)
+
+    # Move the train and val annotation directories into ann_dir
+    shutil.move(os.path.join(Dataset_Path, 'train'), os.path.join(ann_dir, 'train'))
+    shutil.move(os.path.join(Dataset_Path, 'val'), os.path.join(ann_dir, 'val'))
+
+if __name__ == '__main__':
+    main()

run/test.py

@@ -0,0 +1,58 @@
+import os
+import numpy as np
+import cv2
+from tqdm import tqdm
+import argparse
+from mmseg.apis import init_model, inference_model
+
+
+def process_single_img(img_path, model, outpath, palette_dict):
+    
+    img_bgr = cv2.imread(img_path)
+
+    result = inference_model(model, img_bgr)
+    pred_mask = result.pred_sem_seg.data[0].cpu().numpy()
+
+    # Map the predicted integer ID to the color of the corresponding category
+    pred_mask_bgr = np.zeros((pred_mask.shape[0], pred_mask.shape[1], 3))
+    for idx in palette_dict.keys():
+        pred_mask_bgr[np.where(pred_mask==idx)] = palette_dict[idx]
+    pred_mask_bgr = pred_mask_bgr.astype('uint8')
+    
+    save_path = os.path.join(outpath, os.path.basename(img_path))
+    cv2.imwrite(save_path, pred_mask_bgr)
+
+
+
+def main(args):
+    # Initialize model
+    model = init_model(args.config_file, args.checkpoint_file, device=args.device)
+
+    # Define class palette
+    palette = [
+        ['background', [0, 0, 0]],
+        ['red', [0, 0, 255]]
+    ]
+    palette_dict = {idx: each[1] for idx, each in enumerate(palette)}
+
+    # Create output directory if not exists
+    if not os.path.exists(args.outpath):
+        os.mkdir(args.outpath)
+
+    # Process each image in the given directory
+    for img_name in tqdm(os.listdir(args.data_folder)):
+        img_path = os.path.join(args.data_folder, img_name)
+        process_single_img(img_path, model, args.outpath, palette_dict)
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description="Process images for semantic segmentation inference.")
+    parser.add_argument('-d','--data_folder', type=str, required=True, help="Path to the folder containing input images.")
+    parser.add_argument('-m','--config_file', type=str, required=True, help="Path to the model config file.")
+    parser.add_argument('-pth','--checkpoint_file', type=str, required=True, help="Path to the model checkpoint file.")
+    parser.add_argument('-o','--outpath', type=str, help="Path to save the output images.")
+    parser.add_argument('--device', type=str, default='cuda:0', help="Device to run the model (e.g., 'cuda:0', 'cpu').")
+    
+    args = parser.parse_args()
+    main(args)
+