initial

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text
+*.png filter=lfs diff=lfs merge=lfs -text

README.md

@@ -5,7 +5,7 @@ colorFrom: yellow
 colorTo: green
 sdk: gradio
 sdk_version: 5.6.0
-app_file: app.py
+app_file: main.py
 pinned: false
 license: apache-2.0
 short_description: Root analysis using deep learning

dependecies/segroot/binarize_crop.py

@@ -0,0 +1,72 @@
+from pathlib import Path
+from PIL import Image
+import matplotlib.pyplot as plt
+import skimage.io as io
+from skimage.morphology import dilation
+import pickle
+import numpy as np
+import argparse
+from dataloader import pad_pair_256
+
+parser = argparse.ArgumentParser()  
+parser.add_argument(
+    "--dilate",
+    default=0,
+    type=int,
+    help="dilation degree of masks")
+
+args = parser.parse_args()
+data_dir = Path('../data/data_raw')
+mask_dir = Path('../data/masks')
+mask_dir.mkdir(exist_ok=True, parents=True)
+
+imgs = sorted(list(data_dir.glob('*Untitled.jpg')))
+print('original images count : ', len(imgs))
+masks = sorted(list(data_dir.glob('*Untitled-mask.jpg')))
+print('original masks count : ', len(masks))
+
+# generate binary masks for every annotated images
+for m in masks:
+    mask = io.imread(m.as_posix(), as_gray=True)
+    # binarize
+    mask[mask > 0.5 ] = 1.0
+    mask[mask <= 0.5] = 0.0
+    for i in range(args.dilate):
+        mask = dilation(mask)
+    print('binary masks dilated !!!')
+    plt.imsave((mask_dir / m.parts[-1]).as_posix(), mask, cmap='gray')
+print('binary masks generated !!!')
+
+# save idx info in a dictionary
+info_dict = {k: v.parts[-1] for k, v in enumerate(imgs)}
+with open('../data/info.pkl', 'wb') as handle:
+    pickle.dump(info_dict, handle)
+print('index info saved!!!')
+
+# crop the padded image to generate 256*256 subimages
+new_masks = sorted(list(mask_dir.glob('*Untitled-mask.jpg')))
+print('new_mask length : ',len(new_masks))
+
+subimg_path = Path('../data/subimg')
+subimg_path.mkdir(exist_ok=True, parents=True)
+submask_path = Path('../data/submask')
+submask_path.mkdir(exist_ok=True, parents=True)
+
+for idx, (mask_path, img_path) in enumerate(zip(new_masks, imgs)):
+    mask = Image.open(mask_path)
+    img = Image.open(img_path)
+    new_img, new_mask = pad_pair_256(img, mask)
+    new_img, new_mask = np.array(new_img), np.array(new_mask)
+    # padded shape (2560, 2304)
+    w, h, _ = new_img.shape
+    for i in range(int(w/256)):
+        for j in range(int(h/256)):
+            subimg = new_img[i*256:(i+1)*256, j*256:(j+1)*256, :]
+            subimg_fn = '{}/{}-{}-{}.png'.format(
+                Path('../data/subimg').as_posix(), idx, i, j)
+            plt.imsave(subimg_fn, subimg)
+            submask_fn = '{}/{}-{}-{}.png'.format(
+                Path('../data/submask').as_posix(), idx, i, j)
+            submask = new_mask[i*256:(i+1)*256, j*256:(j+1)*256]
+            plt.imsave(submask_fn, submask, cmap='gray')
+    print('No.{} image & mask cropped!!!'.format(idx))

dependecies/segroot/dataloader.py

