Upload 7 files

TEED/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022 Xavier Soria Poma
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

TEED/README.md

@@ -0,0 +1,68 @@
+[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/tiny-and-efficient-model-for-the-edge/edge-detection-on-uded)](https://paperswithcode.com/sota/edge-detection-on-uded?p=tiny-and-efficient-model-for-the-edge)
+
+# Tiny and Efficient Model for the Edge Detection Generalization (Paper)
+
+## Overview
+
+<div style="text-align:center"><img src='imgs/teedBanner.png' width=800>
+</div>
+
+
+
+Tiny and Efficient Edge Detector (TEED) is a light convolutional neural
+network with only $58K$ parameters, less than $0.2$% of the 
+state-of-the-art models. Training on the [BIPED](https://www.kaggle.com/datasets/xavysp/biped)
+dataset takes *less than 30 minutes*, with each epoch requiring 
+*less than 5 minutes*. Our proposed model is easy to train
+and it quickly converges within very first few epochs, while the 
+predicted edge-maps are crisp and of high quality, see image above.
+[This paper has been accepted by ICCV 2023-Workshop RCV](https://arxiv.org/abs/2308.06468).
+
+... In construction
+
+    git clone https://github.com/xavysp/TEED.git
+    cd TEED
+
+Then,
+
+## Testing with TEED
+
+Copy and paste your images into data/ folder, and:
+
+    python main.py --choose_test_data=-1
+
+## Training with TEED
+
+Set the following lines in main.py:
+
+    25: is_testing =False
+    # training with BIPED
+    223: TRAIN_DATA = DATASET_NAMES[0] 
+
+then run
+    
+    python main.py
+
+Check the configurations of the datasets in dataset.py
+
+
+## UDED dataset
+
+Here the [link](https://github.com/xavysp/UDED) to access the UDED dataset for edge detection
+
+## Citation
+
+If you like TEED, why not starring the project on GitHub!
+
+[![GitHub stars](https://img.shields.io/github/stars/xavysp/TEED.svg?style=social&label=Star&maxAge=3600)](https://GitHub.com/xavysp/TEED/stargazers/)
+
+Please cite our Dataset if you find helpful in your academic/scientific publication,
+```
+@InProceedings{Soria_2023teed,
+    author    = {Soria, Xavier and Li, Yachuan and Rouhani, Mohammad and Sappa, Angel D.},
+    title     = {Tiny and Efficient Model for the Edge Detection Generalization},
+    booktitle = {Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV) Workshops},
+    month     = {October},
+    year      = {2023},
+    pages     = {1364-1373}
+}

TEED/dataset.py

@@ -0,0 +1,581 @@
+import os
+import random
+
+import cv2
+import numpy as np
+import torch
+from torch.utils.data import Dataset
+import json
+
+DATASET_NAMES = [
+    'BIPED',
+    'BIPED-B2',
+    'BIPED-B3',
+    'BIPED-B5',
+    'BIPED-B6',
+    'BSDS', # 5
+    'BRIND', # 6
+    'ICEDA', #7
+    'BSDS300',
+    'CID', #9
+    'DCD',
+    'MDBD', #11
+    'PASCAL',
+    'NYUD', #13
+    'BIPBRI',
+    'UDED', # 15 just for testing
+    'DMRIR',
+    'CLASSIC'
+]  # 8
+# [108, 109.451,112.230,137.86]
+BIPED_mean = [103.939,116.779,123.68,137.86]
+
+def dataset_info(dataset_name, is_linux=True):
+    if is_linux:
+
+        config = {
+            'UDED': {
+                'img_height': 512,  # 321
+                'img_width': 512,  # 481
+                'train_list': None,
+                'test_list': 'test_pair.lst',
+                'data_dir': '/root/workspace/datasets/UDED',  # mean_rgb
+                'yita': 0.5,
+                'mean': [104.007, 116.669, 122.679, 137.86]# [104.007, 116.669, 122.679, 137.86]
+            }, #[98.939,111.779,117.68,137.86]
+            'BSDS': {
+                'img_height': 512, #321
+                'img_width': 512, #481
+                'train_list': 'train_pair.lst',
+                'test_list': 'test_pair.lst',
+                'data_dir': '/root/workspace/datasets/BSDS',  # mean_rgb
+                'yita': 0.5,
+                'mean': [104.007, 116.669, 122.679, 137.86]
+            },
+            'BRIND': {
+                'img_height': 512,  # 321
+                'img_width': 512,  # 481
+                'train_list': 'train_pair_all.lst',
+                # all train_pair_all.lst
+                # less train_pair.lst
+                'test_list': 'test_pair.lst',
+                'data_dir': '/root/workspace/datasets/BRIND',  # mean_rgb
+                'yita': 0.5,
+                'mean': [104.007, 116.669, 122.679, 137.86]
+            },
+            'ICEDA': {
+            'img_height': 1024,  # 321
+            'img_width': 1408,  # 481
+            'train_list': None,
+            'test_list': 'test_pair.lst',
+            'data_dir': '/root/workspace/datasets/ICEDA',  # mean_rgb
+            'yita': 0.5,
+            'mean': [104.007, 116.669, 122.679, 137.86]
+        },
+            'BSDS300': {
+                'img_height': 512, #321
+                'img_width': 512, #481
+                'test_list': 'test_pair.lst',
+                'train_list': None,
+                'data_dir': '/root/workspace/datasets/BSDS300',  # NIR
+                'yita': 0.5,
+                'mean': [104.007, 116.669, 122.679, 137.86]
+            },
+            'PASCAL': {
+                'img_height': 416, # 375
+                'img_width': 512, #500
+                'test_list': 'test_pair.lst',
+                'train_list': None,
+                'data_dir': '/root/datasets/PASCAL',  # mean_rgb
+                'yita': 0.3,
+                'mean': [104.007, 116.669, 122.679, 137.86]
+            },
+            'CID': {
+                'img_height': 512,
+                'img_width': 512,
+                'test_list': 'test_pair.lst',
+                'train_list': None,
+                'data_dir': '/root/datasets/CID',  # mean_rgb
+                'yita': 0.3,
+                'mean': [104.007, 116.669, 122.679, 137.86]
+            },
+            'NYUD': {
+                'img_height': 448,#425
+                'img_width': 560,#560
+                'test_list': 'test_pair.lst',
+                'train_list': None,
+                'data_dir': '/root/datasets/NYUD',  # mean_rgb
+                'yita': 0.5,
+                'mean': [104.007, 116.669, 122.679, 137.86]
+            },
+            'MDBD': {
+                'img_height': 720,
+                'img_width': 1280,
+                'test_list': 'test_pair.lst',
+                'train_list': 'train_pair.lst',
+                'data_dir': '/root/workspace/datasets/MDBD',  # mean_rgb
+                'yita': 0.3,
+                'mean': [104.007, 116.669, 122.679, 137.86]
+            },
+            'BIPED': {
+                'img_height': 720, #720 # 1088
+                'img_width': 1280, # 1280 5 1920
+                'test_list': 'test_pair.lst',
+                'train_list': 'train_pair0.lst', # Base augmentation
+                # 'train_list': 'train_pairB3.lst', # another augmentation
+                # 'train_list': 'train_pairB5.lst', # Last augmentation
+                'data_dir': '/root/workspace/datasets/BIPED',  # mean_rgb
+                'yita': 0.5,
+                'mean':BIPED_mean
+            #
+            },
+            'CLASSIC': {
+                'img_height': 512,#
+                'img_width': 512,# 512
+                'test_list': None,
+                'train_list': None,
+                'data_dir': 'data',  # mean_rgb
+                'yita': 0.5,
+                'mean': [104.007, 116.669, 122.679, 137.86]
+            },
+            'BIPED-B2': {'img_height': 720,  # 720
+                         'img_width': 1280,  # 1280
+                         'test_list': 'test_pair.lst',
+                         'train_list': 'train_rgb.lst',
+                         'data_dir': 'C:/Users/xavysp/dataset/BIPED',  # WIN: '../.../dataset/BIPED/edges'
+                         'yita': 0.5,
+                         'mean':BIPED_mean},
+            'BIPED-B3': {'img_height': 720,  # 720
+                         'img_width': 1280,  # 1280
+                         'test_list': 'test_pair.lst',
+                         'train_list': 'train_rgb.lst',
+                         'data_dir': 'C:/Users/xavysp/dataset/BIPED',  # WIN: '../.../dataset/BIPED/edges'
+                         'yita': 0.5,
+                         'mean':BIPED_mean},
+            'BIPED-B5': {'img_height': 720,  # 720
+                         'img_width': 1280,  # 1280
+                         'test_list': 'test_pair.lst',
+                         'train_list': 'train_rgb.lst',
+                         'data_dir': 'C:/Users/xavysp/dataset/BIPED',  # WIN: '../.../dataset/BIPED/edges'
+                         'yita': 0.5,
+                        'mean':BIPED_mean},
+            'BIPED-B6': {'img_height': 720,  # 720
+                         'img_width': 1280,  # 1280
+                         'test_list': 'test_pair.lst',
+                         'train_list': 'train_rgb.lst',
+                         'data_dir': 'C:/Users/xavysp/dataset/BIPED',  # WIN: '../.../dataset/BIPED/edges'
+                         'yita': 0.5,
+                         'mean':BIPED_mean},
+            'DCD': {
+                'img_height': 352, #240
+                'img_width': 480,# 360
+                'test_list': 'test_pair.lst',
+                'train_list': None,
+                'data_dir': '/opt/dataset/DCD',  # mean_rgb
+                'yita': 0.2,
+                'mean': [104.007, 116.669, 122.679, 137.86]
+            }
+        }
+    else:
+        config = {
+            'UDED': {
+                'img_height': 512,  # 321
+                'img_width': 512,  # 481
+                'train_list': None,
+                'test_list': 'test_pair.lst',
+                'data_dir': 'C:/dataset/UDED',  # mean_rgb
+                'yita': 0.5,
+                'mean':[104.007, 116.669, 122.679, 137.86] # [183.939,196.779,203.68,137.86] # [104.007, 116.669, 122.679, 137.86]
+            },
+            'BSDS': {'img_height': 480,  # 321
+                     'img_width': 480,  # 481
+                     'test_list': 'test_pair.lst',
+                     'data_dir': 'C:/dataset/BSDS',  # mean_rgb
+                     'yita': 0.5,
+                    'mean':[103.939, 116.669, 122.679, 137.86] },
+            # [103.939, 116.669, 122.679, 137.86]
+            #[159.510, 159.451,162.230,137.86]
+            'BRIND': {
+                'img_height': 512,  # 321
+                'img_width': 512,  # 481
+                'train_list': 'train_pair_all.lst',
+                # all train_pair_all.lst
+                # less train_pair.lst
+                'test_list': 'test_pair.lst',
+                'data_dir': 'C:/dataset/BRIND',  # mean_rgb
+                'yita': 0.5,
+                'mean': [104.007, 116.669, 122.679, 137.86]
+            },
+            'ICEDA': {
+            'img_height': 1024,  # 321
+            'img_width': 1408,  # 481
+            'train_list': None,
+            'test_list': 'test_pair.lst',
+            'data_dir': 'C:/dataset/ICEDA',  # mean_rgb
+            'yita': 0.5,
+            'mean': [104.007, 116.669, 122.679, 137.86]
+        },
+            'BSDS300': {'img_height': 512,  # 321
+                        'img_width': 512,  # 481
+                        'test_list': 'test_pair.lst',
+                        'data_dir': 'C:/Users/xavysp/dataset/BSDS300',  # NIR
+                        'yita': 0.5,
+                    'mean': [104.007, 116.669, 122.679, 137.86]},
+            'PASCAL': {'img_height': 375,
+                       'img_width': 500,
+                       'test_list': 'test_pair.lst',
+                       'data_dir': 'C:/dataset/PASCAL',  # mean_rgb
+                       'yita': 0.3,
+                    'mean': [104.007, 116.669, 122.679, 137.86]},
+            'CID': {'img_height': 512,
+                    'img_width': 512,
+                    'test_list': 'test_pair.lst',
+                    'data_dir': 'C:/dataset/CID',  # mean_rgb
+                    'yita': 0.3,
+                    'mean': [104.007, 116.669, 122.679, 137.86]},
+            'NYUD': {'img_height': 425,
+                     'img_width': 560,
+                     'test_list': 'test_pair.lst',
+                     'data_dir': 'C:/dataset/NYUD',  # mean_rgb
+                     'yita': 0.5,
+                    'mean': [104.007, 116.669, 122.679, 137.86]},
+            'MDBD': {'img_height': 720,
+                         'img_width': 1280,
+                         'test_list': 'test_pair.lst',
+                         'train_list': 'train_pair.lst',
+                         'data_dir': 'C:/dataset/MDBD',  # mean_rgb
+                         'yita': 0.3,
+                         'mean': [104.007, 116.669, 122.679, 137.86]},
+            'BIPED': {'img_height': 720,  # 720
+                      'img_width': 1280,  # 1280
+                      'test_list': 'test_pair.lst',
+                      'train_list': 'train_pair0.lst',
+                      # 'train_list': 'train_rgb.lst',
+                      'data_dir': 'C:/dataset/BIPED',  # WIN: '../.../dataset/BIPED/edges'
+                      'yita': 0.5,
+                      'mean':BIPED_mean},
+            'BIPED-B2': {'img_height': 720,  # 720
+                      'img_width': 1280,  # 1280
+                      'test_list': 'test_pair.lst',
+                      'train_list': 'train_rgb.lst',
+                      'data_dir': 'C:/dataset/BIPED',  # WIN: '../.../dataset/BIPED/edges'
+                      'yita': 0.5,
+                      'mean':BIPED_mean},
+            'BIPED-B3': {'img_height': 720,  # 720
+                      'img_width': 1280,  # 1280
+                      'test_list': 'test_pair.lst',
+                      'train_list': 'train_rgb.lst',
+                      'data_dir': 'C:/dataset/BIPED',  # WIN: '../.../dataset/BIPED/edges'
+                      'yita': 0.5,
+                      'mean':BIPED_mean},
+            'BIPED-B5': {'img_height': 720,  # 720
+                         'img_width': 1280,  # 1280
+                         'test_list': 'test_pair.lst',
+                         'train_list': 'train_rgb.lst',
+                         'data_dir': 'C:/Users/xavysp/dataset/BIPED',  # WIN: '../.../dataset/BIPED/edges'
+                         'yita': 0.5,
+                         'mean':BIPED_mean},
+            'BIPED-B6': {'img_height': 720,  # 720
+                         'img_width': 1280,  # 1280
+                         'test_list': 'test_pair.lst',
+                         'train_list': 'train_rgb.lst',
+                         'data_dir': 'C:/Users/xavysp/dataset/BIPED',  # WIN: '../.../dataset/BIPED/edges'
+                         'yita': 0.5,
+                         'mean':BIPED_mean},
+            'CLASSIC': {'img_height': 512,
+                        'img_width': 512,
+                        'test_list': None,
+                        'train_list': None,
+                        'data_dir': 'teed_tmp',  # mean_rgb
+                        'yita': 0.5,
+                        'mean': [104.007, 116.669, 122.679, 137.86]},
+            'DCD': {'img_height': 240,
+                    'img_width': 360,
+                    'test_list': 'test_pair.lst',
+                    'data_dir': 'C:/dataset/DCD',  # mean_rgb
+                    'yita': 0.2,
+                    'mean': [104.007, 116.669, 122.679, 137.86]}
+        }
+    return config[dataset_name]
+
+class TestDataset(Dataset):
+    def __init__(self,
+                 data_root,
+                 test_data,
+                 img_height,
+                 img_width,
+                 test_list=None,
+                 arg=None
+                 ):
+        if test_data not in DATASET_NAMES:
+            raise ValueError(f"Unsupported dataset: {test_data}")
+
+        self.data_root = data_root
+        self.test_data = test_data
+        self.test_list = test_list
+        self.args = arg
+        self.up_scale = arg.up_scale
+        self.mean_bgr = arg.mean_test if len(arg.mean_test) == 3 else arg.mean_test[:3]
+        self.img_height = img_height
+        self.img_width = img_width
+        self.data_index = self._build_index()
+
+
+    def _build_index(self):
+        sample_indices = []
+        if self.test_data == "CLASSIC":
+            # for single image testing
+            images_path = os.listdir(self.data_root)
+            labels_path = None
+            sample_indices = [images_path, labels_path]
+        else:
+            # image and label paths are located in a list file
+
+            if not self.test_list:
+                raise ValueError(
+                    f"Test list not provided for dataset: {self.test_data}")
+
+            list_name = os.path.join(self.data_root, self.test_list)
+            if self.test_data.upper() in ['BIPED', 'BRIND','UDED','ICEDA']:
+
+                with open(list_name,encoding='utf-8') as f:
+                    files = json.load(f)
+                for pair in files:
+                    tmp_img = pair[0]
+                    tmp_gt = pair[1]
+                    sample_indices.append(
+                        (os.path.join(self.data_root, tmp_img),
+                         os.path.join(self.data_root, tmp_gt),))
+            else:
+                with open(list_name, 'r') as f:
+                    files = f.readlines()
+                files = [line.strip() for line in files]
+                pairs = [line.split() for line in files]
+
+                for pair in pairs:
+                    tmp_img = pair[0]
+                    tmp_gt = pair[1]
+                    sample_indices.append(
+                        (os.path.join(self.data_root, tmp_img),
+                         os.path.join(self.data_root, tmp_gt),))
+        return sample_indices
+
+    def __len__(self):
+        return len(self.data_index[0]) if self.test_data.upper() == 'CLASSIC' else len(self.data_index)
+
+    def __getitem__(self, idx):
+        # get data sample
+        # image_path, label_path = self.data_index[idx]
+        if self.data_index[1] is None:
+            image_path = self.data_index[0][idx] if len(self.data_index[0]) > 1 else self.data_index[0][idx - 1]
+        else:
+            image_path = self.data_index[idx][0]
+        label_path = None if self.test_data == "CLASSIC" else self.data_index[idx][1]
+        img_name = os.path.basename(image_path)
+        # print(img_name)
+        file_name = os.path.splitext(img_name)[0] + ".png"
+
+        # base dir
+        if self.test_data.upper() == 'BIPED':
+            img_dir = os.path.join(self.data_root, 'imgs', 'test')
+            gt_dir = os.path.join(self.data_root, 'edge_maps', 'test')
+        elif self.test_data.upper() == 'CLASSIC':
+            img_dir = self.data_root
+            gt_dir = None
+        else:
+            img_dir = self.data_root
+            gt_dir = self.data_root
+
+        # load data
+        image = cv2.imdecode(np.fromfile(os.path.join(img_dir, image_path), np.uint8), cv2.IMREAD_COLOR)
+        if not self.test_data == "CLASSIC":
+            label = cv2.imread(os.path.join(
+                gt_dir, label_path), cv2.IMREAD_COLOR)
+        else:
+            label = None
+
+        im_shape = [image.shape[0], image.shape[1]]
+        image, label = self.transform(img=image, gt=label)
+
+        return dict(images=image, labels=label, file_names=file_name, image_shape=im_shape)
+
+    def transform(self, img, gt):
+        # gt[gt< 51] = 0 # test without gt discrimination
+        # up scale test image
+        if self.up_scale:
+            # For TEED BIPBRIlight Upscale
+            img = cv2.resize(img,(0,0),fx=1.3,fy=1.3)
+
+        if img.shape[0] < 512 or img.shape[1] < 512:
+            #TEED BIPED standard proposal if you want speed up the test, comment this block
+            img = cv2.resize(img, (0, 0), fx=1.5, fy=1.5)
+        # else:
+        #     img = cv2.resize(img, (0, 0), fx=1.1, fy=1.1)
+
+        # Make sure images and labels are divisible by 2^4=16
+        if img.shape[0] % 8 != 0 or img.shape[1] % 8 != 0:
+            img_width = ((img.shape[1] // 8) + 1) * 8
+            img_height = ((img.shape[0] // 8) + 1) * 8
+            img = cv2.resize(img, (img_width, img_height))
+            # gt = cv2.resize(gt, (img_width, img_height))
+        else:
+            pass
+        #     img_width = self.args.test_img_width
+        #     img_height = self.args.test_img_height
+        #     img = cv2.resize(img, (img_width, img_height))
+        #     gt = cv2.resize(gt, (img_width, img_height))
+        # # For FPS
+        # img = cv2.resize(img, (496,320))
+
+        img = np.array(img, dtype=np.float32)
+        # if self.rgb:
+        #     img = img[:, :, ::-1]  # RGB->BGR
+
+        img -= self.mean_bgr
+        img = img.transpose((2, 0, 1))
+        img = torch.from_numpy(img.copy()).float()
+
+        if self.test_data == "CLASSIC":
+            gt = np.zeros((img.shape[:2]))
+            gt = torch.from_numpy(np.array([gt])).float()
+        else:
+            gt = np.array(gt, dtype=np.float32)
+            if len(gt.shape) == 3:
+                gt = gt[:, :, 0]
+            gt /= 255.
+            gt = torch.from_numpy(np.array([gt])).float()
+
+        return img, gt
+
+# *************************************************
+# ************* training **************************
+# *************************************************
+class BipedDataset(Dataset):
+    train_modes = ['train', 'test', ]
+    dataset_types = ['rgbr', ]
+    data_types = ['aug', ]
+
+    def __init__(self,
+                 data_root,
+                 img_height,
+                 img_width,
+                 train_mode='train',
+                 dataset_type='rgbr',
+                 #  is_scaling=None,
+                 # Whether to crop image or otherwise resize image to match image height and width.
+                 crop_img=False,
+                 arg=None
+                 ):
+        self.data_root = data_root
+        self.train_mode = train_mode
+        self.dataset_type = dataset_type
+        self.data_type = 'aug'  # be aware that this might change in the future
+        self.img_height = img_height
+        self.img_width = img_width
+        self.mean_bgr = arg.mean_train if len(arg.mean_train) == 3 else arg.mean_train[:3]
+        self.crop_img = crop_img
+        self.arg = arg
+
+        self.data_index = self._build_index()
+
+    def _build_index(self):
+        assert self.train_mode in self.train_modes, self.train_mode
+        assert self.dataset_type in self.dataset_types, self.dataset_type
+        assert self.data_type in self.data_types, self.data_type
+
+        data_root = os.path.abspath(self.data_root)
+        sample_indices = []
+
+        file_path = os.path.join(data_root, self.arg.train_list)
+        if self.arg.train_data.lower() == 'bsds':
+
+            with open(file_path, 'r') as f:
+                files = f.readlines()
+            files = [line.strip() for line in files]
+
+            pairs = [line.split() for line in files]
+            for pair in pairs:
+                tmp_img = pair[0]
+                tmp_gt = pair[1]
+                sample_indices.append(
+                    (os.path.join(data_root, tmp_img),
+                     os.path.join(data_root, tmp_gt),))
+        else:
+            with open(file_path) as f:
+                files = json.load(f)
+            for pair in files:
+                tmp_img = pair[0]
+                tmp_gt = pair[1]
+                sample_indices.append(
+                    (os.path.join(data_root, tmp_img),
+                     os.path.join(data_root, tmp_gt),))
+
+        return sample_indices
+
+    def __len__(self):
+        return len(self.data_index)
+
+    def __getitem__(self, idx):
+        # get data sample
+        image_path, label_path = self.data_index[idx]
+
+        # load data
+        image = cv2.imdecode(np.fromfile(image_path, np.uint8), cv2.IMREAD_COLOR)
+        label = cv2.imdecode(np.fromfile(label_path), cv2.IMREAD_GRAYSCALE)
+        image, label = self.transform(img=image, gt=label)
+        return dict(images=image, labels=label)
+
+    def transform(self, img, gt):
+        gt = np.array(gt, dtype=np.float32)
+        if len(gt.shape) == 3:
+            gt = gt[:, :, 0]
+
+        gt /= 255.  # for LDC input and BDCN
+
+        img = np.array(img, dtype=np.float32)
+        img -= self.mean_bgr
+        i_h, i_w, _ = img.shape
+        #  400 for BIPEd and 352 for BSDS check with 384
+        crop_size = self.img_height if self.img_height == self.img_width else None  # 448# MDBD=480 BIPED=480/400 BSDS=352
+        #
+        # # for BSDS 352/BRIND
+        # if i_w > crop_size and i_h > crop_size:  # later 400, before crop_size
+        #     i = random.randint(0, i_h - crop_size)
+        #     j = random.randint(0, i_w - crop_size)
+        #     img = img[i:i + crop_size, j:j + crop_size]
+        #     gt = gt[i:i + crop_size, j:j + crop_size]
+
+        # for BIPED/MDBD
+        # Second augmentation
+        if i_w> 400 and i_h>400: #before 420
+            h,w = gt.shape
+            if np.random.random() > 0.4: #before i_w> 500 and i_h>500:
+
+                LR_img_size = crop_size #l BIPED=256, 240 200 # MDBD= 352 BSDS= 176
+                i = random.randint(0, h - LR_img_size)
+                j = random.randint(0, w - LR_img_size)
+                # if img.
+                img = img[i:i + LR_img_size , j:j + LR_img_size ]
+                gt = gt[i:i + LR_img_size , j:j + LR_img_size ]
+            else:
+                LR_img_size = 300# 256 300 400  # l BIPED=208-352, # MDBD= 352-480- BSDS= 176-320
+                i = random.randint(0, h - LR_img_size)
+                j = random.randint(0, w - LR_img_size)
+                # if img.
+                img = img[i:i + LR_img_size, j:j + LR_img_size]
+                gt = gt[i:i + LR_img_size, j:j + LR_img_size]
+                img = cv2.resize(img, dsize=(crop_size, crop_size), )
+                gt = cv2.resize(gt, dsize=(crop_size, crop_size))
+
+        else:
+            # New addidings
+            img = cv2.resize(img, dsize=(crop_size, crop_size))
+            gt = cv2.resize(gt, dsize=(crop_size, crop_size))
+        # BRIND Best for TEDD+BIPED
+        gt[gt > 0.1] +=0.2#0.4
+        gt = np.clip(gt, 0., 1.)
+
+        img = img.transpose((2, 0, 1))
+        img = torch.from_numpy(img.copy()).float()
+        gt = torch.from_numpy(np.array([gt])).float()
+        return img, gt

TEED/devices.py

@@ -0,0 +1,271 @@
+import sys
+import contextlib
+from functools import lru_cache
+
+import torch
+from modules import errors, shared, npu_specific
+
+if sys.platform == "darwin":
+    from modules import mac_specific
+
+if shared.cmd_opts.use_ipex:
+    from modules import xpu_specific
+
+
+def has_xpu() -> bool:
+    return shared.cmd_opts.use_ipex and xpu_specific.has_xpu
+
+
+def has_mps() -> bool:
+    if sys.platform != "darwin":
+        return False
+    else:
+        return mac_specific.has_mps
+
+
+def cuda_no_autocast(device_id=None) -> bool:
+    if device_id is None:
+        device_id = get_cuda_device_id()
+    return (
+        torch.cuda.get_device_capability(device_id) == (7, 5)
+        and torch.cuda.get_device_name(device_id).startswith("NVIDIA GeForce GTX 16")
+    )
+
+
+def get_cuda_device_id():
+    return (
+        int(shared.cmd_opts.device_id)
+        if shared.cmd_opts.device_id is not None and shared.cmd_opts.device_id.isdigit()
+        else 0
+    ) or torch.cuda.current_device()
+
+
+def get_cuda_device_string():
+    if shared.cmd_opts.device_id is not None:
+        return f"cuda:{shared.cmd_opts.device_id}"
+
+    return "cuda"
+
+
+def get_optimal_device_name():
+    if torch.cuda.is_available():
+        return get_cuda_device_string()
+
+    if has_mps():
+        return "mps"
+
+    if has_xpu():
+        return xpu_specific.get_xpu_device_string()
+
+    if npu_specific.has_npu:
+        return npu_specific.get_npu_device_string()
+
+    return "cpu"
+
+
+def get_optimal_device():
+    return torch.device(get_optimal_device_name())
+
+
+def get_device_for(task):
+    if task in shared.cmd_opts.use_cpu or "all" in shared.cmd_opts.use_cpu:
+        return cpu
+
+    return get_optimal_device()
+
+
+def torch_gc():
+
+    if torch.cuda.is_available():
+        with torch.cuda.device(get_cuda_device_string()):
+            torch.cuda.empty_cache()
+            torch.cuda.ipc_collect()
+
+    if has_mps():
+        mac_specific.torch_mps_gc()
+
+    if has_xpu():
+        xpu_specific.torch_xpu_gc()
+
+    if npu_specific.has_npu:
+        torch_npu_set_device()
+        npu_specific.torch_npu_gc()
+
+
+def torch_npu_set_device():
+    # Work around due to bug in torch_npu, revert me after fixed, @see https://gitee.com/ascend/pytorch/issues/I8KECW?from=project-issue
+    if npu_specific.has_npu:
+        torch.npu.set_device(0)
+
+
+def enable_tf32():
+    if torch.cuda.is_available():
+
+        # enabling benchmark option seems to enable a range of cards to do fp16 when they otherwise can't
+        # see https://github.com/AUTOMATIC1111/stable-diffusion-webui/pull/4407
+        if cuda_no_autocast():
+            torch.backends.cudnn.benchmark = True
+
+        torch.backends.cuda.matmul.allow_tf32 = True
+        torch.backends.cudnn.allow_tf32 = True
+
+
+errors.run(enable_tf32, "Enabling TF32")
+
+cpu: torch.device = torch.device("cpu")
+fp8: bool = False
+device: torch.device = None
+device_interrogate: torch.device = None
+device_gfpgan: torch.device = None
+device_esrgan: torch.device = None
+device_codeformer: torch.device = None
+dtype: torch.dtype = torch.float16
+dtype_vae: torch.dtype = torch.float16
+dtype_unet: torch.dtype = torch.float16
+dtype_inference: torch.dtype = torch.float16
+unet_needs_upcast = False
+
+
+def cond_cast_unet(input):
+    return input.to(dtype_unet) if unet_needs_upcast else input
+
+
+def cond_cast_float(input):
+    return input.float() if unet_needs_upcast else input
+
+
+nv_rng = None
+patch_module_list = [
+    torch.nn.Linear,
+    torch.nn.Conv2d,
+    torch.nn.MultiheadAttention,
+    torch.nn.GroupNorm,
+    torch.nn.LayerNorm,
+]
+
+
+def manual_cast_forward(target_dtype):
+    def forward_wrapper(self, *args, **kwargs):
+        if any(
+            isinstance(arg, torch.Tensor) and arg.dtype != target_dtype
+            for arg in args
+        ):
+            args = [arg.to(target_dtype) if isinstance(arg, torch.Tensor) else arg for arg in args]
+            kwargs = {k: v.to(target_dtype) if isinstance(v, torch.Tensor) else v for k, v in kwargs.items()}
+
+        org_dtype = target_dtype
+        for param in self.parameters():
+            if param.dtype != target_dtype:
+                org_dtype = param.dtype
+                break
+
+        if org_dtype != target_dtype:
+            self.to(target_dtype)
+        result = self.org_forward(*args, **kwargs)
+        if org_dtype != target_dtype:
+            self.to(org_dtype)
+
+        if target_dtype != dtype_inference:
+            if isinstance(result, tuple):
+                result = tuple(
+                    i.to(dtype_inference)
+                    if isinstance(i, torch.Tensor)
+                    else i
+                    for i in result
+                )
+            elif isinstance(result, torch.Tensor):
+                result = result.to(dtype_inference)
+        return result
+    return forward_wrapper
+
+
+@contextlib.contextmanager
+def manual_cast(target_dtype):
+    applied = False
+    for module_type in patch_module_list:
+        if hasattr(module_type, "org_forward"):
+            continue
+        applied = True
+        org_forward = module_type.forward
+        if module_type == torch.nn.MultiheadAttention:
+            module_type.forward = manual_cast_forward(torch.float32)
+        else:
+            module_type.forward = manual_cast_forward(target_dtype)
+        module_type.org_forward = org_forward
+    try:
+        yield None
+    finally:
+        if applied:
+            for module_type in patch_module_list:
+                if hasattr(module_type, "org_forward"):
+                    module_type.forward = module_type.org_forward
+                    delattr(module_type, "org_forward")
+
+
+def autocast(disable=False):
+    if disable:
+        return contextlib.nullcontext()
+
+    if fp8 and device==cpu:
+        return torch.autocast("cpu", dtype=torch.bfloat16, enabled=True)
+
+    if fp8 and dtype_inference == torch.float32:
+        return manual_cast(dtype)
+
+    if dtype == torch.float32 or dtype_inference == torch.float32:
+        return contextlib.nullcontext()
+
+    if has_xpu() or has_mps() or cuda_no_autocast():
+        return manual_cast(dtype)
+
+    return torch.autocast("cuda")
+
+
+def without_autocast(disable=False):
+    return torch.autocast("cuda", enabled=False) if torch.is_autocast_enabled() and not disable else contextlib.nullcontext()
+
+
+class NansException(Exception):
+    pass
+
+
+def test_for_nans(x, where):
+    if shared.cmd_opts.disable_nan_check:
+        return
+
+    if not torch.all(torch.isnan(x)).item():
+        return
+
+    if where == "unet":
+        message = "A tensor with all NaNs was produced in Unet."
+
+        if not shared.cmd_opts.no_half:
+            message += " This could be either because there's not enough precision to represent the picture, or because your video card does not support half type. Try setting the \"Upcast cross attention layer to float32\" option in Settings > Stable Diffusion or using the --no-half commandline argument to fix this."
+
+    elif where == "vae":
+        message = "A tensor with all NaNs was produced in VAE."
+
+        if not shared.cmd_opts.no_half and not shared.cmd_opts.no_half_vae:
+            message += " This could be because there's not enough precision to represent the picture. Try adding --no-half-vae commandline argument to fix this."
+    else:
+        message = "A tensor with all NaNs was produced."
+
+    message += " Use --disable-nan-check commandline argument to disable this check."
+
+    raise NansException(message)
+
+
+@lru_cache
+def first_time_calculation():
+    """
+    just do any calculation with pytorch layers - the first time this is done it allocaltes about 700MB of memory and
+    spends about 2.7 seconds doing that, at least with NVidia.
+    """
+
+    x = torch.zeros((1, 1)).to(device, dtype)
+    linear = torch.nn.Linear(1, 1).to(device, dtype)
+    linear(x)
+
+    x = torch.zeros((1, 1, 3, 3)).to(device, dtype)
+    conv2d = torch.nn.Conv2d(1, 1, (3, 3)).to(device, dtype)
+    conv2d(x)

TEED/loss2.py

@@ -0,0 +1,92 @@
+import torch
+import torch.nn.functional as F
+from TEED.utils.AF.Fsmish import smish as Fsmish
+
+def bdcn_loss2(inputs, targets, l_weight=1.1):
+    # bdcn loss modified in DexiNed
+
+    targets = targets.long()
+    mask = targets.float()
+    num_positive = torch.sum((mask > 0.0).float()).float() # >0.1
+    num_negative = torch.sum((mask <= 0.0).float()).float() # <= 0.1
+
+    mask[mask > 0.] = 1.0 * num_negative / (num_positive + num_negative) #0.1
+    mask[mask <= 0.] = 1.1 * num_positive / (num_positive + num_negative)  # before mask[mask <= 0.1]
+    inputs= torch.sigmoid(inputs)
+    cost = torch.nn.BCELoss(mask, reduction='none')(inputs, targets.float())
+    cost = torch.sum(cost.float().mean((1, 2, 3))) # before sum
+    return l_weight*cost
+
+# ------------ cats losses ----------
+def bdrloss(prediction, label, radius,device='cpu'):
+    '''
+    The boundary tracing loss that handles the confusing pixels.
+    '''
+
+    filt = torch.ones(1, 1, 2*radius+1, 2*radius+1)
+    filt.requires_grad = False
+    filt = filt.to(device)
+
+    bdr_pred = prediction * label
+    pred_bdr_sum = label * F.conv2d(bdr_pred, filt, bias=None, stride=1, padding=radius)
+
+    texture_mask = F.conv2d(label.float(), filt, bias=None, stride=1, padding=radius)
+    mask = (texture_mask != 0).float()
+    mask[label == 1] = 0
+    pred_texture_sum = F.conv2d(prediction * (1-label) * mask, filt, bias=None, stride=1, padding=radius)
+
+    softmax_map = torch.clamp(pred_bdr_sum / (pred_texture_sum + pred_bdr_sum + 1e-10), 1e-10, 1 - 1e-10)
+    cost = -label * torch.log(softmax_map)
+    cost[label == 0] = 0
+
+    return torch.sum(cost.float().mean((1, 2, 3)))
+
+def textureloss(prediction, label, mask_radius, device='cpu'):
+    '''
+    The texture suppression loss that smooths the texture regions.
+    '''
+    filt1 = torch.ones(1, 1, 3, 3)
+    filt1.requires_grad = False
+    filt1 = filt1.to(device)
+    filt2 = torch.ones(1, 1, 2*mask_radius+1, 2*mask_radius+1)
+    filt2.requires_grad = False
+    filt2 = filt2.to(device)
+
+    pred_sums = F.conv2d(prediction.float(), filt1, bias=None, stride=1, padding=1)
+    label_sums = F.conv2d(label.float(), filt2, bias=None, stride=1, padding=mask_radius)
+
+    mask = 1 - torch.gt(label_sums, 0).float()
+
+    loss = -torch.log(torch.clamp(1-pred_sums/9, 1e-10, 1-1e-10))
+    loss[mask == 0] = 0
+
+    return torch.sum(loss.float().mean((1, 2, 3)))
+
+
+def cats_loss(prediction, label, l_weight=[0.,0.], device='cpu'):
+    # tracingLoss
+
+    tex_factor,bdr_factor = l_weight
+    balanced_w = 1.1
+    label = label.float()
+    prediction = prediction.float()
+    with torch.no_grad():
+        mask = label.clone()
+
+        num_positive = torch.sum((mask == 1).float()).float()
+        num_negative = torch.sum((mask == 0).float()).float()
+        beta = num_negative / (num_positive + num_negative)
+        mask[mask == 1] = beta
+        mask[mask == 0] = balanced_w * (1 - beta)
+        mask[mask == 2] = 0
+
+    prediction = torch.sigmoid(prediction)
+
+    cost = torch.nn.functional.binary_cross_entropy(
+        prediction.float(), label.float(), weight=mask, reduction='none')
+    cost = torch.sum(cost.float().mean((1, 2, 3)))  # by me
+    label_w = (label != 0).float()
+    textcost = textureloss(prediction.float(), label_w.float(), mask_radius=4, device=device)
+    bdrcost = bdrloss(prediction.float(), label_w.float(), radius=4, device=device)
+
+    return cost + bdr_factor * bdrcost + tex_factor * textcost

TEED/ted.py

@@ -0,0 +1,297 @@
+# TEED: is a Tiny but Efficient Edge Detection, it comes from the LDC-B3
+# with a Slightly modification
+# LDC parameters:
+# 155665
+# TED > 58K
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from TEED.utils.AF.Fsmish import smish as Fsmish
+from TEED.utils.AF.Xsmish import Smish
+from TEED.utils.img_processing import count_parameters
+
+
+def weight_init(m):
+    if isinstance(m, (nn.Conv2d,)):
+        torch.nn.init.xavier_normal_(m.weight, gain=1.0)
+
+        if m.bias is not None:
+            torch.nn.init.zeros_(m.bias)
+
+    # for fusion layer
+    if isinstance(m, (nn.ConvTranspose2d,)):
+        torch.nn.init.xavier_normal_(m.weight, gain=1.0)
+        if m.bias is not None:
+            torch.nn.init.zeros_(m.bias)
+
+class CoFusion(nn.Module):
+    # from LDC
+
+    def __init__(self, in_ch, out_ch):
+        super(CoFusion, self).__init__()
+        self.conv1 = nn.Conv2d(in_ch, 32, kernel_size=3,
+                               stride=1, padding=1) # before 64
+        self.conv3= nn.Conv2d(32, out_ch, kernel_size=3,
+                               stride=1, padding=1)# before 64  instead of 32
+        self.relu = nn.ReLU()
+        self.norm_layer1 = nn.GroupNorm(4, 32) # before 64
+
+    def forward(self, x):
+        # fusecat = torch.cat(x, dim=1)
+        attn = self.relu(self.norm_layer1(self.conv1(x)))
+        attn = F.softmax(self.conv3(attn), dim=1)
+        return ((x * attn).sum(1)).unsqueeze(1)
+
+
+class CoFusion2(nn.Module):
+        # TEDv14-3
+    def __init__(self, in_ch, out_ch):
+        super(CoFusion2, self).__init__()
+        self.conv1 = nn.Conv2d(in_ch, 32, kernel_size=3,
+                               stride=1, padding=1) # before 64
+        # self.conv2 = nn.Conv2d(32, 32, kernel_size=3,
+        #                        stride=1, padding=1)# before 64
+        self.conv3 = nn.Conv2d(32, out_ch, kernel_size=3,
+                               stride=1, padding=1)# before 64  instead of 32
+        self.smish= Smish()#nn.ReLU(inplace=True)
+
+
+    def forward(self, x):
+        # fusecat = torch.cat(x, dim=1)
+        attn = self.conv1(self.smish(x))
+        attn = self.conv3(self.smish(attn)) # before , )dim=1)
+
+        # return ((fusecat * attn).sum(1)).unsqueeze(1)
+        return ((x * attn).sum(1)).unsqueeze(1)
+
+class DoubleFusion(nn.Module):
+    # TED fusion before the final edge map prediction
+    def __init__(self, in_ch, out_ch):
+        super(DoubleFusion, self).__init__()
+        self.DWconv1 = nn.Conv2d(in_ch, in_ch*8, kernel_size=3,
+                               stride=1, padding=1, groups=in_ch) # before 64
+        self.PSconv1 = nn.PixelShuffle(1)
+
+        self.DWconv2 = nn.Conv2d(24, 24*1, kernel_size=3,
+                               stride=1, padding=1,groups=24)# before 64  instead of 32
+
+        self.AF= Smish()#XAF() #nn.Tanh()# XAF() #   # Smish()#
+
+
+    def forward(self, x):
+        # fusecat = torch.cat(x, dim=1)
+        attn = self.PSconv1(self.DWconv1(self.AF(x))) # #TEED best res TEDv14 [8, 32, 352, 352]
+
+        attn2 = self.PSconv1(self.DWconv2(self.AF(attn))) # #TEED best res TEDv14[8, 3, 352, 352]
+
+        return Fsmish(((attn2 +attn).sum(1)).unsqueeze(1)) #TED best res
+
+class _DenseLayer(nn.Sequential):
+    def __init__(self, input_features, out_features):
+        super(_DenseLayer, self).__init__()
+
+        self.add_module('conv1', nn.Conv2d(input_features, out_features,
+                                           kernel_size=3, stride=1, padding=2, bias=True)),
+        self.add_module('smish1', Smish()),
+        self.add_module('conv2', nn.Conv2d(out_features, out_features,
+                                           kernel_size=3, stride=1, bias=True))
+    def forward(self, x):
+        x1, x2 = x
+
+        new_features = super(_DenseLayer, self).forward(Fsmish(x1))  # F.relu()
+
+        return 0.5 * (new_features + x2), x2
+
+
+class _DenseBlock(nn.Sequential):
+    def __init__(self, num_layers, input_features, out_features):
+        super(_DenseBlock, self).__init__()
+        for i in range(num_layers):
+            layer = _DenseLayer(input_features, out_features)
+            self.add_module('denselayer%d' % (i + 1), layer)
+            input_features = out_features
+
+
+class UpConvBlock(nn.Module):
+    def __init__(self, in_features, up_scale):
+        super(UpConvBlock, self).__init__()
+        self.up_factor = 2
+        self.constant_features = 16
+
+        layers = self.make_deconv_layers(in_features, up_scale)
+        assert layers is not None, layers
+        self.features = nn.Sequential(*layers)
+
+    def make_deconv_layers(self, in_features, up_scale):
+        layers = []
+        all_pads=[0,0,1,3,7]
+        for i in range(up_scale):
+            kernel_size = 2 ** up_scale
+            pad = all_pads[up_scale]  # kernel_size-1
+            out_features = self.compute_out_features(i, up_scale)
+            layers.append(nn.Conv2d(in_features, out_features, 1))
+            layers.append(Smish())
+            layers.append(nn.ConvTranspose2d(
+                out_features, out_features, kernel_size, stride=2, padding=pad))
+            in_features = out_features
+        return layers
+
+    def compute_out_features(self, idx, up_scale):
+        return 1 if idx == up_scale - 1 else self.constant_features
+
+    def forward(self, x):
+        return self.features(x)
+
+
+class SingleConvBlock(nn.Module):
+    def __init__(self, in_features, out_features, stride, use_ac=False):
+        super(SingleConvBlock, self).__init__()
+        # self.use_bn = use_bs
+        self.use_ac=use_ac
+        self.conv = nn.Conv2d(in_features, out_features, 1, stride=stride,
+                              bias=True)
+        if self.use_ac:
+            self.smish = Smish()
+
+    def forward(self, x):
+        x = self.conv(x)
+        if self.use_ac:
+            return self.smish(x)
+        else:
+            return x
+
+class DoubleConvBlock(nn.Module):
+    def __init__(self, in_features, mid_features,
+                 out_features=None,
+                 stride=1,
+                 use_act=True):
+        super(DoubleConvBlock, self).__init__()
+
+        self.use_act = use_act
+        if out_features is None:
+            out_features = mid_features
+        self.conv1 = nn.Conv2d(in_features, mid_features,
+                               3, padding=1, stride=stride)
+        self.conv2 = nn.Conv2d(mid_features, out_features, 3, padding=1)
+        self.smish= Smish()#nn.ReLU(inplace=True)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.smish(x)
+        x = self.conv2(x)
+        if self.use_act:
+            x = self.smish(x)
+        return x
+
+
+class TED(nn.Module):
+    """ Definition of  Tiny and Efficient Edge Detector
+    model
+    """
+
+    def __init__(self):
+        super(TED, self).__init__()
+        self.block_1 = DoubleConvBlock(3, 16, 16, stride=2,)
+        self.block_2 = DoubleConvBlock(16, 32, use_act=False)
+        self.dblock_3 = _DenseBlock(1, 32, 48) # [32,48,100,100] before (2, 32, 64)
+
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+
+        # skip1 connection, see fig. 2
+        self.side_1 = SingleConvBlock(16, 32, 2)
+
+        # skip2 connection, see fig. 2
+        self.pre_dense_3 = SingleConvBlock(32, 48, 1)  # before (32, 64, 1)
+
+        # USNet
+        self.up_block_1 = UpConvBlock(16, 1)
+        self.up_block_2 = UpConvBlock(32, 1)
+        self.up_block_3 = UpConvBlock(48, 2) # (32, 64, 1)
+
+        self.block_cat = DoubleFusion(3,3) # TEED: DoubleFusion
+
+        self.apply(weight_init)
+
+    def slice(self, tensor, slice_shape):
+        t_shape = tensor.shape
+        img_h, img_w = slice_shape
+        if img_w!=t_shape[-1] or img_h!=t_shape[2]:
+            new_tensor = F.interpolate(
+                tensor, size=(img_h, img_w), mode='bicubic',align_corners=False)
+
+        else:
+            new_tensor=tensor
+        # tensor[..., :height, :width]
+        return new_tensor
+    def resize_input(self,tensor):
+        t_shape = tensor.shape
+        if t_shape[2] % 8 != 0 or t_shape[3] % 8 != 0:
+            img_w= ((t_shape[3]// 8) + 1) * 8
+            img_h = ((t_shape[2] // 8) + 1) * 8
+            new_tensor = F.interpolate(
+                tensor, size=(img_h, img_w), mode='bicubic', align_corners=False)
+        else:
+            new_tensor = tensor
+        return new_tensor
+
+    def crop_bdcn(data1, h, w, crop_h, crop_w):
+        # Based on BDCN Implementation @ https://github.com/pkuCactus/BDCN
+        _, _, h1, w1 = data1.size()
+        assert (h <= h1 and w <= w1)
+        data = data1[:, :, crop_h:crop_h + h, crop_w:crop_w + w]
+        return data
+
+
+    def forward(self, x, single_test=False):
+        assert x.ndim == 4, x.shape
+         # supose the image size is 352x352
+
+        # Block 1
+        block_1 = self.block_1(x) # [8,16,176,176]
+        block_1_side = self.side_1(block_1) # 16 [8,32,88,88]
+
+        # Block 2
+        block_2 = self.block_2(block_1) # 32 # [8,32,176,176]
+        block_2_down = self.maxpool(block_2) # [8,32,88,88]
+        block_2_add = block_2_down + block_1_side # [8,32,88,88]
+
+        # Block 3
+        block_3_pre_dense = self.pre_dense_3(block_2_down) # [8,64,88,88] block 3 L connection
+        block_3, _ = self.dblock_3([block_2_add, block_3_pre_dense]) # [8,64,88,88]
+
+        # upsampling blocks
+        out_1 = self.up_block_1(block_1)
+        out_2 = self.up_block_2(block_2)
+        out_3 = self.up_block_3(block_3)
+
+        results = [out_1, out_2, out_3]
+
+        # concatenate multiscale outputs
+        block_cat = torch.cat(results, dim=1)  # Bx6xHxW
+        block_cat = self.block_cat(block_cat)  # Bx1xHxW DoubleFusion
+
+        results.append(block_cat)
+        return results
+
+
+if __name__ == '__main__':
+    batch_size = 8
+    img_height = 352
+    img_width = 352
+
+    # device = "cuda" if torch.cuda.is_available() else "cpu"
+    device = "cpu"
+    input = torch.rand(batch_size, 3, img_height, img_width).to(device)
+    # target = torch.rand(batch_size, 1, img_height, img_width).to(device)
+    print(f"input shape: {input.shape}")
+    model = TED().to(device)
+    output = model(input)
+    print(f"output shapes: {[t.shape for t in output]}")
+
+    # for i in range(20000):
+    #     print(i)
+    #     output = model(input)
+    #     loss = nn.MSELoss()(output[-1], target)
+    #     loss.backward()

TEED/util.py

@@ -0,0 +1,78 @@
+import numpy as np
+import cv2
+import os
+
+
+def load_model(filename: str, remote_url: str, model_dir: str) -> str:
+    """
+    Load the model from the specified filename and remote URL if it doesn't exist locally.
+
+    Args:
+        filename (str): The filename of the model.
+        remote_url (str): The remote URL of the model.
+    """
+    local_path = os.path.join(model_dir, filename)
+    if not os.path.exists(local_path):
+        from scripts.utils import load_file_from_url
+        load_file_from_url(remote_url, model_dir=model_dir)
+    return local_path
+
+
+def HWC3(x):
+    assert x.dtype == np.uint8
+    if x.ndim == 2:
+        x = x[:, :, None]
+    assert x.ndim == 3
+    H, W, C = x.shape
+    assert C == 1 or C == 3 or C == 4
+    if C == 3:
+        return x
+    if C == 1:
+        return np.concatenate([x, x, x], axis=2)
+    if C == 4:
+        color = x[:, :, 0:3].astype(np.float32)
+        alpha = x[:, :, 3:4].astype(np.float32) / 255.0
+        y = color * alpha + 255.0 * (1.0 - alpha)
+        y = y.clip(0, 255).astype(np.uint8)
+        return y
+
+
+def make_noise_disk(H, W, C, F):
+    noise = np.random.uniform(low=0, high=1, size=((H // F) + 2, (W // F) + 2, C))
+    noise = cv2.resize(noise, (W + 2 * F, H + 2 * F), interpolation=cv2.INTER_CUBIC)
+    noise = noise[F: F + H, F: F + W]
+    noise -= np.min(noise)
+    noise /= np.max(noise)
+    if C == 1:
+        noise = noise[:, :, None]
+    return noise
+
+
+def nms(x, t, s):
+    x = cv2.GaussianBlur(x.astype(np.float32), (0, 0), s)
+
+    f1 = np.array([[0, 0, 0], [1, 1, 1], [0, 0, 0]], dtype=np.uint8)
+    f2 = np.array([[0, 1, 0], [0, 1, 0], [0, 1, 0]], dtype=np.uint8)
+    f3 = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=np.uint8)
+    f4 = np.array([[0, 0, 1], [0, 1, 0], [1, 0, 0]], dtype=np.uint8)
+
+    y = np.zeros_like(x)
+
+    for f in [f1, f2, f3, f4]:
+        np.putmask(y, cv2.dilate(x, kernel=f) == x, x)
+
+    z = np.zeros_like(y, dtype=np.uint8)
+    z[y > t] = 255
+    return z
+
+
+def min_max_norm(x):
+    x -= np.min(x)
+    x /= np.maximum(np.max(x), 1e-5)
+    return x
+
+
+def safe_step(x, step=2):
+    y = x.astype(np.float32) * float(step + 1)
+    y = y.astype(np.int32).astype(np.float32) / float(step)
+    return y