Additional Changes for Core ML Conversion

.gitignore

@@ -1,4 +1,6 @@
 *.pyc
 __pycache__
 .DS_STORE
-.idea
+.idea
+
+/vision_datasets

README_Segmentation.md

@@ -103,6 +103,10 @@ CUDA_VISIBLE_DEVICES=0 python test_segmentation.py --model espnetv2 --s 2.0 --da
 CUDA_VISIBLE_DEVICES=0 python test_segmentation.py --model espnetv2 --s 2.0 --dataset city --data-path ../vision_datasets/cityscapes/ --split val --im-size 1024 512 --weights-test model/segmentation/model_zoo/espnetv2/espnetv2_s_2.0_city_1024x512.pth
 ```
 
+```
+CUDA_VISIBLE_DEVICES=0 python test_segmentation.py --model espnetv2 --s 2.0 --dataset city --data-path ./vision_datasets/cityscapes/ --split val --im-size 1024 512 --weights-test ./model/segmentation/model_zoo/espnetv2/espnetv2_s_2.0_city_1024x512.pth
+```
+
 For evaluation on the PASCAL VOC dataset, use below command:
 ```
 CUDA_VISIBLE_DEVICES=0 python test_segmentation.py --model espnetv2 --s 2.0 --dataset pascal --data-path ../vision_datasets/pascal_voc/VOCdevkit/ --split val --im-size 384 384 

core_ml_conversion.py