@@ -0,0 +1,151 @@
+import os
+import itertools
+import pickle
+import torch
+from torchvision import models
+from pathlib import Path
+from PIL import Image
+from torch.utils.data import Dataset, DataLoader, Sampler
+
+import dependecies.segroot.paired_transforms_pt04 as p_tr
+
+train_transform = p_tr.Compose([
+    p_tr.RandomCrop(256),
+    p_tr.RandomRotation((90, 90)),
+    p_tr.RandomRotation((180, 180)),
+    p_tr.RandomRotation((270, 270)),
+    p_tr.RandomHorizontalFlip(),
+    p_tr.RandomVerticalFlip(),
+    p_tr.ToTensor()
+])
+
+# normalize = p_tr.Normalize([0.35042979, 0.44016893, 0.2340332],
+#                            [0.20999724, 0.25972678, 0.13885915])
+normalize = p_tr.Normalize([0.5, 0.5, 0.5],
+                           [0.5, 0.5, 0.5])
+
+
+def pad_pair_256(image, gt):
+    w, h = image.size
+    new_w = ((w - 1) // 256 + 1) * 256
+    new_h = ((h - 1) // 256 + 1) * 256
+    new_image = Image.new("RGB", (new_w, new_h))
+    new_image.paste(image, ((new_w - w) // 2, (new_h - h) // 2))
+    new_gt = Image.new("L", (new_w, new_h))
+    new_gt.paste(gt, ((new_w - w) // 2, (new_h - h) // 2))
+    return new_image, new_gt
+
+
+def convert_png(image, gt):
+    new_image = Image.new('RGB', (256, 256))
+    new_image.paste(image)
+    new_gt = Image.new('L', (256, 256))
+    new_gt.paste(gt)
+    return new_image, new_gt
+
+
+def get_paths(root_dir, im_ids):
+    imgs = []
+    for i in im_ids:
+        tmp = Path(root_dir).glob('*{}-*.png'.format(i))
+        tmp = [p for p in tmp if p.parts[-1].startswith(str(i)+'-')]
+        imgs = imgs + list(tmp)
+    return imgs
+
+
+class LoopSampler(Sampler):
+    def __init__(self, data_source):
+        self.data_source = data_source
+
+    def __iter__(self):
+        return itertools.cycle(range(len(self.data_source)))
+
+    def __len__(self):
+        return len(self.data_source)
+
+
+class TrainDataset(Dataset):
+    def __init__(self, im_ids):
+        self.root_dir = '../data/data_raw'
+        self.mask_dir = '../data/mask'
+        self.im_ids = im_ids
+        with open('../data/info.pkl', 'rb') as handle:
+            self.info = pickle.load(handle)
+        self.fns = [self.info[im_id] for im_id in im_ids]
+
+    def __getitem__(self, index):
+        im_fn = self.fns[index]
+        im_name = os.path.join(self.root_dir, im_fn)
+        gt_name = os.path.join(
+            self.mask_dir, im_fn.split('.jpg')[0] + '-mask.jpg')
+        image = Image.open(im_name)
+        gt = Image.open(gt_name)
+        image, gt = pad_pair_256(image, gt)
+
+        image, gt = train_transform(image, gt)
+        image = normalize(image)
+
+        return image, gt
+
+    def __len__(self):
+        return len(self.im_ids)
+
+
+class StaticTrainDataset(Dataset):
+    def __init__(self, im_ids):
+        self.subimgs = sorted(get_paths('../data/subimg', im_ids))
+        self.submasks = sorted(get_paths('../data/submask', im_ids))
+        self.im_ids = im_ids
+
+    def __getitem__(self, index):
+        im_name = self.subimgs[index]
+        gt_name = self.submasks[index]
+        image = Image.open(im_name)
+        gt = Image.open(gt_name)
+        image, gt = convert_png(image, gt)
+
+        image, gt = train_transform(image, gt)
+        image = normalize(image)
+
+        return image, gt
+
+    def __len__(self):
+        return len(self.im_ids * 90)
+
+
+class TrainDataLoader():
+    def __init__(self, dataset, batch_size, num_workers=0):
+        self.dataset = dataset
+        self.dataloader = DataLoader(self.dataset, batch_size=batch_size,
+                                     num_workers=num_workers, sampler=LoopSampler(self.dataset))
+        self.dl = iter(self.dataloader)
+
+    def next_batch(self):
+        image, gt = next(self.dl)
+        return image, gt
+
+
+class TestDataset(Dataset):
+    def __init__(self, im_ids):
+        self.root_dir = '../data/data_raw'
+        self.mask_dir = '../data/masks'
+        with open('../data/info.pkl', 'rb') as handle:
+            self.info = pickle.load(handle)
+        self.im_ids = im_ids
+        self.fns = [self.info[im_id] for im_id in im_ids]
+
+    def __getitem__(self, index):
+        im_fn = self.fns[index]
+        im_name = os.path.join(self.root_dir, im_fn)
+        gt_name = os.path.join(
+                self.mask_dir, im_fn.split('.jpg')[0] + '-mask.jpg')
+        image = Image.open(im_name)
+        gt = Image.open(gt_name)
+        image, gt = pad_pair_256(image, gt)
+
+        image, gt = p_tr.ToTensor()(image, gt)
+        image = normalize(image)
+        return image, gt
+
+    def __len__(self):
+        return len(self.fns)

dependecies/segroot/main_segroot.py

@@ -0,0 +1,112 @@
+import torch
+from torch.utils.data import DataLoader
+import numpy as np
+import random
+from tqdm import tqdm
+import argparse
+
+from model import SegRoot
+from dataloader import StaticTrainDataset, TestDataset, TrainDataset, LoopSampler
+from utils import (
+    dice_score,
+    init_weights,
+    evaluate,
+    get_ids,
+    load_vgg16,
+    set_random_seed,
+)
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--seed", default=42, type=int, help="set random seed")
+parser.add_argument("--width", default=8, type=int, help="width of SegRoot")
+parser.add_argument("--depth", default=5, type=int, help="depth of SegRoot")
+parser.add_argument("--bs", default=64, type=int, help="batch size of dataloaders")
+parser.add_argument("--lr", default=1e-2, type=float, help="learning rate")
+parser.add_argument("--epochs", default=200, type=int, help="max epochs of training")
+parser.add_argument(
+    "--verbose", default=5, type=int, help="intervals to save and validate model"
+)
+parser.add_argument(
+    "--dynamic", action="store_true", help="use dynamic sub-images during training"
+)
+
+
+def train_one_epoch(model, train_iter, optimizer, device):
+    model.train()
+    for p in model.parameters():
+        p.requires_grad = True
+    for x, y in train_iter:
+        x, y = x.to(device), y.to(device)
+        bs = x.shape[0]
+        optimizer.zero_grad()
+        y_pred = model(x)
+        loss = 1 - dice_score(y, y_pred)
+        loss = torch.sum(loss) / bs
+        loss.backward()
+        optimizer.step()
+
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+    seed = args.seed
+    bs = args.bs
+    lr = args.lr
+    width = args.width
+    depth = args.depth
+    epochs = args.epochs
+    verbose = args.verbose
+
+    # set random seed
+    set_random_seed(seed)
+    # define the device for training
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    # get training ids
+    train_ids, valid_ids, test_ids = get_ids(65)
+    # define dataloaders
+    if args.dynamic:
+        train_data = TrainDataset(train_ids)
+        train_iter = DataLoader(
+            train_data, batch_size=bs, num_workers=6, sampler=LoopSampler
+        )
+    else:
+        train_data = StaticTrainDataset(train_ids)
+        train_iter = DataLoader(train_data, batch_size=bs, num_workers=6, shuffle=True)
+
+    train_tdata = TestDataset(train_ids)
+    valid_tdata = TestDataset(valid_ids)
+    test_tdata = TestDataset(test_ids)
+
+    # define model
+    model = SegRoot(width, depth).to(device)
+    model = model.apply(init_weights)
+
+    # define optimizer and lr_scheduler
+    optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=5e-4)
+    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
+        optimizer, mode="max", factor=0.5, verbose=True, patience=5
+    )
+
+    print(f"Start training SegRoot-({width},{depth}))......")
+    print(f"Random seed is {seed}, batch size is {bs}......")
+    print(f"learning rate is {lr}, max epochs is {epochs}......")
+    best_valid = float("-inf")
+    for epoch in tqdm(range(epochs)):
+        train_one_epoch(model, train_iter, optimizer, device)
+        if epoch % verbose == 0:
+            train_dice = evaluate(model, train_tdata, device)
+            valid_dice = evaluate(model, valid_tdata, device)
+            scheduler.step(valid_dice)
+            print(
+                "Epoch {:05d}, train dice: {:.4f}, valid dice: {:.4f}".format(
+                    epoch, train_dice, valid_dice
+                )
+            )
+            if valid_dice > best_valid:
+                best_valid = valid_dice
+                test_dice = evaluate(model, test_tdata, device)
+                print("New best validation, test dice: {:.4f}".format(test_dice))
+                torch.save(
+                    model.state_dict(),
+                    f"../weights/best_segroot-({args.width},{args.depth}).pt",
+                )

dependecies/segroot/model.py

@@ -0,0 +1,124 @@
+import torch
+import torchvision
+from torch import nn
+from torch.nn import functional as F
+
+
+class ConvBNRelu(nn.Module):
+    def __init__(self, in_ch, out_ch):
+        super(ConvBNRelu, self).__init__()
+        self.conv = nn.Conv2d(in_ch, out_ch, 3, padding=1, bias=True)
+        self.bn = nn.BatchNorm2d(out_ch)
+        self.activation = nn.ReLU()
+
+    def forward(self, x):
+        x = self.conv(x)
+        x = self.bn(x)
+        x = self.activation(x)
+        # print(x.shape)
+        return x
+
+
+class FirstBlock(nn.Module):
+    def __init__(self, in_ch, out_ch):
+        super(FirstBlock, self).__init__()
+        self.conv1 = ConvBNRelu(in_ch, out_ch)
+        self.conv2 = ConvBNRelu(out_ch, out_ch)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.conv2(x)
+        return x
+
+
+class DownBlock(nn.Module):
+    def __init__(self, in_ch, out_ch):
+        super(DownBlock, self).__init__()
+        self.conv1 = ConvBNRelu(in_ch, out_ch)
+        self.conv2 = ConvBNRelu(out_ch, out_ch)
+
+    def forward(self, x):
+        x = F.max_pool2d(x,kernel_size=2,stride=2)
+        x = self.conv1(x)
+        x = self.conv2(x)
+        return x
+
+class Encoder(nn.Module):
+    def __init__(self, in_ch, out_ch, block_num=2):
+        super(Encoder, self).__init__()
+        layers = []
+        layers += [ConvBNRelu(in_ch, out_ch)]
+        for i in range(block_num-1):
+            layers += [ConvBNRelu(out_ch, out_ch)]
+        # layers += [nn.Dropout2d(0.5)]
+        self.features = nn.Sequential(*layers)
+    
+    def forward(self, x):
+        x = self.features(x)
+        x, indices = F.max_pool2d(x, kernel_size=2, stride=2, return_indices=True) 
+        return x, indices
+
+class Decoder(nn.Module):
+    def __init__(self, in_ch, out_ch, block_num=2):
+        super(Decoder, self).__init__()
+        layers = []
+        layers += [ConvBNRelu(in_ch, out_ch)]
+        for i in range(block_num-1):
+            layers += [ConvBNRelu(out_ch, out_ch)]
+        # layers += [nn.Dropout2d(0.5)]
+        self.features = nn.Sequential(*layers)
+        
+    def forward(self, x, indices):
+        x = F.max_unpool2d(x, indices=indices, kernel_size=2, stride=2)
+        x = self.features(x)
+        return x
+
+class SegRoot(nn.Module):
+    def __init__(self, width=8, depth=5, num_classes=2):
+        super(SegRoot, self).__init__()
+        chs = []
+        for i in range(depth-1):
+            chs.append(width * (2**i))
+        chs.append(chs[-1])
+        self.e_ch_info = [3,] + chs
+        self.e_bl_info = [2,2,3,3]
+        for _ in range(depth - 4):
+            self.e_bl_info += [3,]
+        self.d_ch_info = chs[::-1] + [4,]
+        self.d_bl_info = self.e_bl_info[::-1]
+        # using same setup with Unet
+        if width == 4:
+            self.e_ch_info = [3,4,8,16,32,64]
+            self.d_ch_info = [64,32,16,8,4,4]
+        self.num_classes = num_classes
+        self.encoders = nn.ModuleList()
+        self.decoders = nn.ModuleList()
+        
+        for i in range(1,len(self.e_ch_info)):
+            self.encoders.append(Encoder(self.e_ch_info[i-1], self.e_ch_info[i], self.e_bl_info[i-1]))
+            self.decoders.append(Decoder(self.d_ch_info[i-1], self.d_ch_info[i], self.d_bl_info[i-1]))
+        
+        # self.classifier = nn.Conv2d(self.d_ch_info[-1], num_classes, kernel_size=3, padding=1)
+        self.classifier = nn.Conv2d(self.d_ch_info[-1], 1, 1)
+        
+    def forward(self, x):
+        indices = []
+        bs = x.shape[0]
+        for i in range(len(self.e_bl_info)):
+            x, ind = self.encoders[i](x)
+            indices.append(ind)
+            
+        indices = indices[::-1]    
+        for i in range(len(self.e_bl_info)):
+            x = self.decoders[i](x, indices[i])
+        
+        x = self.classifier(x)
+        # x = F.softmax(x,dim=1)
+        x = torch.sigmoid(x)
+        return x
+
+
+if __name__ == '__main__':
+    x = torch.zeros((1, 3, 256, 256))
+    net = SegRoot(8,5)
+    print(net(x).shape)

dependecies/segroot/paired_transforms_pt04.py

@@ -0,0 +1,1027 @@
+from __future__ import division
+import torch
+import math
+import random
+from PIL import Image, ImageOps, ImageEnhance
+try:
+    import accimage
+except ImportError:
+    accimage = None
+import numpy as np
+import numbers
+import types
+import collections
+import warnings
+
+from torchvision.transforms import functional as F
+
+__all__ = ["Compose", "ToTensor", "ToPILImage", "Normalize", "Resize", "Scale", "CenterCrop", "Pad",
+           "Lambda", "RandomApply", "RandomChoice", "RandomOrder", "RandomCrop", "RandomHorizontalFlip",
+           "RandomVerticalFlip", "RandomResizedCrop", "RandomSizedCrop", "FiveCrop", "TenCrop", "LinearTransformation",
+           "ColorJitter", "RandomRotation", "RandomAffine", "Grayscale", "RandomGrayscale"]
+
+_pil_interpolation_to_str = {
+    Image.NEAREST: 'PIL.Image.NEAREST',
+    Image.BILINEAR: 'PIL.Image.BILINEAR',
+    Image.BICUBIC: 'PIL.Image.BICUBIC',
+    Image.LANCZOS: 'PIL.Image.LANCZOS',
+}
+
+
+class Compose(object):
+    """Composes several transforms together.
+    Args:
+        transforms (list of ``Transform`` objects): list of transforms to compose.
+    Example:
+        >>> transforms.Compose([
+        >>>     transforms.CenterCrop(10),
+        >>>     transforms.ToTensor(),
+        >>> ])
+    """
+
+    def __init__(self, transforms):
+        self.transforms = transforms
+
+    def __call__(self, img, target = None):
+        if target is not None:
+            for t in self.transforms:
+                img, target = t(img, target)
+            return img, target
+
+        for t in self.transforms:
+            img = t(img)
+        return img
+
+    def __repr__(self):
+        format_string = self.__class__.__name__ + '('
+        for t in self.transforms:
+            format_string += '\n'
+            format_string += '    {0}'.format(t)
+        format_string += '\n)'
+        return format_string
+
+
+class ToTensor(object):
+    """Convert a ``PIL Image`` or ``numpy.ndarray`` to tensor.
+    Converts a PIL Image or numpy.ndarray (H x W x C) in the range
+    [0, 255] to a torch.FloatTensor of shape (C x H x W) in the range [0.0, 1.0].
+    """
+
+    def __call__(self, pic, pic2=None):
+        """
+        Args:
+            pic (PIL Image or numpy.ndarray): Image to be converted to tensor.
+            pic2 (PIL Image): (optional) Second image to be converted also.
+        Returns:
+            Tensor(s): Converted image(s).
+        """
+        if pic2 is not None:
+            return F.to_tensor(pic), F.to_tensor(pic2)
+        return F.to_tensor(pic)
+
+    def __repr__(self):
+        return self.__class__.__name__ + '()'
+
+
+class ToPILImage(object):
+    """Convert a tensor or an ndarray to PIL Image.
+    Converts a torch.*Tensor of shape C x H x W or a numpy ndarray of shape
+    H x W x C to a PIL Image while preserving the value range.
+    Args:
+        mode (`PIL.Image mode`_): color space and pixel depth of input data (optional).
+            If ``mode`` is ``None`` (default) there are some assumptions made about the input data:
+            1. If the input has 3 channels, the ``mode`` is assumed to be ``RGB``.
+            2. If the input has 4 channels, the ``mode`` is assumed to be ``RGBA``.
+            3. If the input has 1 channel, the ``mode`` is determined by the data type (i,e,
+            ``int``, ``float``, ``short``).
+    .. _PIL.Image mode: http://pillow.readthedocs.io/en/3.4.x/handbook/concepts.html#modes
+    """
+    def __init__(self, mode=None):
+        self.mode = mode
+
+    def __call__(self, pic, pic2=None):
+        """
+        Args:
+            pic (Tensor or numpy.ndarray): Image to be converted to PIL Image.
+        Returns:
+            PIL Image: Image converted to PIL Image.
+        """
+        if pic2 is not None:
+            return F.to_pil_image(pic), F.to_pil_image(pic2)
+        return F.to_pil_image(pic, self.mode)
+
+    def __repr__(self):
+        format_string = self.__class__.__name__ + '('
+        if self.mode is not None:
+            format_string += 'mode={0}'.format(self.mode)
+        format_string += ')'
+        return format_string
+
+
+class Normalize(object):
+    """Normalize a tensor image with mean and standard deviation.
+    Given mean: ``(M1,...,Mn)`` and std: ``(S1,..,Sn)`` for ``n`` channels, this transform
+    will normalize each channel of the input ``torch.*Tensor`` i.e.
+    ``input[channel] = (input[channel] - mean[channel]) / std[channel]``
+    Args:
+        mean (sequence): Sequence of means for each channel.
+        std (sequence): Sequence of standard deviations for each channel.
+    """
+
+    def __init__(self, mean, std):
+        self.mean = mean
+        self.std = std
+
+    def __call__(self, tensor):
+        """
+        Args:
+            tensor (Tensor): Tensor image of size (C, H, W) to be normalized.
+        Returns:
+            Tensor: Normalized Tensor image.
+        """
+        return F.normalize(tensor, self.mean, self.std)
+
+    def __repr__(self):
+        return self.__class__.__name__ + '(mean={0}, std={1})'.format(self.mean, self.std)
+
+
+class Resize(object):
+    """Resize the input PIL Image to the given size.
+    Args:
+        size (sequence or int): Desired output size. If size is a sequence like
+            (h, w), output size will be matched to this. If size is an int,
+            smaller edge of the image will be matched to this number.
+            i.e, if height > width, then image will be rescaled to
+            (size * height / width, size)
+        interpolation (int, optional): Desired interpolation. Default is
+            ``PIL.Image.BILINEAR``
+        interpolation_tg (int, optional): Desired interpolation for target. Default is
+            ``PIL.Image.NEAREST``
+    """
+
+    def __init__(self, size, interpolation=Image.BILINEAR, interpolation_tg = Image.NEAREST):
+        assert isinstance(size, int) or (isinstance(size, collections.Iterable) and len(size) == 2)
+        self.size = size
+        self.interpolation = interpolation
+        self.interpolation_tg = interpolation_tg
+
+    def __call__(self, img, target = None):
+        """
+        Args:
+            img (PIL Image): Image to be scaled.
+            target (PIL Image): (optional) Target to be scaled
+        Returns:
+            PIL Image: Rescaled image(s).
+        """
+        if target is not None:
+            return F.resize(img, self.size, self.interpolation), F.resize(target, self.size, self.interpolation_tg)
+        return F.resize(img, self.size, self.interpolation)
+
+    def __repr__(self):
+        interpolate_str = _pil_interpolation_to_str[self.interpolation]
+        return self.__class__.__name__ + '(size={0}, interpolation={1})'.format(self.size, interpolate_str)
+
+
+class Scale(Resize):
+    """
+    Note: This transform is deprecated in favor of Resize.
+    """
+    def __init__(self, *args, **kwargs):
+        warnings.warn("The use of the transforms.Scale transform is deprecated, " +
+                      "please use transforms.Resize instead.")
+        super(Scale, self).__init__(*args, **kwargs)
+
+
+class CenterCrop(object):
+    """Crops the given PIL Image at the center.
+    Args:
+        size (sequence or int): Desired output size of the crop. If size is an
+            int instead of sequence like (h, w), a square crop (size, size) is
+            made.
+    """
+
+    def __init__(self, size):
+        if isinstance(size, numbers.Number):
+            self.size = (int(size), int(size))
+        else:
+            self.size = size
+
+    def __call__(self, img, target=None):
+        """
+        Args:
+            img (PIL Image): Image to be cropped.
+            target (PIL Image): (optional) Target to be cropped
+        Returns:
+            PIL Image: Cropped image(s).
+        """
+        if target is not None:
+            return F.center_crop(img, self.size), F.center_crop(target, self.size)
+        return F.center_crop(img, self.size)
+
+    def __repr__(self):
+        return self.__class__.__name__ + '(size={0})'.format(self.size)
+
+
+class Pad(object):
+    """Pad the given PIL Image on all sides with the given "pad" value.
+    Args:
+        padding (int or tuple): Padding on each border. If a single int is provided this
+            is used to pad all borders. If tuple of length 2 is provided this is the padding
+            on left/right and top/bottom respectively. If a tuple of length 4 is provided
+            this is the padding for the left, top, right and bottom borders
+            respectively.
+        fill: Pixel fill value for constant fill. Default is 0. If a tuple of
+            length 3, it is used to fill R, G, B channels respectively.
+            This value is only used when the padding_mode is constant
+        padding_mode: Type of padding. Should be: constant, edge, reflect or symmetric. Default is constant.
+            constant: pads with a constant value, this value is specified with fill
+            edge: pads with the last value at the edge of the image
+            reflect: pads with reflection of image (without repeating the last value on the edge)
+                padding [1, 2, 3, 4] with 2 elements on both sides in reflect mode
+                will result in [3, 2, 1, 2, 3, 4, 3, 2]
+            symmetric: pads with reflection of image (repeating the last value on the edge)
+                padding [1, 2, 3, 4] with 2 elements on both sides in symmetric mode
+                will result in [2, 1, 1, 2, 3, 4, 4, 3]
+    """
+
+    def __init__(self, padding, fill=0, padding_mode='constant'):
+        assert isinstance(padding, (numbers.Number, tuple))
+        assert isinstance(fill, (numbers.Number, str, tuple))
+        assert padding_mode in ['constant', 'edge', 'reflect', 'symmetric']
+        if isinstance(padding, collections.Sequence) and len(padding) not in [2, 4]:
+            raise ValueError("Padding must be an int or a 2, or 4 element tuple, not a " +
+                             "{} element tuple".format(len(padding)))
+
+        self.padding = padding
+        self.fill = fill
+        self.padding_mode = padding_mode
+
+    def __call__(self, img):
+        """
+        Args:
+            img (PIL Image): Image to be padded.
+        Returns:
+            PIL Image: Padded image.
+        """
+        return F.pad(img, self.padding, self.fill, self.padding_mode)
+
+    def __repr__(self):
+        return self.__class__.__name__ + '(padding={0}, fill={1}, padding_mode={2})'.\
+            format(self.padding, self.fill, self.padding_mode)
+
+
+class Lambda(object):
+    """Apply a user-defined lambda as a transform.
+    Args:
+        lambd (function): Lambda/function to be used for transform.
+    """
+
+    def __init__(self, lambd):
+        assert isinstance(lambd, types.LambdaType)
+        self.lambd = lambd
+
+    def __call__(self, img):
+        return self.lambd(img)
+
+    def __repr__(self):
+        return self.__class__.__name__ + '()'
+
+
+class RandomTransforms(object):
+    """Base class for a list of transformations with randomness
+    Args:
+        transforms (list or tuple): list of transformations
+    """
+
+    def __init__(self, transforms):
+        assert isinstance(transforms, (list, tuple))
+        self.transforms = transforms
+
+    def __call__(self, *args, **kwargs):
+        raise NotImplementedError()
+
+    def __repr__(self):
+        format_string = self.__class__.__name__ + '('
+        for t in self.transforms:
+            format_string += '\n'
+            format_string += '    {0}'.format(t)
+        format_string += '\n)'
+        return format_string
+
+
+class RandomApply(RandomTransforms):
+    """Apply randomly a list of transformations with a given probability
+    Args:
+        transforms (list or tuple): list of transformations
+        p (float): probability
+    """
+
+    def __init__(self, transforms, p=0.5):
+        super(RandomApply, self).__init__(transforms)
+        self.p = p
+
+    def __call__(self, img):
+        if self.p < random.random():
+            return img
+        for t in self.transforms:
+            img = t(img)
+        return img
+
+    def __repr__(self):
+        format_string = self.__class__.__name__ + '('
+        format_string += '\n    p={}'.format(self.p)
+        for t in self.transforms:
+            format_string += '\n'
+            format_string += '    {0}'.format(t)
+        format_string += '\n)'
+        return format_string
+
+
+class RandomOrder(RandomTransforms):
+    """Apply a list of transformations in a random order
+    """
+    def __call__(self, img):
+        order = list(range(len(self.transforms)))
+        random.shuffle(order)
+        for i in order:
+            img = self.transforms[i](img)
+        return img
+
+
+class RandomChoice(RandomTransforms):
+    """Apply single transformation randomly picked from a list
+    """
+    def __call__(self, img):
+        t = random.choice(self.transforms)
+        return t(img)
+
+
+class RandomCrop(object):
+    """Crop the given PIL Image at a random location.
+    Args:
+        size (sequence or int): Desired output size of the crop. If size is an
+            int instead of sequence like (h, w), a square crop (size, size) is
+            made.
+        padding (int or sequence, optional): Optional padding on each border
+            of the image. Default is 0, i.e no padding. If a sequence of length
+            4 is provided, it is used to pad left, top, right, bottom borders
+            respectively.
+        pad_if_needed (boolean): It will pad the image if smaller than the
+            desired size to avoid raising an exception.
+    """
+
+    def __init__(self, size, padding=0, pad_if_needed=False):
+        if isinstance(size, numbers.Number):
+            self.size = (int(size), int(size))
+        else:
+            self.size = size
+        self.padding = padding
+        self.pad_if_needed = pad_if_needed
+
+    @staticmethod
+    def get_params(img, output_size):
+        """Get parameters for ``crop`` for a random crop.
+        Args:
+            img (PIL Image): Image to be cropped.
+            output_size (tuple): Expected output size of the crop.
+        Returns:
+            tuple: params (i, j, h, w) to be passed to ``crop`` for random crop.
+        """
+        w, h = img.size
+        th, tw = output_size
+        if w == tw and h == th:
+            return 0, 0, h, w
+
+        i = random.randint(0, h - th)
+        j = random.randint(0, w - tw)
+        return i, j, th, tw
+
+    def __call__(self, img, target = None):
+        """
+        Args:
+            img (PIL Image): Image to be cropped.
+            target (PIL Image): (optional) Target to be cropped
+        Returns:
+            PIL Images: Cropped image(s).
+        """
+        if self.padding > 0:
+            img = F.pad(img, self.padding)
+            if target is not None:
+                target = F.pad(target, self.padding)
+
+        # pad the width if needed
+        if self.pad_if_needed and img.size[0] < self.size[1]:
+            img = F.pad(img, (int((1 + self.size[1] - img.size[0]) / 2), 0))
+            if target is not None:
+                target = F.pad(target, (int((1 + self.size[1] - target.size[0]) / 2), 0))
+        # pad the height if needed
+        if self.pad_if_needed and img.size[1] < self.size[0]:
+            img = F.pad(img, (0, int((1 + self.size[0] - img.size[1]) / 2)))
+            if target is not None:
+                target = F.pad(target, (0, int((1 + self.size[0] - target.size[1]) / 2)))
+        i, j, h, w = self.get_params(img, self.size)
+
+        if target is not None:
+            return F.crop(img, i, j, h, w), F.crop(target, i, j, h, w)
+        else:
+            return F.crop(img, i, j, h, w)
+
+    def __repr__(self):
+        return self.__class__.__name__ + '(size={0}, padding={1})'.format(self.size, self.padding)
+
+
+class RandomHorizontalFlip(object):
+    """Horizontally flip the given PIL Image randomly with a given probability.
+    Args:
+        p (float): probability of the image being flipped. Default value is 0.5
+    """
+
+    def __init__(self, p=0.5):
+        self.p = p
+
+    def __call__(self, img, target=None):
+        """
+        Args:
+            img (PIL Image): Image to be flipped.
+            target (PIL Image): (optional) Target to be flipped
+        Returns:
+            PIL Image: Randomly flipped image(s).
+        """
+        if random.random() < self.p:
+            if target is not None:
+                return F.hflip(img), F.hflip(target)
+            else:
+                return F.hflip(img)
+
+        if target is not None:
+            return img, target
+        return img
+
+    def __repr__(self):
+        return self.__class__.__name__ + '(p={})'.format(self.p)
+
+
+class RandomVerticalFlip(object):
+    """Vertically flip the given PIL Image randomly with a given probability.
+    Args:
+        p (float): probability of the image being flipped. Default value is 0.5
+    """
+
+    def __init__(self, p=0.5):
+        self.p = p
+
+    def __call__(self, img, target=None):
+        """
+        Args:
+            img (PIL Image): Image to be flipped.
+            target (PIL Image): (optional) Target to be flipped
+        Returns:
+            PIL Image: Randomly flipped image(s).
+        """
+        if random.random() < self.p:
+            if target is not None:
+                return F.vflip(img), F.vflip(target)
+            else:
+                return F.vflip(img)
+
+        if target is not None:
+            return img, target
+        return img
+
+    def __repr__(self):
+        return self.__class__.__name__ + '(p={})'.format(self.p)
+
+
+class RandomResizedCrop(object):
+    """Crop the given PIL Image to random size and aspect ratio.
+    A crop of random size (default: of 0.08 to 1.0) of the original size and a random
+    aspect ratio (default: of 3/4 to 4/3) of the original aspect ratio is made. This crop
+    is finally resized to given size.
+    This is popularly used to train the Inception networks.
+    Args:
+        size: expected output size of each edge
+        scale: range of size of the origin size cropped
+        ratio: range of aspect ratio of the origin aspect ratio cropped
+        interpolation: Default: PIL.Image.BILINEAR
+    """
+
+    def __init__(self, size, scale=(0.08, 1.0), ratio=(3. / 4., 4. / 3.),
+                 interpolation=Image.BILINEAR, interpolation_tg = Image.NEAREST):
+        self.size = (size, size)
+        self.interpolation = interpolation
+        self.interpolation_tg = interpolation_tg
+        self.scale = scale
+        self.ratio = ratio
+
+    @staticmethod
+    def get_params(img, scale, ratio):
+        """Get parameters for ``crop`` for a random sized crop.
+        Args:
+            img (PIL Image): Image to be cropped.
+            scale (tuple): range of size of the origin size cropped
+            ratio (tuple): range of aspect ratio of the origin aspect ratio cropped
+        Returns:
+            tuple: params (i, j, h, w) to be passed to ``crop`` for a random
+                sized crop.
+        """
+        for attempt in range(10):
+            area = img.size[0] * img.size[1]
+            target_area = random.uniform(*scale) * area
+            aspect_ratio = random.uniform(*ratio)
+
+            w = int(round(math.sqrt(target_area * aspect_ratio)))
+            h = int(round(math.sqrt(target_area / aspect_ratio)))
+
+            if random.random() < 0.5:
+                w, h = h, w
+
+            if w <= img.size[0] and h <= img.size[1]:
+                i = random.randint(0, img.size[1] - h)
+                j = random.randint(0, img.size[0] - w)
+                return i, j, h, w
+
+        # Fallback
+        w = min(img.size[0], img.size[1])
+        i = (img.size[1] - w) // 2
+        j = (img.size[0] - w) // 2
+        return i, j, w, w
+
+    def __call__(self, img, target = None):
+        """
+        Args:
+            img (PIL Image): Image to be cropped and resized.
+            target (PIL Image): (optional) Target to be cropped and resized.
+        Returns:
+            PIL Image: Randomly cropped and resized image(s).
+        """
+        i, j, h, w = self.get_params(img, self.scale, self.ratio)
+        if target is not None:            
+            return F.resized_crop(img, i, j, h, w, self.size, self.interpolation), \
+                   F.resized_crop(target, i, j, h, w, self.size, self.interpolation_tg)
+        return F.resized_crop(img, i, j, h, w, self.size, self.interpolation)
+
+    def __repr__(self):
+        interpolate_str = _pil_interpolation_to_str[self.interpolation]
+        format_string = self.__class__.__name__ + '(size={0}'.format(self.size)
+        format_string += ', scale={0}'.format(tuple(round(s, 4) for s in self.scale))
+        format_string += ', ratio={0}'.format(tuple(round(r, 4) for r in self.ratio))
+        format_string += ', interpolation={0})'.format(interpolate_str)
+        return format_string
+
+
+class RandomSizedCrop(RandomResizedCrop):
+    """
+    Note: This transform is deprecated in favor of RandomResizedCrop.
+    """
+    def __init__(self, *args, **kwargs):
+        warnings.warn("The use of the transforms.RandomSizedCrop transform is deprecated, " +
+                      "please use transforms.RandomResizedCrop instead.")
+        super(RandomSizedCrop, self).__init__(*args, **kwargs)
+
+
+class FiveCrop(object):
+    """Crop the given PIL Image into four corners and the central crop
+    .. Note::
+         This transform returns a tuple of images and there may be a mismatch in the number of
+         inputs and targets your Dataset returns. See below for an example of how to deal with
+         this.
+    Args:
+         size (sequence or int): Desired output size of the crop. If size is an ``int``
+            instead of sequence like (h, w), a square crop of size (size, size) is made.
+    Example:
+         >>> transform = Compose([
+         >>>    FiveCrop(size), # this is a list of PIL Images
+         >>>    Lambda(lambda crops: torch.stack([ToTensor()(crop) for crop in crops])) # returns a 4D tensor
+         >>> ])
+         >>> #In your test loop you can do the following:
+         >>> input, target = batch # input is a 5d tensor, target is 2d
+         >>> bs, ncrops, c, h, w = input.size()
+         >>> result = model(input.view(-1, c, h, w)) # fuse batch size and ncrops
+         >>> result_avg = result.view(bs, ncrops, -1).mean(1) # avg over crops
+    """
+
+    def __init__(self, size):
+        self.size = size
+        if isinstance(size, numbers.Number):
+            self.size = (int(size), int(size))
+        else:
+            assert len(size) == 2, "Please provide only two dimensions (h, w) for size."
+            self.size = size
+
+    def __call__(self, img, target=None):
+        if target is not None:
+            return F.five_crop(img, self.size), F.five_crop(target, self.size)
+        return F.five_crop(img, self.size)
+
+    def __repr__(self):
+        return self.__class__.__name__ + '(size={0})'.format(self.size)
+
+
+class TenCrop(object):
+    """Crop the given PIL Image into four corners and the central crop plus the flipped version of
+    these (horizontal flipping is used by default)
+    .. Note::
+         This transform returns a tuple of images and there may be a mismatch in the number of
+         inputs and targets your Dataset returns. See below for an example of how to deal with
+         this.
+    Args:
+        size (sequence or int): Desired output size of the crop. If size is an
+            int instead of sequence like (h, w), a square crop (size, size) is
+            made.
+        vertical_flip(bool): Use vertical flipping instead of horizontal
+    Example:
+         >>> transform = Compose([
+         >>>    TenCrop(size), # this is a list of PIL Images
+         >>>    Lambda(lambda crops: torch.stack([ToTensor()(crop) for crop in crops])) # returns a 4D tensor
+         >>> ])
+         >>> #In your test loop you can do the following:
+         >>> input, target = batch # input is a 5d tensor, target is 2d
+         >>> bs, ncrops, c, h, w = input.size()
+         >>> result = model(input.view(-1, c, h, w)) # fuse batch size and ncrops
+         >>> result_avg = result.view(bs, ncrops, -1).mean(1) # avg over crops
+    """
+
+    def __init__(self, size, vertical_flip=False):
+        self.size = size
+        if isinstance(size, numbers.Number):
+            self.size = (int(size), int(size))
+        else:
+            assert len(size) == 2, "Please provide only two dimensions (h, w) for size."
+            self.size = size
+        self.vertical_flip = vertical_flip
+
+    def __call__(self, img, target = None):
+        if target is not None:
+            return F.ten_crop(img, self.size), F.ten_crop(target, self.size)
+        return F.ten_crop(img, self.size, self.vertical_flip)
+
+    def __repr__(self):
+        return self.__class__.__name__ + '(size={0}, vertical_flip={1})'.format(self.size, self.vertical_flip)
+
+
+class LinearTransformation(object):
+    """Transform a tensor image with a square transformation matrix computed
+    offline.
+    Given transformation_matrix, will flatten the torch.*Tensor, compute the dot
+    product with the transformation matrix and reshape the tensor to its
+    original shape.
+    Applications:
+    - whitening: zero-center the data, compute the data covariance matrix
+                 [D x D] with np.dot(X.T, X), perform SVD on this matrix and
+                 pass it as transformation_matrix.
+    Args:
+        transformation_matrix (Tensor): tensor [D x D], D = C x H x W
+    """
+
+    def __init__(self, transformation_matrix):
+        if transformation_matrix.size(0) != transformation_matrix.size(1):
+            raise ValueError("transformation_matrix should be square. Got " +
+                             "[{} x {}] rectangular matrix.".format(*transformation_matrix.size()))
+        self.transformation_matrix = transformation_matrix
+
+    def __call__(self, tensor, target_tensor=None):
+        """
+        Args:
+            tensor (Tensor): Tensor image of size (C, H, W) to be whitened.
+        Returns:
+            Tensor: Transformed image.
+        """
+        if target_tensor is not None:
+            raise NotImplementedError("LinearTransformation not implemented for tensor pairs.")
+        if tensor.size(0) * tensor.size(1) * tensor.size(2) != self.transformation_matrix.size(0):
+            raise ValueError("tensor and transformation matrix have incompatible shape." +
+                             "[{} x {} x {}] != ".format(*tensor.size()) +
+                             "{}".format(self.transformation_matrix.size(0)))
+        flat_tensor = tensor.view(1, -1)
+        transformed_tensor = torch.mm(flat_tensor, self.transformation_matrix)
+        tensor = transformed_tensor.view(tensor.size())
+        return tensor
+
+    def __repr__(self):
+        format_string = self.__class__.__name__ + '('
+        format_string += (str(self.transformation_matrix.numpy().tolist()) + ')')
+        return format_string
+
+
+class ColorJitter(object):
+    """Randomly change the brightness, contrast and saturation of an image.
+    Args:
+        brightness (float): How much to jitter brightness. brightness_factor
+            is chosen uniformly from [max(0, 1 - brightness), 1 + brightness].
+        contrast (float): How much to jitter contrast. contrast_factor
+            is chosen uniformly from [max(0, 1 - contrast), 1 + contrast].
+        saturation (float): How much to jitter saturation. saturation_factor
+            is chosen uniformly from [max(0, 1 - saturation), 1 + saturation].
+        hue(float): How much to jitter hue. hue_factor is chosen uniformly from
+            [-hue, hue]. Should be >=0 and <= 0.5.
+    """
+    def __init__(self, brightness=0, contrast=0, saturation=0, hue=0):
+        self.brightness = brightness
+        self.contrast = contrast
+        self.saturation = saturation
+        self.hue = hue
+
+    @staticmethod
+    def get_params(brightness, contrast, saturation, hue):
+        """Get a randomized transform to be applied on image.
+        Arguments are same as that of __init__.
+        Returns:
+            Transform which randomly adjusts brightness, contrast and
+            saturation in a random order.
+        """
+        transforms = []
+        if brightness > 0:
+            brightness_factor = random.uniform(max(0, 1 - brightness), 1 + brightness)
+            transforms.append(Lambda(lambda img: F.adjust_brightness(img, brightness_factor)))
+
+        if contrast > 0:
+            contrast_factor = random.uniform(max(0, 1 - contrast), 1 + contrast)
+            transforms.append(Lambda(lambda img: F.adjust_contrast(img, contrast_factor)))
+
+        if saturation > 0:
+            saturation_factor = random.uniform(max(0, 1 - saturation), 1 + saturation)
+            transforms.append(Lambda(lambda img: F.adjust_saturation(img, saturation_factor)))
+
+        if hue > 0:
+            hue_factor = random.uniform(-hue, hue)
+            transforms.append(Lambda(lambda img: F.adjust_hue(img, hue_factor)))
+
+        random.shuffle(transforms)
+        transform = Compose(transforms)
+
+        return transform
+
+    def __call__(self, img, target = None):
+        """
+        Args:
+            img (PIL Image): Input image.
+        Returns:
+            PIL Image: Color jittered image.
+        """
+        transform = self.get_params(self.brightness, self.contrast,
+                                    self.saturation, self.hue)
+
+        if target is not None:
+            return transform(img), target
+        return transform(img)
+
+    def __repr__(self):
+        format_string = self.__class__.__name__ + '('
+        format_string += 'brightness={0}'.format(self.brightness)
+        format_string += ', contrast={0}'.format(self.contrast)
+        format_string += ', saturation={0}'.format(self.saturation)
+        format_string += ', hue={0})'.format(self.hue)
+        return format_string
+
+
+class RandomRotation(object):
+    """Rotate the image by angle.
+    Args:
+        degrees (sequence or float or int): Range of degrees to select from.
+            If degrees is a number instead of sequence like (min, max), the range of degrees
+            will be (-degrees, +degrees).
+        resample ({PIL.Image.NEAREST, PIL.Image.BILINEAR, PIL.Image.BICUBIC}, optional):
+            An optional resampling filter.
+            See http://pillow.readthedocs.io/en/3.4.x/handbook/concepts.html#filters
+            If omitted, or if the image has mode "1" or "P", it is set to PIL.Image.NEAREST.
+        expand (bool, optional): Optional expansion flag.
+            If true, expands the output to make it large enough to hold the entire rotated image.
+            If false or omitted, make the output image the same size as the input image.
+            Note that the expand flag assumes rotation around the center and no translation.
+        center (2-tuple, optional): Optional center of rotation.
+            Origin is the upper left corner.
+            Default is the center of the image.
+    """
+
+    def __init__(self, degrees, resample=False, resample_tg=False, expand=False, center=None):
+        if isinstance(degrees, numbers.Number):
+            if degrees < 0:
+                raise ValueError("If degrees is a single number, it must be positive.")
+            self.degrees = (-degrees, degrees)
+        else:
+            if len(degrees) != 2:
+                raise ValueError("If degrees is a sequence, it must be of len 2.")
+            self.degrees = degrees
+
+        self.resample = resample
+        self.resample_tg = resample_tg
+        self.expand = expand
+        self.center = center
+
+    @staticmethod
+    def get_params(degrees):
+        """Get parameters for ``rotate`` for a random rotation.
+        Returns:
+            sequence: params to be passed to ``rotate`` for random rotation.
+        """
+        angle = random.uniform(degrees[0], degrees[1])
+
+        return angle
+
+    def __call__(self, img, target=None):
+        """
+            img (PIL Image): Image to be rotated.
+            target (PIL Image): (optional) Target to be rotated
+        Returns:
+            PIL Image: Rotated image(s).
+        """
+
+        angle = self.get_params(self.degrees)
+
+        if target is not None:
+            return F.rotate(img, angle, self.resample, self.expand, self.center), \
+                   F.rotate(target, angle, self.resample_tg, self.expand, self.center)
+                   # resample = False is by default nearest, appropriate for targets
+
+    def __repr__(self):
+        format_string = self.__class__.__name__ + '(degrees={0}'.format(self.degrees)
+        format_string += ', resample={0}'.format(self.resample)
+        format_string += ', expand={0}'.format(self.expand)
+        if self.center is not None:
+            format_string += ', center={0}'.format(self.center)
+        format_string += ')'
+        return format_string
+
+
+class RandomAffine(object):
+    """Random affine transformation of the image keeping center invariant
+    Args:
+        degrees (sequence or float or int): Range of degrees to select from.
+            If degrees is a number instead of sequence like (min, max), the range of degrees
+            will be (-degrees, +degrees). Set to 0 to desactivate rotations.
+        translate (tuple, optional): tuple of maximum absolute fraction for horizontal
+            and vertical translations. For example translate=(a, b), then horizontal shift
+            is randomly sampled in the range -img_width * a < dx < img_width * a and vertical shift is
+            randomly sampled in the range -img_height * b < dy < img_height * b. Will not translate by default.
+        scale (tuple, optional): scaling factor interval, e.g (a, b), then scale is
+            randomly sampled from the range a <= scale <= b. Will keep original scale by default.
+        shear (sequence or float or int, optional): Range of degrees to select from.
+            If degrees is a number instead of sequence like (min, max), the range of degrees
+            will be (-degrees, +degrees). Will not apply shear by default
+        resample ({PIL.Image.NEAREST, PIL.Image.BILINEAR, PIL.Image.BICUBIC}, optional):
+            An optional resampling filter.
+            See http://pillow.readthedocs.io/en/3.4.x/handbook/concepts.html#filters
+            If omitted, or if the image has mode "1" or "P", it is set to PIL.Image.NEAREST.
+        fillcolor (int): Optional fill color for the area outside the transform in the output image. (Pillow>=5.0.0)
+    """
+
+    def __init__(self, degrees, translate=None, scale=None, shear=None, resample=False, resample_tg=False, fillcolor=0):
+        if isinstance(degrees, numbers.Number):
+            if degrees < 0:
+                raise ValueError("If degrees is a single number, it must be positive.")
+            self.degrees = (-degrees, degrees)
+        else:
+            assert isinstance(degrees, (tuple, list)) and len(degrees) == 2, \
+                "degrees should be a list or tuple and it must be of length 2."
+            self.degrees = degrees
+
+        if translate is not None:
+            assert isinstance(translate, (tuple, list)) and len(translate) == 2, \
+                "translate should be a list or tuple and it must be of length 2."
+            for t in translate:
+                if not (0.0 <= t <= 1.0):
+                    raise ValueError("translation values should be between 0 and 1")
+        self.translate = translate
+
+        if scale is not None:
+            assert isinstance(scale, (tuple, list)) and len(scale) == 2, \
+                "scale should be a list or tuple and it must be of length 2."
+            for s in scale:
+                if s <= 0:
+                    raise ValueError("scale values should be positive")
+        self.scale = scale
+
+        if shear is not None:
+            if isinstance(shear, numbers.Number):
+                if shear < 0:
+                    raise ValueError("If shear is a single number, it must be positive.")
+                self.shear = (-shear, shear)
+            else:
+                assert isinstance(shear, (tuple, list)) and len(shear) == 2, \
+                    "shear should be a list or tuple and it must be of length 2."
+                self.shear = shear
+        else:
+            self.shear = shear
+
+        self.resample = resample
+        self.resample_tg = resample_tg
+        self.fillcolor = fillcolor
+
+    @staticmethod
+    def get_params(degrees, translate, scale_ranges, shears, img_size):
+        """Get parameters for affine transformation
+        Returns:
+            sequence: params to be passed to the affine transformation
+        """
+        angle = random.uniform(degrees[0], degrees[1])
+        if translate is not None:
+            max_dx = translate[0] * img_size[0]
+            max_dy = translate[1] * img_size[1]
+            translations = (np.round(random.uniform(-max_dx, max_dx)),
+                            np.round(random.uniform(-max_dy, max_dy)))
+        else:
+            translations = (0, 0)
+
+        if scale_ranges is not None:
+            scale = random.uniform(scale_ranges[0], scale_ranges[1])
+        else:
+            scale = 1.0
+
+        if shears is not None:
+            shear = random.uniform(shears[0], shears[1])
+        else:
+            shear = 0.0
+
+        return angle, translations, scale, shear
+
+    def __call__(self, img, target=None):
+        """
+            img (PIL Image): Image to be rotated.
+            target (PIL Image): (optional) Target to be rotated
+        Returns:
+            PIL Image: Rotated image(s).
+        """
+        ret = self.get_params(self.degrees, self.translate, self.scale, self.shear, img.size)
+        if target is not None:
+            return F.affine(img, *ret, resample=self.resample, fillcolor=self.fillcolor), \
+                   F.affine(target, *ret, resample=self.resample_tg, fillcolor=self.fillcolor)
+                   # resample = False is by default nearest, appropriate for targets
+        return F.affine(img, *ret, resample=self.resample, fillcolor=self.fillcolor)
+
+    def __repr__(self):
+        s = '{name}(degrees={degrees}'
+        if self.translate is not None:
+            s += ', translate={translate}'
+        if self.scale is not None:
+            s += ', scale={scale}'
+        if self.shear is not None:
+            s += ', shear={shear}'
+        if self.resample > 0:
+            s += ', resample={resample}'
+        if self.fillcolor != 0:
+            s += ', fillcolor={fillcolor}'
+        s += ')'
+        d = dict(self.__dict__)
+        d['resample'] = _pil_interpolation_to_str[d['resample']]
+        return s.format(name=self.__class__.__name__, **d)
+
+
+class Grayscale(object):
+    """Convert image to grayscale.
+    Args:
+        num_output_channels (int): (1 or 3) number of channels desired for output image
+    Returns:
+        PIL Image: Grayscale version of the input.
+        - If num_output_channels == 1 : returned image is single channel
+        - If num_output_channels == 3 : returned image is 3 channel with r == g == b
+    """
+
+    def __init__(self, num_output_channels=1):
+        self.num_output_channels = num_output_channels
+
+    def __call__(self, img, target = None):
+        """
+        Args:
+            img (PIL Image): Image to be converted to grayscale.
+        Returns:
+            PIL Image: Randomly grayscaled image.
+        """
+        if target is not None:
+            return F.to_grayscale(img, num_output_channels=self.num_output_channels), target
+        return F.to_grayscale(img, num_output_channels=self.num_output_channels)
+
+    def __repr__(self):
+        return self.__class__.__name__ + '(num_output_channels={0})'.format(self.num_output_channels)
+
+
+class RandomGrayscale(object):
+    """Randomly convert image to grayscale with a probability of p (default 0.1).
+    Args:
+        p (float): probability that image should be converted to grayscale.
+    Returns:
+        PIL Image: Grayscale version of the input image with probability p and unchanged
+        with probability (1-p).
+        - If input image is 1 channel: grayscale version is 1 channel
+        - If input image is 3 channel: grayscale version is 3 channel with r == g == b
+    """
+
+    def __init__(self, p=0.1):
+        self.p = p
+
+    def __call__(self, img, target = None):
+        """
+        Args:
+            img (PIL Image): Image to be converted to grayscale.
+        Returns:
+            PIL Image: Randomly grayscaled image.
+        """
+        num_output_channels = 1 if img.mode == 'L' else 3
+        if random.random() < self.p:
+            if target is not None:
+                return F.to_grayscale(img, num_output_channels=num_output_channels), target
+        if target is not None:
+            return img, target
+        return img
+
+    def __repr__(self):
+        return self.__class__.__name__ + '(p={0})'.format(self.p)

dependecies/segroot/predict_imgs.py

@@ -0,0 +1,121 @@
+import argparse
+from pathlib import Path
+from PIL import Image
+import torch
+import torchvision
+from skimage.morphology import erosion
+import matplotlib.pyplot as plt
+import time
+
+from segroot.utils import init_weights
+from segroot.dataloader import pad_pair_256, normalize
+from segroot.model import SegRoot
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "--image", default="test.jpg", type=str, help="filename of one test image"
+)
+parser.add_argument(
+    "--thres", default=0.9, type=float, help="threshold of the final binarization"
+)
+parser.add_argument(
+    "--all", action="store_true", help="make prediction on all images in the folder"
+)
+parser.add_argument(
+    "--data_dir",
+    default="../data/prediction",
+    type=Path,
+    help="define the data directory",
+)
+parser.add_argument(
+    "--weights",
+    default="../weights/best_segnet-(8,5)-0.6441.pt",
+    type=Path,
+    help="path of pretrained weights",
+)
+parser.add_argument("--width", default=8, type=int, help="width of SegRoot")
+parser.add_argument("--depth", default=5, type=int, help="depth of SegRoot")
+
+
+def pad_256(img_path):
+    image = Image.open(img_path)
+    W, H = image.size
+    img, _ = pad_pair_256(image, image)
+    NW, NH = img.size
+    img = torchvision.transforms.ToTensor()(img)
+    img = normalize(img)
+    return img, (H, W, NH, NW)
+
+
+def predict(model, test_img, device):
+    for p in model.parameters():
+        p.requires_grad = False
+
+    model.eval()
+    # test_img.shape = (3, 2304, 2560)
+    test_img = test_img.unsqueeze(0)
+    output = model(test_img)
+    # output.shape = (1, 1, 2304, 2560)
+    output = torch.squeeze(output)
+    torch.cuda.empty_cache()
+    return output
+
+
+def predict_gen(model, img_path, thres, device, info):
+    img, dims = pad_256(img_path)
+    H, W, NH, NW = dims
+    img = img.to(device)
+    prediction = predict(model, img, device)
+    prediction[prediction >= thres] = 1.0
+    prediction[prediction < thres] = 0.0
+    if device.type == "cpu":
+        prediction = prediction.detach().numpy()
+    else:
+        prediction = prediction.cpu().detach().numpy()
+    prediction = erosion(prediction)
+    # reverse padding
+    prediction = prediction[
+        (NH - H) // 2 : (NH - H) // 2 + H, (NW - W) // 2 : (NW - W) // 2 + W
+    ]
+    save_path = img_path.parent / (
+        img_path.parts[-1].split(".jpg")[0] + "-pre-mask-segnet-({},5).jpg".format(info)
+    )
+    plt.imsave(save_path.as_posix(), prediction, cmap="gray")
+    print("{} generated!".format(save_path.parts[-1]))
+
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+    # define model
+    print("using segnet, width : {}, depth : {}".format(args.width, args.depth))
+    model = SegRoot(args.width, args.depth).to(device)
+    weights_path = args.weights
+
+    if device.type == "cpu":
+        print("load weights to cpu")
+        print(weights_path.as_posix())
+        model.load_state_dict(torch.load(weights_path.as_posix(), map_location="cpu"))
+    else:
+        print("load weights to gpu")
+        print(weights_path.as_posix())
+        model.load_state_dict(torch.load(weights_path.as_posix()))
+
+    # define the prediction's saving directory
+    pre_dir = Path("../data/prediction")
+    pre_dir.mkdir(parents=True, exist_ok=True)
+    if not args.all:
+        # load and pad image
+        img_path = pre_dir / args.image
+        start_time = time.time()
+        predict_gen(model, img_path, args.thres, device, 8)
+        end_time = time.time()
+        print("{:.4f}s for one image".format(end_time - start_time))
+    else:
+        img_paths = args.data_dir.glob("*.jpg")
+        for img_path in img_paths:
+            start_time = time.time()
+            predict_gen(model, img_path, args.thres, device, 8)
+            end_time = time.time()
+            print("{:.4f}s for one image".format(end_time - start_time))
+

dependecies/segroot/run_all_experiments.sh

@@ -0,0 +1,6 @@
+# !/bin/sh
+python -u train_segroot.py --width 2 > "log_SegRoot(2,5).txt"
+python -u train_segroot.py --width 16 --depth 4 --lr 1e-3 > "log_SegRoot(16,4).txt"
+python -u train_segroot.py --width 32 --depth 5 --lr 1e-4 --bs 32 > "log_SegRoot(32,5).txt"
+python -u train_segroot.py --width 64 --depth 4 --lr 1e-4 --bs 16 > "log_SegRoot(64,4).txt"
+python -u train_segroot.py --width 64 --depth 5 --lr 2e-5 --bs 8 --epochs 100 --verbose 2 > "log_SegRoot(64,5).txt"

dependecies/segroot/utils.py

@@ -0,0 +1,109 @@
+import pickle
+import torch
+from torchvision import models
+import random
+import logging
+import numpy as np
+import json
+
+def set_random_seed(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+    torch.backends.cudnn.deterministic = True
+
+def set_logger(log_path):
+    logger = logging.getLogger()
+    logger.setLevel(logging.INFO)
+
+    if not logger.handlers:
+        # Logging to a file
+        file_handler = logging.FileHandler(log_path)
+        file_handler.setFormatter(logging.Formatter('%(asctime)s:%(levelname)s: %(message)s'))
+        logger.addHandler(file_handler)
+
+        # Logging to console
+        stream_handler = logging.StreamHandler()
+        stream_handler.setFormatter(logging.Formatter('%(message)s'))
+        logger.addHandler(stream_handler)
+
+def to_np(x):
+    return x.data.cpu().numpy()
+
+def get_ids(length_dataset):
+    ids = list(range(length_dataset))
+
+    random.shuffle(ids)
+    train_split = round(0.6 * length_dataset)
+    t_v_spplit = (length_dataset - train_split) // 2
+    train_ids = ids[:train_split]
+    valid_ids = ids[train_split:train_split+t_v_spplit]
+    test_ids = ids[train_split+t_v_spplit:]
+    return train_ids, valid_ids, test_ids
+
+def dice_score(y, y_pred, smooth=1.0, thres=0.9):
+    n = y.shape[0]
+    y = y.view(n, -1)
+    y_pred = y_pred.view(n, -1)
+    # y_pred_[y_pred>=thres] = 1.0
+    # y_pred_[y_pred<thres] = 0.0 
+    num = 2 * torch.sum(y * y_pred, dim=1, keepdim=True) + smooth
+    den = torch.sum(y, dim=1, keepdim=True) + \
+        torch.sum(y_pred, dim=1, keepdim=True) + smooth
+    score = num / den
+    return score
+
+def init_weights(m):
+    if isinstance(m, torch.nn.Conv2d):
+        torch.nn.init.kaiming_uniform_(m.weight, nonlinearity='relu')
+        # torch.nn.init.constant_(m.bias, 0)
+    elif isinstance(m, torch.nn.BatchNorm2d):
+        torch.nn.init.constant_(m.weight, 1)
+
+def load_vgg16(segnet):
+    vgg16 = models.vgg16_bn(pretrained=True)
+    with open('paired_weight_vgg16.plk', 'rb') as handle:
+        paired = pickle.load(handle)
+    segnet_p = dict(segnet.state_dict())
+    vgg16_p = vgg16.state_dict()
+
+    for k, v in paired.items():
+        for n, p in vgg16_p.items():
+            if n == v:
+                segnet_p[k].data.copy_(p.data)
+    segnet.load_state_dict(segnet_p)
+    return segnet
+
+def train_one_epoch(model, train_iter, optimizer, device):
+    model.train()
+    for p in model.parameters():
+        p.requires_grad = True
+    for x, y in train_iter:
+        x, y = x.to(device), y.to(device)
+        bs = x.shape[0]
+        optimizer.zero_grad()
+        y_pred = model(x)
+        loss = 1 - dice_score(y, y_pred)
+        loss = torch.sum(loss) / bs
+        loss.backward()
+        optimizer.step()
+
+def evaluate(model, dataset, device, thres=0.9):
+    model.eval()
+    torch.cuda.empty_cache()    
+    num, den = 0, 0
+    # shutdown the autograd
+    with torch.no_grad():
+        for i in range(len(dataset)):
+            x, y = dataset[i]
+            x, y = x.unsqueeze(0).to(device), y.unsqueeze(0).to(device)
+            y_pred = model(x)
+            y = y.cpu().detach().numpy()
+            y_pred = y_pred.cpu().detach().numpy()
+            y_pred[y_pred>=thres] = 1.0
+            y_pred[y_pred<thres] = 0.0
+            num += 2 * (y_pred * y).sum()
+            den += y_pred.sum() + y.sum()
+    torch.cuda.empty_cache() 
+    return num / den

flagged/log.csv

@@ -0,0 +1,2 @@
+input_img,Model,output,flag,username,timestamp
+flagged\input_img\a7a20e8c8e03de5e007f\example_1.jpg,segroot_finetuned,,,,2024-11-20 11:20:45.490192

main.py

@@ -0,0 +1,188 @@
+import gradio as gr
+
+from processsors import RootSegmentor
+from processsors import *
+
+from gradio_imageslider import ImageSlider
+
+import cv2 as cv
+
+PRELOAD_MODELS = False
+
+if PRELOAD_MODELS:
+    root_segmentor = RootSegmentor()
+    
+    
+def process(input_img, model_type):
+    
+    print(model_type)
+    
+    if PRELOAD_MODELS:
+        global root_segmentor
+    else:
+        root_segmentor = RootSegmentor(model_type)
+        
+    result = root_segmentor.predict(input_img)
+    
+    return result
+
+def just_show(files, should_process, model_type):
+    
+    imgs = []
+    
+    img = merge_images(files)
+    
+    
+    
+    imgs.append(img)
+    
+    if should_process:    
+        root_segmentor = RootSegmentor(model_type)
+    
+    results = []
+
+    for file in files:
+        print(type(file))
+        print(file)
+        img = cv.imread(file)
+        img = cv.cvtColor(img, cv.COLOR_BGR2RGB)
+        #imgs.append(img)
+        
+        if should_process: 
+        
+            result = root_segmentor.predict(img)
+            results.append(result)
+            #imgs.append(results)
+        
+    if should_process: 
+        img_res = merge_images(results)
+        imgs.append(img_res)
+
+    return imgs
+
+def slider_test(img1, img2):
+ 
+    return [img1,img2]
+
+def download_result():
+    
+    #print(filepath)
+    return
+    
+
+def gui():
+
+  with gr.Blocks(title="Root analysis", theme=gr.themes.Soft()) as demo:
+
+    big_block = gr.HTML("""
+
+    <style>
+        body {
+            font-family: Arial, sans-serif;
+            background-color: white
+            margin: 0;
+        }
+
+        header {
+            display: flex;
+            justify-content: space-between;
+            align-items: center;
+            padding: 5px;
+            color: #fff;
+        }
+
+        hr {
+            border: 1px solid #ddd;
+            margin: 5px;
+        }
+
+    </style>
+
+    <header>
+        <div style="display: flex; align-items: center;">
+            <div style="text-align: left;">
+            <h1>Root Analysis</h1>
+            <p>Root segmentation using underground root scanner images.</p>
+            <h3>Tropical Forages Program</h3>
+            <p><b>Authors: </b>Andres Felipe Ruiz-Hurtado, Juan Andrés Cardoso Arango</p>
+            <p></p>            
+        </div>
+        </div>
+        <div style="background-color: white; padding: 5px; border-radius: 15px; box-shadow: 0px 4px 8px rgba(0, 0, 0, 0.1);">
+                        <img src='file/logo.png' alt="Logo" width="200" height="100">
+                    </div>
+    </header>   
+    
+    """)
+    
+    #<iframe style="height:600px;width: 100%;" src="/file=slides.html" title="description"></iframe>
+
+    
+    #<iframe style="height:600px;width: 100%;" src="https://revealjs.com/demo/?view" title="description"></iframe>
+
+    with gr.Tab("Single Image"):
+    
+        model_selector = gr.Dropdown(
+                ["segroot_finetuned", "segroot", "segroot_finetuned_dec", "seg_model"], label="Model"
+                , info="AI model"
+                ,value="segroot_finetuned"
+            )
+
+        input_img=gr.Image(render=False)
+        output_img=gr.Image(render=False)
+
+        gr.Interface(
+            fn=process,
+            inputs=[input_img,model_selector],
+            outputs=output_img,
+            examples=[["example_1.jpg"],["example_2.jpg"],["example_3.jpg"]]
+        )
+
+        #examples = gr.Examples([["Chicago"], ["Little Rock"], ["San Francisco"]], textbox)
+
+        with gr.Row():
+            img_comp = ImageSlider(label="Root Segmentation")
+        with gr.Row():
+            compare_button = gr.Button("Compare")
+            compare_button.click(fn=slider_test, inputs=[input_img,output_img], outputs=img_comp, api_name="slider_test")
+        
+    with gr.Tab("Multiple Images"):
+
+    #img_comp = ImageSlider(label="Blur image", type="pil")
+    
+        gallery = gr.Gallery(show_fullscreen_button=True, render=False) 
+        
+        gr.Interface(
+            fn=just_show
+            ,inputs=[gr.File(file_count="multiple"),gr.Checkbox(label="Process", info="Check if you want to process"),model_selector]
+            ,outputs= gallery
+            , examples=[[["example_1.jpg", "example_2.jpg", "example_3.jpg"]]]
+        )
+
+    with gr.Tab("Compare"):
+
+        img_comp = ImageSlider(label="Root Segmentation")
+        img_comp.upload(inputs=img_comp, outputs=img_comp)
+
+    
+    #d = gr.DownloadButton("Download the file")
+    #d.click(download_result, gallery, None)
+    
+    # with gr.Row():    
+    #     img1=gr.Image()
+    #     img2=gr.Image()
+    # with gr.Row():
+    #     img_comp = ImageSlider(label="Blur image", type="pil")
+    # with gr.Row():
+    #     compare_button = gr.Button("Compare")
+    #     compare_button.click(fn=slider_test, inputs=[img1,img2], outputs=img_comp, api_name="slider_test")
+    
+    # with gr.Group():
+    #     img_comp = ImageSlider(label="Blur image", type="pil")
+    #     #img1.upload(slider_test, inputs=[img1,img2], outputs=img_comp)
+    #     gr.Interface(slider_test, inputs=[img1,img2], outputs=img_comp)
+    
+    demo.launch(allowed_paths=["logo.png"], share=False)
+
+if __name__ == "__main__":
+    gui()

models/best_segnet-(8,5)-0.6441.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:dffa166609b5ab3241d1b175bffbf454377beaa4a7fb46bd74e38605e2f71d03
+size 1611034

models/roots_model.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:48254d394d1b11fd9bcfd42bcc754bb1fba5a2052848f5ad70b259972bce4681
+size 58655218

models/segroot-(8,5)_finetuned.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cbb992086ea1900ef24e110d7b454126d6214ac8f14687348ec021cf860f4eca
+size 1640578

processsors.py

@@ -0,0 +1,210 @@
+import torch
+import torchvision
+
+
+from PIL import Image
+import numpy as np
+
+from skimage.morphology import erosion
+
+from dependecies.segroot.model import SegRoot
+from dependecies.segroot.dataloader import pad_pair_256, normalize
+from torchvision.transforms import v2 as transforms
+
+
+import onnxruntime as ort
+import cv2 as cv
+
+import os
+
+MODELS_PATH = r"./models"
+
+def pad_256(img_path):
+    image = Image.open(img_path)
+    W, H = image.size
+    img, _ = pad_pair_256(image, image)
+    NW, NH = img.size
+    img = torchvision.transforms.ToTensor()(img)
+    img = normalize(img)
+    return img, (H, W, NH, NW)
+
+def pad_256_np(np_img):
+    #image = Image.open(img_path)
+    image = Image.fromarray(np_img)
+    W, H = image.size
+    img, _ = pad_pair_256(image, image)
+    NW, NH = img.size
+    img = torchvision.transforms.ToTensor()(img)
+    img = normalize(img)
+    return img, (H, W, NH, NW)
+
+def merge_images(files, path=""):
+    
+    is_array = False
+    if type(files[0]) == np.ndarray:
+        is_array = True
+        
+
+    final_img = []
+    resize_factor = 0.4
+    offset0 = 930
+    offset1 = 305 
+    for index, file in enumerate(files):
+
+        if is_array:
+            img = file
+        else:
+            img = cv.imread(file)
+            img = cv.cvtColor(img, cv.COLOR_BGR2RGB)
+        #img = cv.resize(img, (0,0), fx=resize_factor, fy=resize_factor)
+        img = cv.rotate(img, cv.ROTATE_90_CLOCKWISE)
+        
+        if index == 0:
+            img = img[0:img.shape[0]-offset0,0:img.shape[1]]
+            final_img = img
+        elif index == len(file)-1:
+            final_img = cv.vconcat([final_img, img])
+        else:
+            #final_img = np.concatenate((final_img, img), axis=1)
+            img = img[0:img.shape[0]-offset1,0:img.shape[1]]
+            final_img = cv.vconcat([final_img, img])
+        
+    final_img = cv.resize(final_img, (0,0), fx=resize_factor, fy=resize_factor)  
+     
+    #cv.imwrite(path, final_img)
+    print(final_img.shape)
+    
+    return final_img
+
+class RootSegmentor():
+    
+    def __init__(self, model_type):
+        
+        
+        self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        
+        self.model_type = model_type
+        
+        if model_type != "seg_model":        
+            self.initialize()
+
+        return
+    
+    def initialize(self):
+        
+        width = 8
+        depth = 5
+        
+        if self.model_type == "segroot":
+            #weights_path = os.path.join(r"D:\local_mydev\roots_finetuning\SegRoot0\weights\best_segnet-(8,5)-0.6441.pt"
+            #weights_path = r"D:\local_mydev\SegRoot\weights\best_segnet-(8,5)-0.6441.pt"
+            #weights_path = r"\\CATALOGUE.CGIARAD.ORG\AcceleratedBreedingInitiative\4.Scripts\AndresRuiz\local_mydata_backup\model\roots\best_segnet-(8,5)-0.6441.pt"
+            #weights_path = os.path.join(MODELS_PATH, r"AcceleratedBreedingInitiative\4.Scripts\AndresRuiz\local_mydata_backup\model\roots\best_segnet-(8,5)-0.6441.pt")
+            weights_path = os.path.join(MODELS_PATH, r"best_segnet-(8,5)-0.6441.pt")
+        elif self.model_type == "segroot_finetuned":
+            #weights_path = r"\\CATALOGUE.CGIARAD.ORG\AcceleratedBreedingInitiative\4.Scripts\AndresRuiz\local_mydata_backup\model\roots\segroot-(8,5)_finetuned.pt"
+            #weights_path = os.path.join(MODELS_PATH, r"AcceleratedBreedingInitiative\4.Scripts\AndresRuiz\local_mydata_backup\model\roots\segroot-(8,5)_finetuned.pt")
+            weights_path = os.path.join(MODELS_PATH, r"segroot-(8,5)_finetuned.pt")
+        elif self.model_type == "segroot_finetuned_dec":
+            #weights_path = r"\\CATALOGUE.CGIARAD.ORG\AcceleratedBreedingInitiative\4.Scripts\AndresRuiz\local_mydata_backup\model\roots\segroot-(8,5)_finetuned_dec_full.pt"
+            #weights_path = r"\\CATALOGUE.CGIARAD.ORG\AcceleratedBreedingInitiative\4.Scripts\AndresRuiz\local_mydata_backup\model\roots\segroot-(8,5)_finetuned_clas.pt"
+            #weights_path = os.path.join(MODELS_PATH, r"AcceleratedBreedingInitiative\4.Scripts\AndresRuiz\local_mydata_backup\model\roots\segroot-(8,5)_finetuned_clas.pt")
+            weights_path = os.path.join(MODELS_PATH, r"segroot-(8,5)_finetuned.pt")
+            
+        self.model = SegRoot(width, depth).to(self.device)
+
+        if self.device.type == "cpu":
+            print("load weights to cpu")
+            #print(weights_path.as_posix())
+            self.model.load_state_dict(torch.load(weights_path, map_location="cpu"))
+        else:
+            print("load weights to gpu")
+            #print(weights_path.as_posix())
+            self.model.load_state_dict(torch.load(weights_path))
+            
+        for p in self.model.parameters():
+            p.requires_grad = False
+            
+        self.model.eval()
+                
+        return    
+    
+    def predict(self, img_path):
+        
+        if self.model_type == "seg_model":
+            
+            print(str(type(img_path)))
+            
+            if type(img_path) == np.ndarray:
+                img = img_path  
+            else:
+                img = cv.imread(img_path)
+                img = cv.cvtColor(img, cv.COLOR_BGR2RGB)
+            
+            weights_path = r"\\CATALOGUE.CGIARAD.ORG\AcceleratedBreedingInitiative\4.Scripts\AndresRuiz\local_mydata_backup\model\roots\roots_model.onnx"
+            weights_path = os.path.join(MODELS_PATH,"roots_model.onnx") 
+            ort_sess = ort.InferenceSession(weights_path
+                                ,providers=ort.get_available_providers()
+                                )
+            
+            dim = img.shape
+            
+            transforms_list = []
+            transforms_list.append(transforms.ToTensor())
+            transforms_list.append(transforms.Resize((800,800)))
+            #transforms_list.append(transforms.CenterCrop(800))
+            #transforms_list.append(transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]))
+
+            apply_t =  transforms.Compose(transforms_list) 
+
+            img = apply_t(img)
+
+            outputs = ort_sess.run(None, {'input': [img.numpy()]})
+            
+            print(outputs)
+
+            #np_res = outputs[0][0]
+
+            output_image = outputs[0][:,:,1]
+            final = cv.resize(output_image, (dim[0], dim[1]))
+                
+            return final
+            
+        else:
+
+            thres = 0.9
+            
+            print(str(type(img_path)))
+            
+            if type(img_path) == np.ndarray:
+                img, dims = pad_256_np(img_path)    
+            else:
+                img, dims = pad_256(img_path)
+        
+            H, W, NH, NW = dims
+            
+            img = img.to(self.device)
+            
+            img = img.unsqueeze(0)
+            output = self.model(img)
+
+            output = torch.squeeze(output)
+            torch.cuda.empty_cache()
+
+            prediction = output
+            
+            prediction[prediction >= thres] = 1.0
+            prediction[prediction < thres] = 0.0
+            
+            if self.device.type == "cpu":
+                prediction = prediction.detach().numpy()
+            else:
+                prediction = prediction.cpu().detach().numpy()
+                
+            prediction = erosion(prediction)
+            # reverse padding
+            prediction = prediction[
+                (NH - H) // 2 : (NH - H) // 2 + H, (NW - W) // 2 : (NW - W) // 2 + W
+            ]
+
+            return prediction

requirements.txt

@@ -0,0 +1,11 @@
+matplotlib
+numpy
+opencv-python
+pillow
+scikit-image
+scikit-learn
+torch
+torchvision
+gradio
+onnxruntime
+rasterio