@@ -0,0 +1,261 @@
+import torch
+import glob
+import os
+import numpy as np
+from argparse import ArgumentParser
+from PIL import Image
+from torchvision.transforms import functional as F
+from tqdm import tqdm
+from utilities.print_utils import *
+from transforms.classification.data_transforms import MEAN, STD
+from utilities.utils import model_parameters, compute_flops
+import coremltools as ct
+
+# # Load the model and its weights
+# model = MyModel()
+# model.load_state_dict(torch.load('path_to_your_model.pth'))
+
+
+# # Example input to use for tracing
+# example_input = torch.rand(1, 3, 224, 224)  # Adjust the input shape as per your model
+
+# # Convert the model to TorchScript
+# traced_model = torch.jit.trace(model, example_input)
+
+# # Convert the TorchScript model to Core ML
+# mlmodel = ct.convert(
+#     traced_model,
+#     inputs=[ct.TensorType(shape=example_input.shape)],
+# )
+
+# # Save the Core ML model
+# mlmodel.save('path_to_save_your_model.mlmodel')
+
+
+def relabel(img):
+    '''
+    This function relabels the predicted labels so that cityscape dataset can process
+    :param img:
+    :return:
+    '''
+    img[img == 19] = 255
+    img[img == 18] = 33
+    img[img == 17] = 32
+    img[img == 16] = 31
+    img[img == 15] = 28
+    img[img == 14] = 27
+    img[img == 13] = 26
+    img[img == 12] = 25
+    img[img == 11] = 24
+    img[img == 10] = 23
+    img[img == 9] = 22
+    img[img == 8] = 21
+    img[img == 7] = 20
+    img[img == 6] = 19
+    img[img == 5] = 17
+    img[img == 4] = 13
+    img[img == 3] = 12
+    img[img == 2] = 11
+    img[img == 1] = 8
+    img[img == 0] = 7
+    img[img == 255] = 0
+    return img
+
+def data_transform(img, im_size):
+    img = img.resize(im_size, Image.BILINEAR)
+    img = F.to_tensor(img)  # convert to tensor (values between 0 and 1)
+    img = F.normalize(img, MEAN, STD)  # normalize the tensor
+    return img
+
+def evaluate(args, model, image_list, device):
+    im_size = tuple(args.im_size)
+    print(im_size)
+
+    # get color map for pascal dataset
+    if args.dataset == 'pascal':
+        from utilities.color_map import VOCColormap
+        cmap = VOCColormap().get_color_map_voc()
+    else:
+        cmap = None
+
+    model.eval()
+    for i, imgName in tqdm(enumerate(image_list)):
+        img = Image.open(imgName).convert('RGB')
+        w, h = img.size
+        print(f'w: {w}, h: {h}')
+
+        img = data_transform(img, im_size)
+        print(f'Data transform {img.shape}')
+        img = img.unsqueeze(0)  # add a batch dimension
+        print(f'Unsqueeze {img.shape}')
+        img = img.to(device)
+        img_out = model(img)
+        print(f'Model {img_out.shape}')
+        img_out = img_out.squeeze(0)  # remove the batch dimension
+        print(f'Squeeze {img_out.shape}')
+        img_out = img_out.max(0)[1].byte()  # get the label map
+        print(f'Max {img_out.shape}')
+        img_out = img_out.to(device='cpu').numpy()
+        print(img.shape, img_out.shape)
+
+        # Prints
+        # w: 2048, h: 1024
+        # Data transform torch.Size([3, 256, 512])
+        # Unsqueeze torch.Size([1, 3, 256, 512])
+        # Model torch.Size([1, 20, 256, 512])
+        # Squeeze torch.Size([20, 256, 512])
+        # Max torch.Size([256, 512])
+        # torch.Size([1, 3, 256, 512]) (256, 512)
+        return
+
+        if args.dataset == 'city':
+            # cityscape uses different IDs for training and testing
+            # so, change from Train IDs to actual IDs
+            img_out = relabel(img_out)
+
+        img_out = Image.fromarray(img_out)
+        # resize to original size
+        img_out = img_out.resize((w, h), Image.NEAREST)
+
+        # pascal dataset accepts colored segmentations
+        if args.dataset == 'pascal':
+            img_out.putpalette(cmap)
+
+        # save the segmentation mask
+        name = imgName.split('/')[-1]
+        img_extn = imgName.split('.')[-1]
+        name = '{}/{}'.format(args.savedir, name.replace(img_extn, 'png'))
+        img_out.save(name)
+
+def convert_to_coreml(model, args, device):
+    im_size = tuple(args.im_size)
+    example_input = Image.new('RGB', im_size)  #np.random.rand(im_size[0], im_size[1], 3)
+    # print(example_input.size)
+    # img = example_input.convert('RGB')
+    img = data_transform(example_input, im_size)
+    img = img.unsqueeze(0)  # add a batch dimension
+    img = img.to(device)
+
+    model.eval()
+    traced_model = torch.jit.trace(model, img)
+
+    mlmodel = ct.convert(
+        traced_model,
+        inputs=[ct.TensorType(shape=img.shape)],
+        convert_to='neuralnetwork',
+        #minimum_deployment_target=ct.target.iOS13
+    )
+    try:
+        mlmodel.save('{}/{}'.format(args.savedir, 'model.mlpackage'))#'path_to_save_your_model.mlmodel')
+    except Exception as e:
+        print_error_message('Error while saving the model: {}'.format(e))
+        return
+    print_info_message('CoreML model saved successfully in {}'.format(args.savedir))
+
+
+def main(args):
+    # print(os.path.abspath(os.path.join(args.data_path, "leftImg8bit")))
+    # print([x[0] for x in os.walk(os.path.join(args.data_path, "leftImg8bit"))])
+    # read all the images in the folder
+    if args.dataset == 'city':
+        image_path = os.path.join(args.data_path, "leftImg8bit", args.split, "*", "*.png")
+        image_list = glob.glob(image_path)
+        from data_loader.segmentation.cityscapes import CITYSCAPE_CLASS_LIST
+        seg_classes = len(CITYSCAPE_CLASS_LIST)
+    elif args.dataset == 'pascal':
+        from data_loader.segmentation.voc import VOC_CLASS_LIST
+        seg_classes = len(VOC_CLASS_LIST)
+        data_file = os.path.join(args.data_path, 'VOC2012', 'list', '{}.txt'.format(args.split))
+        if not os.path.isfile(data_file):
+            print_error_message('{} file does not exist'.format(data_file))
+        image_list = []
+        with open(data_file, 'r') as lines:
+            for line in lines:
+                rgb_img_loc = '{}/{}/{}'.format(args.data_path, 'VOC2012', line.split()[0])
+                if not os.path.isfile(rgb_img_loc):
+                    print_error_message('{} image file does not exist'.format(rgb_img_loc))
+                image_list.append(rgb_img_loc)
+    else:
+        print_error_message('{} dataset not yet supported'.format(args.dataset))
+
+    if args.model == 'espnetv2':
+        from model.segmentation.espnetv2 import espnetv2_seg
+        args.classes = seg_classes
+        model = espnetv2_seg(args)
+    elif args.model == 'dicenet':
+        from model.segmentation.dicenet import dicenet_seg
+        model = dicenet_seg(args, classes=seg_classes)
+    else:
+        print_error_message('{} network not yet supported'.format(args.model))
+        exit(-1)
+
+    # mdoel information
+    num_params = model_parameters(model)
+    flops = compute_flops(model, input=torch.Tensor(1, 3, args.im_size[0], args.im_size[1]))
+    print_info_message('FLOPs for an input of size {}x{}: {:.2f} million'.format(args.im_size[0], args.im_size[1], flops))
+    print_info_message('# of parameters: {}'.format(num_params))
+
+    if args.weights_test:
+        print_info_message('Loading model weights')
+        weight_dict = torch.load(args.weights_test, map_location=torch.device('cpu'))
+        model.load_state_dict(weight_dict)
+        print_info_message('Weight loaded successfully')
+    else:
+        print_error_message('weight file does not exist or not specified. Please check: {}', format(args.weights_test))
+
+    num_gpus = torch.cuda.device_count()
+    device = 'cuda' if num_gpus > 0 else 'cpu'
+    model = model.to(device=device)
+
+    # evaluate(args, model, image_list, device=device)
+    convert_to_coreml(model, args, device=device)
+
+
+if __name__ == '__main__':
+    from commons.general_details import segmentation_models, segmentation_datasets
+
+    parser = ArgumentParser()
+    # mdoel details
+    parser.add_argument('--model', default="espnetv2", choices=segmentation_models, help='Model name')
+    parser.add_argument('--weights-test', default='', help='Pretrained weights directory.')
+    parser.add_argument('--s', default=2.0, type=float, help='scale')
+    # dataset details
+    parser.add_argument('--data-path', default="", help='Data directory')
+    parser.add_argument('--dataset', default='city', choices=segmentation_datasets, help='Dataset name')
+    # input details
+    parser.add_argument('--im-size', type=int, nargs="+", default=[512, 256], help='Image size for testing (W x H)')
+    parser.add_argument('--split', default='val', choices=['val', 'test'], help='data split')
+    parser.add_argument('--model-width', default=224, type=int, help='Model width')
+    parser.add_argument('--model-height', default=224, type=int, help='Model height')
+    parser.add_argument('--channels', default=3, type=int, help='Input channels')
+    parser.add_argument('--num-classes', default=1000, type=int,
+                        help='ImageNet classes. Required for loading the base network')
+
+    args = parser.parse_args()
+
+    if not args.weights_test:
+        from model.weight_locations.segmentation import model_weight_map
+
+        model_key = '{}_{}'.format(args.model, args.s)
+        dataset_key = '{}_{}x{}'.format(args.dataset, args.im_size[0], args.im_size[1])
+        assert model_key in model_weight_map.keys(), '{} does not exist'.format(model_key)
+        assert dataset_key in model_weight_map[model_key].keys(), '{} does not exist'.format(dataset_key)
+        args.weights_test = model_weight_map[model_key][dataset_key]['weights']
+        if not os.path.isfile(args.weights_test):
+            print_error_message('weight file does not exist: {}'.format(args.weights_test))
+
+    # set-up results path
+    if args.dataset == 'city':
+        args.savedir = '{}_{}_{}/results'.format('results', args.dataset, args.split)
+    elif args.dataset == 'pascal':
+        args.savedir = '{}_{}/results/VOC2012/Segmentation/comp6_{}_cls'.format('results', args.dataset, args.split)
+    else:
+        print_error_message('{} dataset not yet supported'.format(args.dataset))
+
+    if not os.path.isdir(args.savedir):
+        os.makedirs(args.savedir)
+
+    # This key is used to load the ImageNet weights while training. So, set to empty to avoid errors
+    args.weights = ''
+
+    main(args)

core_ml_conversion2.py

@@ -0,0 +1,29 @@
+import torch
+import torch.onnx
+import coremltools as ct
+
+# Other guide: https://github.com/vincentfpgarcia/from-pytorch-to-coreml
+
+# Define a simple PyTorch model for demonstration
+class SimpleModel(torch.nn.Module):
+    def __init__(self):
+        super(SimpleModel, self).__init__()
+        self.conv1 = torch.nn.Conv2d(3, 16, kernel_size=3, stride=2, padding=1)
+        self.fc1 = torch.nn.Linear(16*112*112, 10)  # Assuming input size is (3, 224, 224)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = torch.nn.functional.relu(x)
+        x = torch.flatten(x, 1)
+        x = self.fc1(x)
+        return x
+
+# Instantiate and export the model
+pytorch_model = SimpleModel()
+dummy_input = torch.randn(1, 3, 224, 224)
+torch.onnx.export(pytorch_model, dummy_input, "simple_model.onnx")
+
+# Convert the ONNX model to CoreML
+onnx_model_path = "simple_model.onnx"
+coreml_model = ct.converters.onnx.convert(model=onnx_model_path)
+coreml_model.save("simple_model.mlmodel")

core_ml_conversion3.py

@@ -0,0 +1,18 @@
+# Paste Python code snippet here, complete with any required import statements.
+import coremltools as ct
+import torch
+import torchvision
+
+# Load PyTorch model (and perform tracing)
+torch_model = torchvision.models.mobilenet_v2()
+torch_model.eval() 
+
+example_input = torch.rand(1, 3, 256, 256)
+traced_model = torch.jit.trace(torch_model, example_input)
+
+# Convert using the same API. Note that we need to provide "inputs" for pytorch conversion.
+model_from_torch = ct.convert(traced_model,
+                              inputs=[ct.TensorType(name="input", 
+                                                    shape=example_input.shape)],
+                              debug=True)
+model_from_torch.save('core_ml_exp/model.mlpackage')

core_ml_conversion4.py

@@ -0,0 +1,17 @@
+# Paste Python code snippet here, complete with any required import statements.
+import coremltools as ct
+import torch
+import torchvision
+
+# Load PyTorch model (and perform tracing)
+torch_model = torchvision.models.mobilenet_v2()
+torch_model.eval() 
+
+example_input = torch.rand(1, 3, 256, 256)
+# traced_model = torch.jit.trace(torch_model, example_input)
+torch.onnx.export(torch_model, example_input, "core_ml_exp2/simple_model.onnx")
+
+# Convert the ONNX model to CoreML
+onnx_model_path = "core_ml_exp2/simple_model.onnx"
+coreml_model = ct.converters.onnx.convert(model=onnx_model_path)
+coreml_model.save("core_ml_exp2/simple_model.mlmodel")

core_ml_conversion5.py

@@ -0,0 +1,16 @@
+import torch
+import torchvision
+from onnx_coreml import convert
+import torch
+
+mlmodel = convert(model='./onnx_models/classification/espnetv2/espnetv2_s_1.0_imsize_224x224_imagenet.onnx', 
+                  minimum_ios_deployment_target='13')
+
+# Traceback (most recent call last):
+#   File "/home/ubuntu/ML/EdgeNets/core_ml_conversion5.py", line 3, in <module>
+#     from onnx_coreml import convert
+#   File "/home/ubuntu/anaconda3/envs/edgenets3/lib/python3.9/site-packages/onnx_coreml/__init__.py", line 6, in <module>
+#     from .converter import convert
+#   File "/home/ubuntu/anaconda3/envs/edgenets3/lib/python3.9/site-packages/onnx_coreml/converter.py", line 35, in <module>
+#     from coremltools.converters.nnssa.coreml.graph_pass.mlmodel_passes import remove_disconnected_layers, transform_conv_crop
+# ModuleNotFoundError: No module named 'coremltools.converters.nnssa'

core_ml_exp/model.mlpackage/Manifest.json

@@ -0,0 +1,18 @@
+{
+    "fileFormatVersion": "1.0.0",
+    "itemInfoEntries": {
+        "108e78b9-8dcd-437d-b93f-e6fb00364b7d": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Weights",
+            "name": "weights",
+            "path": "com.apple.CoreML/weights"
+        },
+        "c99da7fc-5195-43f0-a82e-c3070f27a17f": {
+            "author": "com.apple.CoreML",
+            "description": "CoreML Model Specification",
+            "name": "model.mlmodel",
+            "path": "com.apple.CoreML/model.mlmodel"
+        }
+    },
+    "rootModelIdentifier": "c99da7fc-5195-43f0-a82e-c3070f27a17f"
+}