Initial demo

.gitattributes

@@ -6,6 +6,7 @@
 *.ftz filter=lfs diff=lfs merge=lfs -text
 *.gz filter=lfs diff=lfs merge=lfs -text
 *.h5 filter=lfs diff=lfs merge=lfs -text
+*.jpg filter=lfs diff=lfs merge=lfs -text
 *.joblib filter=lfs diff=lfs merge=lfs -text
 *.lfs.* filter=lfs diff=lfs merge=lfs -text
 *.mlmodel filter=lfs diff=lfs merge=lfs -text
@@ -32,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
+pretrained_models filter=lfs diff=lfs merge=lfs -text
+sample_images filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,2 @@
+Demo.ipynb
+__pycache__

DenseMammogram/.gitignore

@@ -0,0 +1,2 @@
+pretrained_models
+__pycache__

DenseMammogram/README.md

@@ -0,0 +1,54 @@
+# Deep Learning for Detection of Iso-Sense, Obscure Masses in Mammographically Dense Breasts
+[![report](https://img.shields.io/badge/arxiv-report-red)](https://arxiv.org/abs/) [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/) 
+
+## Introduction
+Deep Learning for Detection of Iso-Sense, Obscure Masses in Mammographically Dense Breasts is a paper on object detection method for finding malignant masses in breast mammograms. Our model is particularly useful for dense breasts and iso-dense and obscure masses. In this paper we have included code and pretrained weights for the paper along with all the scripts to replicate numbers in the paper(Our private dataset is not included).  
+
+## Getting Started
+
+
+First clone the repo:
+```bash
+git clone https://github.com/Pranjal2041/DenseMammograms.git
+```
+
+Next setup the enviornment using `conda` or `virtualenv`: 
+```bash 
+1. conda create -n densebreast python=3.7
+conda activate densebreast
+pip install -r requirements.txt
+
+or
+
+2. python -m venv densebreast
+source densebreast/bin/activate
+pip install -r requirements.txt
+```
+
+## Pretrained Weights
+
+You can download the pretrained models from this [url](https://csciitd-my.sharepoint.com/:f:/g/personal/cs5190443_iitd_ac_in/ElTbduIuI49EougSH05Tb4IBhbc5gXCrlok0X_xvAI196g?e=Ss2eS1) in the current directory.
+<br>
+
+## Running the Code
+
+To generate predictions and FROC graphs using the pretrained models, run:
+`python all_graphs.py`
+
+For running individual models on other datasets, geenerate_{dataset}_preds.py have been provided.
+For example to run predictions on inbreast, run:
+`python geenerate_inbreast_preds.py`
+
+
+## Demo
+
+You can either use **Google Colab Demo** or **Huggingface demo**
+
+## Citation
+
+Details Coming Soon!
+
+## License
+
+TODO: Add License
+

DenseMammogram/advanced_config.py

@@ -0,0 +1,36 @@
+import json
+import os
+
+class AdvancedConfig:
+
+    def save(self, file):
+        os.makedirs(os.path.split(file)[0], exist_ok=True)
+        json.dump(self.config, open(file, 'w'), indent=4)
+
+    def read_cfg(self, file):
+        # Its a json file with comments
+        new_lines = []
+        for line in open(file).readlines():
+            if line.find("#")!=-1:
+                new_lines.append(line[:line.find("#")])
+            else:
+                new_lines.append(line)
+        return json.loads('\n'.join(new_lines))
+        
+
+    def merge_config(self, cfg_dict, base_dict):
+        for key in cfg_dict:  
+            if key not in base_dict:
+                # Strange, raise an error
+                raise Exception(f'Key {key} not found in base config')
+            if isinstance(cfg_dict[key], dict):
+                base_dict[key] = self.merge_config(cfg_dict[key], base_dict[key])
+            else:
+                base_dict[key] = cfg_dict[key]
+        return base_dict
+
+    def __init__(self, file, base_file = 'configs/default.cfg') -> None:
+        self.default_config = self.read_cfg(base_file)
+        self.new_config = self.read_cfg(file)
+        self.config = self.merge_config(self.new_config, self.default_config)
+        

DenseMammogram/advanced_logger.py

@@ -0,0 +1,48 @@
+# An Advanced Logger class which writes data in a well formatted manner 
+# to files based on different priorities.
+
+from enum import Enum
+import os
+import datetime
+import time
+
+class LogPriority(Enum):
+    """
+    Enum class for different log priorities.
+    """
+    LOW = 0
+    MEDIUM = 1
+    HIGH = 2
+    STATS = 3
+
+class AdvancedLogger:
+
+    def __init__(self, base_dir):
+        self.base_dir =  base_dir
+        self.files = []
+        self.file_names = []
+        for p in LogPriority:
+            self.file_names.append(os.path.join(self.base_dir, f'Log_{p.name}' + '.log'))
+            self.files.append(open(self.file_names[-1], 'w'))
+        self.last_log_time = -1
+
+    def flush(self):
+        for f in self.files:
+                f.close()
+        for i in range(len(self.files)):
+            self.files[i] = open(self.file_names[i], 'a')
+
+    def log(self, *args, priority = LogPriority.LOW):
+        to_log = ' '.join(map(str, args))
+        if priority.value <= LogPriority.MEDIUM.value:
+            # Add current time to to_log
+            now = datetime.datetime.now()
+            to_log = f'[{now.strftime("%H:%M:%S")}]: {to_log}' 
+        print(to_log)
+        for p in range(priority.value+1):
+            self.files[p].write(to_log + '\n')
+        
+        # If time - last_log_time is greater than 10s or Priority is HIGH or above close the file and re-open in append mode
+        if time.time() - self.last_log_time > 10 or priority.value >= LogPriority.HIGH.value: 
+            self.flush()
+            self.last_log_time = time.time()

DenseMammogram/all_graphs.py

@@ -0,0 +1,156 @@
+import os
+from os.path import join
+from merge_predictions import get_image_dict, apply_merge
+from froc_by_pranjal import calc_froc_from_dict, pretty_print_fps
+import numpy as np 
+import matplotlib.pyplot as plt
+
+
+OUT_DIR = 'euro_results_auto'
+numbers_dir = os.path.join(OUT_DIR, 'numbers')
+graphs_dir = os.path.join(OUT_DIR, 'graphs')
+
+BASE_FOLDER = '../bilateral_new/MammoDatasets'
+
+MIN_CLIP_FPI = 0.02
+def plot_froc(input_files, save_file, TITLE = 'FRCNN vs BILATERAL FROC', SHOW = False, CLIP_FPI = 1.2):
+    for file in input_files:
+        lines = open(file).readlines()
+        x = np.array([float(line.split()[0]) for line in lines])
+        y = np.array([float(line.split()[1]) for line in lines])
+        y = y[x<CLIP_FPI]
+        x = x[x<CLIP_FPI]
+        y = y[MIN_CLIP_FPI<x]
+        x = x[MIN_CLIP_FPI<x]
+        plt.plot(x, y, label = input_files[file])
+        plt.legend()
+
+    plt.title(TITLE)
+    plt.xlabel('Average False Positive Per Image')
+    plt.ylabel('Sensitivity')
+
+    if SHOW:
+        plt.show()
+    plt.savefig(save_file)
+    plt.clf()
+
+
+dsets = [('AIIMS_highres_reliable', 'AIIMS'), ('IRCHVal', 'IRCHVal')]
+dsets = dsets[1:]
+for dset in dsets:
+    test_splits = ['test_2', 'test_dense', 'test_iso'][::-1]
+    for test_split in test_splits:
+        main_dataset = join(BASE_FOLDER, dset[0], test_split)
+
+        contrast_datasets =  [join(BASE_FOLDER,f'{dset[1]}_C{i+1}',test_split) for i in range(4)]
+        threshold_datasets = [join(BASE_FOLDER,f'{dset[1]}_T{i+1}',test_split) for i in range(2)]
+        frcnn_preds = 'preds_frcnn_frcnn'
+        contrast_preds = [
+            'preds_frcnn_AIIMS_C1',
+            'preds_frcnn_AIIMS_C2',
+            'preds_frcnn_AIIMS_C3',
+            'preds_frcnn_AIIMS_C4',
+        ]
+        bilateral_preds = 'preds_bilateral_BILATERAL'
+        threshold_preds = [
+            'preds_frcnn_AIIMS_T1',
+            'preds_frcnn_AIIMS_T2',
+        ]
+
+        input_files = []
+        dataset_paths = [join(main_dataset,  '{0}', frcnn_preds)]
+        dataset_paths +=[join(dset, '{0}', preds) for (dset,preds) in zip(contrast_datasets, contrast_preds)]
+        dataset_paths +=[join(dset, '{0}', preds) for (dset,preds) in zip(threshold_datasets, threshold_preds)]
+        dataset_paths +=[join(main_dataset, '{0}', bilateral_preds)]
+
+
+        CONFIGS = {
+            'Baseline' : ('Baseline Model', [0]),
+            'Bilateral' : ('Bilateral Model', [7]),
+            'Contrast' : ('CABD Model', [0,1,2,3,4]),
+            'Threshold' : ('TI Model', [0,5,6]),
+            'Proposed' : ('Proposed Model', [1,2,3,4,5,6,7])
+        }
+
+        # Now handle the directories
+        num_dir = os.path.join(numbers_dir, dset[1], test_split)
+        os.makedirs(num_dir, exist_ok=True)
+
+
+        for config in CONFIGS:
+            title = CONFIGS[config][0]
+            allowed = CONFIGS[config][1]
+
+            weight_map = {
+                0 : 1.,
+                1 : 1,
+                2 : 1.,
+                3 : 1.,
+                4 : .5, # C4
+                5 : 0.5,
+                6 : 0.5,
+                7 : 1
+            }
+
+            weights = [weight_map[x] for x in allowed]
+            
+            # generate the required mp dicts
+            def c2_manp(preds):
+                preds  = list(filter(lambda x: x[0]>0.85,preds)) # keep preds lower than 0.6 confidence
+                return preds
+
+            def c3_manp(preds):
+                preds  = list(filter(lambda x: x[0]>0.85,preds)) # keep preds lower than 0.6 confidence
+                return preds
+
+            def t1_manp(preds):
+                preds  = list(filter(lambda x: x[0]>0.6,preds)) # keep preds lower than 0.6 confidence
+                return preds
+
+            t2_manp = t1_manp
+            mp_dict = {
+                f'{dset[1]}_C2' : c2_manp,
+                f'{dset[1]}_C3' : c3_manp,
+                f'{dset[1]}_T1' : t1_manp,
+                f'{dset[1]}_T2' : t2_manp,
+                f'{dset[1]}_C4' : c3_manp
+            }
+
+            image_dict = get_image_dict(dataset_paths, allowed = allowed, USE_ACR = False, acr_cat = None, mp_dict = mp_dict)
+            image_dict = apply_merge(image_dict, METHOD = 'nms', weights= weights, conf_type='absent_model_aware_avg')
+            
+
+            senses, fps = calc_froc_from_dict(image_dict, fps_req = [0.025,0.05,0.1,0.15,0.2,0.3,1.], save_to = os.path.join(num_dir, f'{title}.txt'))
+
+
+        # Lets plot now
+
+        GRAPHS = [
+            ('Bilateral','Baseline'),
+            ('Contrast','Baseline'),
+            ('Threshold','Baseline'),
+            ('Proposed','Baseline'),
+            ('Proposed', 'Bilateral'),
+            ('Proposed', 'Contrast'),
+            ('Proposed', 'Threshold'),
+        ]
+
+
+        # Now handle the directories
+        graph_dir = os.path.join(graphs_dir, dset[1], test_split)
+        os.makedirs(graph_dir, exist_ok=True)
+
+        for graph in GRAPHS:
+            if graph[0] not in CONFIGS or graph[1] not in CONFIGS: continue
+            file_name1 = f'{CONFIGS[graph[0]][0]}.txt'
+            file_name2 = f'{CONFIGS[graph[1]][0]}.txt'
+
+            title1 = CONFIGS[graph[0]][0]
+            title2 = CONFIGS[graph[1]][0]
+
+            plot_froc({
+                join(num_dir, file_name1): title1,
+                join(num_dir, file_name2) : title2,
+            }, join(graph_dir,f'{title1}_vs_{title2}.png'),f'{title1} vs {title2} FROC', CLIP_FPI = 0.3 if dset[0] == 'IRCHVal' else 0.8)
+
+

DenseMammogram/auc_by_pranjal.py

@@ -0,0 +1,120 @@
+import os
+from os.path import join
+import glob
+from sklearn.metrics import roc_auc_score, roc_curve
+import sys
+
+def file_to_score(file):
+    try:
+        content = open(file, 'r').readlines()
+        st = 0
+        if len(content) == 0:
+            # Empty File Should Return []
+            return 0.
+        if content[0].split()[0].isalpha():
+            st = 1
+        return max([float(line.split()[st]) for line in content])
+    except FileNotFoundError:
+        print(f'No Corresponding Box Found for file {file}, using [] as preds')
+        return []
+    except Exception as e:
+        print('Some Error',e)
+        return []
+
+# Create the image dict
+def generate_image_dict(preds_folder_name='preds_42',
+                        root_fol='/home/krithika_1/densebreeast_datasets/AIIMS_C1',
+                        mal_path=None, ben_path=None, gt_path=None,
+                        mal_img_path = None, ben_img_path = None
+                        ):
+
+    mal_path = join(root_fol, mal_path) if mal_path else join(
+        root_fol, 'mal', preds_folder_name)
+    ben_path = join(root_fol, ben_path) if ben_path else join(
+        root_fol, 'ben', preds_folder_name)
+    mal_img_path = join(root_fol, mal_img_path) if mal_img_path else join(
+        root_fol, 'mal', 'images')
+    ben_img_path = join(root_fol, ben_img_path) if ben_img_path else join(
+        root_fol, 'ben', 'images')
+    gt_path = join(root_fol, gt_path) if gt_path else join(
+        root_fol, 'mal', 'gt')
+
+
+    '''
+        image_dict structure:
+            'image_name(without txt/png)' : {'gt' : [[...]], 'preds' : score}
+    '''
+    image_dict = dict()
+
+    # GT Might be sightly different from images, therefore we will index gts based on
+    # the images folder instead.
+    for file in os.listdir(mal_img_path):
+    # for file in glob.glob(join(gt_path, '*.txt')):
+        if not file.endswith('.png'):
+            continue
+        file = file[:-4] + '.txt'
+        file = join(gt_path, file)
+        key = os.path.split(file)[-1][:-4]
+        image_dict[key] = dict()
+        image_dict[key]['gt'] = 1.
+        image_dict[key]['preds'] = 0.
+
+    for file in glob.glob(join(mal_path, '*.txt')):
+        key = os.path.split(file)[-1][:-4]
+        assert key in image_dict
+        image_dict[key]['preds'] = file_to_score(file)
+
+    for file in os.listdir(ben_img_path):
+    # for file in glob.glob(join(ben_path, '*.txt')):
+        if not file.endswith('.png'):
+            continue
+
+        file = file[:-4] + '.txt'
+        file = join(ben_path, file)
+        key = os.path.split(file)[-1][:-4]
+        # if key == 'Calc-Test_P_00353_LEFT_CC' or key == 'Calc-Training_P_00600_LEFT_CC':
+        #     continue
+        if key in image_dict:
+            print(key)
+            print('SHIT')
+            continue
+        # assert key not in image_dict
+        image_dict[key] = dict()
+        image_dict[key]['preds'] = file_to_score(file)
+        image_dict[key]['gt'] = 0.
+    return image_dict
+
+def get_auc_score_from_imdict(image_dict):
+    keys = list(image_dict.keys())
+    y = [image_dict[k]['gt']for k in keys]    
+    preds = [image_dict[k]['preds']for k in keys]
+    return roc_auc_score(y, preds)
+
+def get_accuracy_from_imdict(image_dict, thresh = 0.3):
+    keys = list(image_dict.keys())
+    ys = [image_dict[k]['gt']for k in keys]    
+    preds = [image_dict[k]['preds']for k in keys]
+    acc = 0
+    for y,pred in zip(ys,preds):
+        if pred < thresh and y == 0.:
+            acc+=1
+        elif pred > thresh and y == 1.:
+            acc+=1
+    return acc/len(preds)
+
+
+def get_auc_score(preds_image_folder, root_fol, retAcc = False, acc_thresh = 0.3):
+    im_dict = generate_image_dict(preds_image_folder, root_fol = root_fol)
+    if retAcc:
+        return get_auc_score_from_imdict(im_dict), get_accuracy_from_imdict(im_dict, acc_thresh)
+    else:
+        return get_auc_score_from_imdict(im_dict)
+
+if __name__ == '__main__':
+    seed = '42' if len(sys.argv)== 1 else sys.argv[1]
+
+    root_fol = '../bilateral_new/MammoDatasets/AIIMS_highres_reliable/test'
+
+    auc_score = get_auc_score(f'preds_{seed}',root_fol)
+    print(f'ROC AUC Score: {auc_score}')
+    

DenseMammogram/dataloaders.py

@@ -0,0 +1,259 @@
+# Get the dataloaders
+# There are only two types of dataloaders, viz. VanillaFRCNN and BilaterialFRCNN
+
+import torch
+import cv2
+import torchvision.transforms as T
+import detection.transforms as transforms
+from torch.utils.data import Dataset,DataLoader
+import detection.utils as utils
+import os
+from tqdm import tqdm
+import pandas as pd
+from os.path import join
+# VanillaFRCNN DataLoaders
+
+class FRCNNDataset(Dataset):
+    def __init__(self,inputs,transform):
+        self.transform = transform
+        self.dataset_dicts = inputs
+
+    def __len__(self):
+        return len(self.dataset_dicts)
+
+
+    def __getitem__(self,index: int):
+        # Select the sample
+        record = self.dataset_dicts[index]
+        # Load input and target
+        img = cv2.imread(record['file_name'])
+
+        target = {k:torch.tensor(v) for k,v in record.items() if k != 'file_name'}
+        if self.transform is not None:
+            img = T.ToPILImage()(img)
+            img,target = self.transform(img,target)
+
+        return img,target
+
+def xml_to_dicts(paths):
+    dataset_dicts = []
+    i=1
+    for path in paths:
+        for image in tqdm(os.listdir(os.path.join(path,'mal/images/'))):
+            xmlfile = os.path.join(path,'mal/gt/',image[:-4]+'.txt')
+            if(not os.path.exists(xmlfile)):
+                continue
+            img = cv2.imread(os.path.join(path,'mal/images/',image))
+            record = {}
+            record['file_name'] = os.path.join(path , 'mal/images/',image)
+            record['image_id'] = i
+            i+=1
+            record['width'] = img.shape[1]
+            record['height'] = img.shape[0]
+            objs = []
+            boxes = []
+            labels = []
+            area = []
+            iscrowd = []
+            f = open(xmlfile,'r')
+            for line in f.readlines():
+                box = list(map(int,map(float,line.split()[1:])))
+                boxes.append(box)
+                labels.append(1)
+                area.append((box[2]-box[0])*(box[3]-box[1]))
+                iscrowd.append(False)
+            f.close()
+            record["boxes"] = boxes
+            record["labels"] = labels
+            record["area"] = area
+            record["iscrowd"] = iscrowd
+            if(len(boxes)>0):
+                dataset_dicts.append(record)
+        for image in tqdm(os.listdir(os.path.join(path,'ben/images/'))):
+            img = cv2.imread(os.path.join(path,'ben/images/',image))
+            record = {}
+            record['file_name'] = os.path.join(path, 'ben/images/',image)
+            record['image_id'] = i
+            i+=1
+            record['width'] = img.shape[1]
+            record['height'] = img.shape[0]
+            record['boxes'] = torch.tensor([[0,0,img.shape[1],img.shape[0]]])
+            record['labels'] = torch.tensor([0])
+            record['area'] = [img.shape[1]*img.shape[0]]
+            record["iscrowd"] = [False]
+            dataset_dicts.append(record)
+    return dataset_dicts
+
+
+
+def get_FRCNN_dataloaders(cfg, batch_size = 2, data_dir = '../bilateral_new',):
+    transform_test = transforms.Compose([transforms.ToTensor()])
+    transform_train = transforms.Compose([transforms.RandomHorizontalFlip(),transforms.ToTensor()])
+    train_paths = [join(data_dir,cfg['AIIMS_DATA'],cfg['AIIMS_TRAIN_SPLIT']),join(data_dir,cfg['DDSM_DATA'],cfg['DDSM_TRAIN_SPLIT']),]
+    val_aiims_path = [join(data_dir,cfg['AIIMS_DATA'],cfg['AIIMS_VAL_SPLIT'])]
+    train_data = FRCNNDataset(xml_to_dicts(train_paths),transform_train)
+    test_aiims = FRCNNDataset(xml_to_dicts(val_aiims_path),transform_test)
+
+    train_loader = DataLoader(train_data,batch_size=batch_size,shuffle=True,drop_last=True,num_workers=4,collate_fn = utils.collate_fn)
+    test_aiims_loader = DataLoader(test_aiims,batch_size=batch_size,shuffle=True,drop_last=True,num_workers=4,collate_fn = utils.collate_fn)
+    #test_ddsm_loader = DataLoader(test_ddsm,batch_size=2,shuffle=True,drop_last=True,num_workers=5,collate_fn = utils.collate_fn)
+
+    return train_loader, test_aiims_loader
+
+# BilaterialFRCNN DataLoaders
+
+def get_direction(dset,file_name):
+        # 1 if right else -1
+        if dset == 'aiims' or dset == 'ddsm':
+            file_name = file_name.lower()
+            r = file_name.find('right')
+            l = file_name.find('left')
+            if l == r and l == -1:
+                raise Exception(f'Unidentifiable Direction {file_name}')
+            if l!=-1 and r!=-1:
+                raise Exception(f'Unidentifiable Direction {file_name}')
+            return 1 if r!=-1 else -1
+        if dset == 'inbreast':
+            dir =file_name.split('_')[3]
+            if dir == 'R': return 1
+            if dir == 'L': return -1
+            raise Exception(f'Unidentifiable Direction {file_name}')
+        if dset == 'irch': 
+            r = file_name.find('_R ')
+            l = file_name.find('_L ')
+            if l == r and l == -1:
+                raise Exception(f'Unidentifiable Direction {file_name}')
+            if l!=-1 and r!=-1:
+                raise Exception(f'Unidentifiable Direction {file_name}')
+            return 1 if r!=-1 else -1
+
+
+class BilateralDataset(torch.utils.data.Dataset):
+
+    def __init__(self,inputs,transform,dset):
+        self.transform = transform
+        self.dataset_dicts = inputs
+        self.dset = dset
+
+    def __len__(self):
+        return len(self.dataset_dicts)
+    
+
+    def __getitem__(self,index: int):
+        # Select the sample
+        record = self.dataset_dicts[index]
+        # Load input and target
+        img1 = cv2.imread(record['file_name'])
+        img2 = cv2.imread(record['file_2'])
+
+        target = {k:torch.tensor(v) for k,v in record.items() if k != 'file_name' and k!='file_2'}
+        if self.transform is not None:
+            img1 = T.ToPILImage()(img1)
+            img2 = T.ToPILImage()(img2)
+            if(get_direction(self.dset,record['file_name'].split('/')[-1])==1):
+                img1,target = transforms.RandomHorizontalFlip(1.0)(img1,target)
+            else:
+                img2,_ = transforms.RandomHorizontalFlip(1.0)(img2)
+            img1,target = self.transform(img1,target)
+            img2,target = self.transform(img2,target)
+
+        images = [img1,img2]
+        return images,target
+
+
+def xml_to_dicts_bilateral(paths,cor_dicts):
+    dataset_dicts = []
+    i=1
+    for path,cor_dict in zip(paths,cor_dicts):
+        for image in tqdm(os.listdir(os.path.join(path,'mal/images/'))):
+            if(not os.path.join(path,'mal/images/',image) in cor_dict):
+                continue
+            if(not os.path.isfile(cor_dict[os.path.join(path,'mal/images/',image)])):
+                continue
+            xmlfile = os.path.join(path,'mal/gt/',image[:-4]+'.txt')
+            if(not os.path.exists(xmlfile)):
+                continue
+            img = cv2.imread(os.path.join(path,'mal/images/',image))
+            
+            record = {}
+            record['file_name'] = os.path.join(path , 'mal/images/',image)
+            record['file_2'] = cor_dict[os.path.join(path,'mal/images/',image)]
+            record['image_id'] = i
+            i+=1
+            record['width'] = img.shape[1]
+            record['height'] = img.shape[0]
+            objs = []
+            boxes = []
+            labels = []
+            area = []
+            iscrowd = []
+            f = open(xmlfile,'r')
+            for line in f.readlines():
+                box = list(map(int,map(float,line.split()[1:])))
+                boxes.append(box)
+                labels.append(1)
+                area.append((box[2]-box[0])*(box[3]-box[1]))
+                iscrowd.append(False)
+
+            f.close()
+            record["boxes"] = boxes
+            record["labels"] = labels
+            record["area"] = area
+            record["iscrowd"] = iscrowd
+            if(len(boxes)>0):
+                dataset_dicts.append(record)
+
+        for image in tqdm(os.listdir(os.path.join(path,'ben/images/'))):
+            if(not os.path.join(path,'ben/images/',image) in cor_dict):
+                continue
+            if(not os.path.isfile(cor_dict[os.path.join(path,'ben/images/',image)])):
+                continue
+            img = cv2.imread(os.path.join(path,'ben/images/',image))
+
+            record = {}
+            record['file_name'] = os.path.join(path , 'ben/images/',image)
+            record['file_2'] = cor_dict[os.path.join(path,'ben/images/',image)]
+            img2 = cv2.imread(cor_dict[os.path.join(path,'ben/images/',image)])
+            record['image_id'] = i
+            i+=1
+            record['width'] = img.shape[1]
+            record['height'] = img.shape[0]
+
+            record["boxes"] = torch.tensor([[0,0,min(img.shape[1],img2.shape[1]),min(img.shape[0],img2.shape[0])]])
+            record['labels'] = torch.tensor([0])
+            record['area'] = [ min(img.shape[1],img2.shape[1]) *min(img.shape[0],img2.shape[0])]
+            record["iscrowd"] = [False]
+            if(len(boxes)>0):
+               dataset_dicts.append(record)
+
+    return dataset_dicts
+
+
+
+def get_dict(data_dir, filename):
+    df = pd.read_csv(filename, header=None, sep=r'\s+', quotechar='"').to_numpy()
+    cor_dict = dict()
+    for a in df:
+      if(a[0]==a[1]):
+        continue
+      cor_dict[a[0]] = a[1]
+    # print(cor_dict)
+    cor_dict = {join(data_dir,k):join(data_dir,v) for k,v in cor_dict.items()}
+    return cor_dict
+
+def get_bilateral_dataloaders(cfg, batch_size = 1, data_dir = '../bilateral_new'):
+    transform_test = transforms.Compose([transforms.ToTensor()])
+    transform_train = transforms.Compose([transforms.RandomHorizontalFlip(),transforms.ToTensor()])
+    train_paths = [join(data_dir,cfg['AIIMS_DATA'],cfg['AIIMS_TRAIN_SPLIT']),join(data_dir,cfg['DDSM_DATA'],cfg['DDSM_TRAIN_SPLIT']),]
+    val_aiims_path = [join(data_dir,cfg['AIIMS_DATA'],cfg['AIIMS_VAL_SPLIT'])]
+    cor_lists_train = [get_dict(data_dir,join(data_dir,cfg['AIIMS_CORRS_LIST'])),get_dict(data_dir,join(data_dir,cfg['DDSM_CORRS_LIST']))]
+    cor_lists_val = [get_dict(data_dir,join(data_dir,cfg['AIIMS_CORRS_LIST']))]
+    cor_lists_train = [get_dict(data_dir,join(data_dir,cfg['AIIMS_CORRS_LIST']))]
+    train_data = BilateralDataset(xml_to_dicts_bilateral(train_paths,cor_lists_train),transform_test,'aiims')
+    val_aiims = BilateralDataset(xml_to_dicts_bilateral(val_aiims_path,cor_lists_val),transform_test,'aiims')
+
+    train_loader = DataLoader(train_data,batch_size=batch_size,shuffle=True,drop_last=True,num_workers=4,collate_fn = utils.collate_fn)
+    val_aiims_loader = DataLoader(val_aiims,batch_size=batch_size,shuffle=True,drop_last=True,num_workers=4,collate_fn = utils.collate_fn)
+    #test_ddsm_loader = DataLoader(test_ddsm,batch_size=2,shuffle=True,drop_last=True,num_workers=5,collate_fn = utils.collate_fn)
+
+    return train_loader, val_aiims_loader

DenseMammogram/detection/README.md

@@ -0,0 +1,81 @@
+# Object detection reference training scripts
+
+This folder contains reference training scripts for object detection.
+They serve as a log of how to train specific models, to provide baseline
+training and evaluation scripts to quickly bootstrap research.
+
+To execute the example commands below you must install the following:
+
+```
+cython
+pycocotools
+matplotlib
+```
+
+You must modify the following flags:
+
+`--data-path=/path/to/coco/dataset`
+
+`--nproc_per_node=<number_of_gpus_available>`
+
+Except otherwise noted, all models have been trained on 8x V100 GPUs. 
+
+### Faster R-CNN ResNet-50 FPN
+```
+torchrun --nproc_per_node=8 train.py\
+    --dataset coco --model fasterrcnn_resnet50_fpn --epochs 26\
+    --lr-steps 16 22 --aspect-ratio-group-factor 3
+```
+
+### Faster R-CNN MobileNetV3-Large FPN
+```
+torchrun --nproc_per_node=8 train.py\
+    --dataset coco --model fasterrcnn_mobilenet_v3_large_fpn --epochs 26\
+    --lr-steps 16 22 --aspect-ratio-group-factor 3
+```
+
+### Faster R-CNN MobileNetV3-Large 320 FPN
+```
+torchrun --nproc_per_node=8 train.py\
+    --dataset coco --model fasterrcnn_mobilenet_v3_large_320_fpn --epochs 26\
+    --lr-steps 16 22 --aspect-ratio-group-factor 3
+```
+
+### RetinaNet
+```
+torchrun --nproc_per_node=8 train.py\
+    --dataset coco --model retinanet_resnet50_fpn --epochs 26\
+    --lr-steps 16 22 --aspect-ratio-group-factor 3 --lr 0.01
+```
+
+### SSD300 VGG16
+```
+torchrun --nproc_per_node=8 train.py\
+    --dataset coco --model ssd300_vgg16 --epochs 120\
+    --lr-steps 80 110 --aspect-ratio-group-factor 3 --lr 0.002 --batch-size 4\
+    --weight-decay 0.0005 --data-augmentation ssd
+```
+
+### SSDlite320 MobileNetV3-Large
+```
+torchrun --nproc_per_node=8 train.py\
+    --dataset coco --model ssdlite320_mobilenet_v3_large --epochs 660\
+    --aspect-ratio-group-factor 3 --lr-scheduler cosineannealinglr --lr 0.15 --batch-size 24\
+    --weight-decay 0.00004 --data-augmentation ssdlite
+```
+
+
+### Mask R-CNN
+```
+torchrun --nproc_per_node=8 train.py\
+    --dataset coco --model maskrcnn_resnet50_fpn --epochs 26\
+    --lr-steps 16 22 --aspect-ratio-group-factor 3
+```
+
+
+### Keypoint R-CNN
+```
+torchrun --nproc_per_node=8 train.py\
+    --dataset coco_kp --model keypointrcnn_resnet50_fpn --epochs 46\
+    --lr-steps 36 43 --aspect-ratio-group-factor 3
+```

DenseMammogram/detection/coco_eval.py

@@ -0,0 +1,191 @@
+import copy
+import io
+from contextlib import redirect_stdout
+
+import numpy as np
+import pycocotools.mask as mask_util
+import torch
+import detection.utils as utils
+from pycocotools.coco import COCO
+from pycocotools.cocoeval import COCOeval
+
+
+class CocoEvaluator:
+    def __init__(self, coco_gt, iou_types):
+        assert isinstance(iou_types, (list, tuple))
+        coco_gt = copy.deepcopy(coco_gt)
+        self.coco_gt = coco_gt
+
+        self.iou_types = iou_types
+        self.coco_eval = {}
+        for iou_type in iou_types:
+            self.coco_eval[iou_type] = COCOeval(coco_gt, iouType=iou_type)
+
+        self.img_ids = []
+        self.eval_imgs = {k: [] for k in iou_types}
+
+    def update(self, predictions):
+        img_ids = list(np.unique(list(predictions.keys())))
+        self.img_ids.extend(img_ids)
+
+        for iou_type in self.iou_types:
+            results = self.prepare(predictions, iou_type)
+            with redirect_stdout(io.StringIO()):
+                coco_dt = COCO.loadRes(self.coco_gt, results) if results else COCO()
+            coco_eval = self.coco_eval[iou_type]
+
+            coco_eval.cocoDt = coco_dt
+            coco_eval.params.imgIds = list(img_ids)
+            img_ids, eval_imgs = evaluate(coco_eval)
+
+            self.eval_imgs[iou_type].append(eval_imgs)
+
+    def synchronize_between_processes(self):
+        for iou_type in self.iou_types:
+            self.eval_imgs[iou_type] = np.concatenate(self.eval_imgs[iou_type], 2)
+            create_common_coco_eval(self.coco_eval[iou_type], self.img_ids, self.eval_imgs[iou_type])
+
+    def accumulate(self):
+        for coco_eval in self.coco_eval.values():
+            coco_eval.accumulate()
+
+    def summarize(self):
+        for iou_type, coco_eval in self.coco_eval.items():
+            print(f"IoU metric: {iou_type}")
+            coco_eval.summarize()
+
+    def prepare(self, predictions, iou_type):
+        if iou_type == "bbox":
+            return self.prepare_for_coco_detection(predictions)
+        if iou_type == "segm":
+            return self.prepare_for_coco_segmentation(predictions)
+        if iou_type == "keypoints":
+            return self.prepare_for_coco_keypoint(predictions)
+        raise ValueError(f"Unknown iou type {iou_type}")
+
+    def prepare_for_coco_detection(self, predictions):
+        coco_results = []
+        for original_id, prediction in predictions.items():
+            if len(prediction) == 0:
+                continue
+
+            boxes = prediction["boxes"]
+            boxes = convert_to_xywh(boxes).tolist()
+            scores = prediction["scores"].tolist()
+            labels = prediction["labels"].tolist()
+
+            coco_results.extend(
+                [
+                    {
+                        "image_id": original_id,
+                        "category_id": labels[k],
+                        "bbox": box,
+                        "score": scores[k],
+                    }
+                    for k, box in enumerate(boxes)
+                ]
+            )
+        return coco_results
+
+    def prepare_for_coco_segmentation(self, predictions):
+        coco_results = []
+        for original_id, prediction in predictions.items():
+            if len(prediction) == 0:
+                continue
+
+            scores = prediction["scores"]
+            labels = prediction["labels"]
+            masks = prediction["masks"]
+
+            masks = masks > 0.5
+
+            scores = prediction["scores"].tolist()
+            labels = prediction["labels"].tolist()
+
+            rles = [
+                mask_util.encode(np.array(mask[0, :, :, np.newaxis], dtype=np.uint8, order="F"))[0] for mask in masks
+            ]
+            for rle in rles:
+                rle["counts"] = rle["counts"].decode("utf-8")
+
+            coco_results.extend(
+                [
+                    {
+                        "image_id": original_id,
+                        "category_id": labels[k],
+                        "segmentation": rle,
+                        "score": scores[k],
+                    }
+                    for k, rle in enumerate(rles)
+                ]
+            )
+        return coco_results
+
+    def prepare_for_coco_keypoint(self, predictions):
+        coco_results = []
+        for original_id, prediction in predictions.items():
+            if len(prediction) == 0:
+                continue
+
+            boxes = prediction["boxes"]
+            boxes = convert_to_xywh(boxes).tolist()
+            scores = prediction["scores"].tolist()
+            labels = prediction["labels"].tolist()
+            keypoints = prediction["keypoints"]
+            keypoints = keypoints.flatten(start_dim=1).tolist()
+
+            coco_results.extend(
+                [
+                    {
+                        "image_id": original_id,
+                        "category_id": labels[k],
+                        "keypoints": keypoint,
+                        "score": scores[k],
+                    }
+                    for k, keypoint in enumerate(keypoints)
+                ]
+            )
+        return coco_results
+
+
+def convert_to_xywh(boxes):
+    xmin, ymin, xmax, ymax = boxes.unbind(1)
+    return torch.stack((xmin, ymin, xmax - xmin, ymax - ymin), dim=1)
+
+
+def merge(img_ids, eval_imgs):
+    all_img_ids = utils.all_gather(img_ids)
+    all_eval_imgs = utils.all_gather(eval_imgs)
+
+    merged_img_ids = []
+    for p in all_img_ids:
+        merged_img_ids.extend(p)
+
+    merged_eval_imgs = []
+    for p in all_eval_imgs:
+        merged_eval_imgs.append(p)
+
+    merged_img_ids = np.array(merged_img_ids)
+    merged_eval_imgs = np.concatenate(merged_eval_imgs, 2)
+
+    # keep only unique (and in sorted order) images
+    merged_img_ids, idx = np.unique(merged_img_ids, return_index=True)
+    merged_eval_imgs = merged_eval_imgs[..., idx]
+
+    return merged_img_ids, merged_eval_imgs
+
+
+def create_common_coco_eval(coco_eval, img_ids, eval_imgs):
+    img_ids, eval_imgs = merge(img_ids, eval_imgs)
+    img_ids = list(img_ids)
+    eval_imgs = list(eval_imgs.flatten())
+
+    coco_eval.evalImgs = eval_imgs
+    coco_eval.params.imgIds = img_ids
+    coco_eval._paramsEval = copy.deepcopy(coco_eval.params)
+
+
+def evaluate(imgs):
+    with redirect_stdout(io.StringIO()):
+        imgs.evaluate()
+    return imgs.params.imgIds, np.asarray(imgs.evalImgs).reshape(-1, len(imgs.params.areaRng), len(imgs.params.imgIds))

DenseMammogram/detection/coco_utils.py

@@ -0,0 +1,249 @@
+import copy
+import os
+
+import torch
+import torch.utils.data
+import torchvision
+import detection.transforms as T
+from pycocotools import mask as coco_mask
+from pycocotools.coco import COCO
+
+
+class FilterAndRemapCocoCategories:
+    def __init__(self, categories, remap=True):
+        self.categories = categories
+        self.remap = remap
+
+    def __call__(self, image, target):
+        anno = target["annotations"]
+        anno = [obj for obj in anno if obj["category_id"] in self.categories]
+        if not self.remap:
+            target["annotations"] = anno
+            return image, target
+        anno = copy.deepcopy(anno)
+        for obj in anno:
+            obj["category_id"] = self.categories.index(obj["category_id"])
+        target["annotations"] = anno
+        return image, target
+
+
+def convert_coco_poly_to_mask(segmentations, height, width):
+    masks = []
+    for polygons in segmentations:
+        rles = coco_mask.frPyObjects(polygons, height, width)
+        mask = coco_mask.decode(rles)
+        if len(mask.shape) < 3:
+            mask = mask[..., None]
+        mask = torch.as_tensor(mask, dtype=torch.uint8)
+        mask = mask.any(dim=2)
+        masks.append(mask)
+    if masks:
+        masks = torch.stack(masks, dim=0)
+    else:
+        masks = torch.zeros((0, height, width), dtype=torch.uint8)
+    return masks
+
+
+class ConvertCocoPolysToMask:
+    def __call__(self, image, target):
+        w, h = image.size
+
+        image_id = target["image_id"]
+        image_id = torch.tensor([image_id])
+
+        anno = target["annotations"]
+
+        anno = [obj for obj in anno if obj["iscrowd"] == 0]
+
+        boxes = [obj["bbox"] for obj in anno]
+        # guard against no boxes via resizing
+        boxes = torch.as_tensor(boxes, dtype=torch.float32).reshape(-1, 4)
+        boxes[:, 2:] += boxes[:, :2]
+        boxes[:, 0::2].clamp_(min=0, max=w)
+        boxes[:, 1::2].clamp_(min=0, max=h)
+
+        classes = [obj["category_id"] for obj in anno]
+        classes = torch.tensor(classes, dtype=torch.int64)
+
+        segmentations = [obj["segmentation"] for obj in anno]
+        masks = convert_coco_poly_to_mask(segmentations, h, w)
+
+        keypoints = None
+        if anno and "keypoints" in anno[0]:
+            keypoints = [obj["keypoints"] for obj in anno]
+            keypoints = torch.as_tensor(keypoints, dtype=torch.float32)
+            num_keypoints = keypoints.shape[0]
+            if num_keypoints:
+                keypoints = keypoints.view(num_keypoints, -1, 3)
+
+        keep = (boxes[:, 3] > boxes[:, 1]) & (boxes[:, 2] > boxes[:, 0])
+        boxes = boxes[keep]
+        classes = classes[keep]
+        masks = masks[keep]
+        if keypoints is not None:
+            keypoints = keypoints[keep]
+
+        target = {}
+        target["boxes"] = boxes
+        target["labels"] = classes
+        target["masks"] = masks
+        target["image_id"] = image_id
+        if keypoints is not None:
+            target["keypoints"] = keypoints
+
+        # for conversion to coco api
+        area = torch.tensor([obj["area"] for obj in anno])
+        iscrowd = torch.tensor([obj["iscrowd"] for obj in anno])
+        target["area"] = area
+        target["iscrowd"] = iscrowd
+
+        return image, target
+
+
+def _coco_remove_images_without_annotations(dataset, cat_list=None):
+    def _has_only_empty_bbox(anno):
+        return all(any(o <= 1 for o in obj["bbox"][2:]) for obj in anno)
+
+    def _count_visible_keypoints(anno):
+        return sum(sum(1 for v in ann["keypoints"][2::3] if v > 0) for ann in anno)
+
+    min_keypoints_per_image = 10
+
+    def _has_valid_annotation(anno):
+        # if it's empty, there is no annotation
+        if len(anno) == 0:
+            return False
+        # if all boxes have close to zero area, there is no annotation
+        if _has_only_empty_bbox(anno):
+            return False
+        # keypoints task have a slight different critera for considering
+        # if an annotation is valid
+        if "keypoints" not in anno[0]:
+            return True
+        # for keypoint detection tasks, only consider valid images those
+        # containing at least min_keypoints_per_image
+        if _count_visible_keypoints(anno) >= min_keypoints_per_image:
+            return True
+        return False
+
+    assert isinstance(dataset, torchvision.datasets.CocoDetection)
+    ids = []
+    for ds_idx, img_id in enumerate(dataset.ids):
+        ann_ids = dataset.coco.getAnnIds(imgIds=img_id, iscrowd=None)
+        anno = dataset.coco.loadAnns(ann_ids)
+        if cat_list:
+            anno = [obj for obj in anno if obj["category_id"] in cat_list]
+        if _has_valid_annotation(anno):
+            ids.append(ds_idx)
+
+    dataset = torch.utils.data.Subset(dataset, ids)
+    return dataset
+
+
+def convert_to_coco_api(ds):
+    coco_ds = COCO()
+    # annotation IDs need to start at 1, not 0, see torchvision issue #1530
+    ann_id = 1
+    dataset = {"images": [], "categories": [], "annotations": []}
+    categories = set()
+    for img_idx in range(len(ds)):
+        # find better way to get target
+        # targets = ds.get_annotations(img_idx)
+        img, targets = ds[img_idx]
+        image_id = targets["image_id"].item()
+        img_dict = {}
+        img_dict["id"] = image_id
+        img_dict["height"] = img.shape[-2]
+        img_dict["width"] = img.shape[-1]
+        dataset["images"].append(img_dict)
+        bboxes = targets["boxes"]
+        bboxes[:, 2:] -= bboxes[:, :2]
+        bboxes = bboxes.tolist()
+        labels = targets["labels"].tolist()
+        areas = targets["area"].tolist()
+        iscrowd = targets["iscrowd"].tolist()
+        if "masks" in targets:
+            masks = targets["masks"]
+            # make masks Fortran contiguous for coco_mask
+            masks = masks.permute(0, 2, 1).contiguous().permute(0, 2, 1)
+        if "keypoints" in targets:
+            keypoints = targets["keypoints"]
+            keypoints = keypoints.reshape(keypoints.shape[0], -1).tolist()
+        num_objs = len(bboxes)
+        for i in range(num_objs):
+            ann = {}
+            ann["image_id"] = image_id
+            ann["bbox"] = bboxes[i]
+            ann["category_id"] = labels[i]
+            categories.add(labels[i])
+            ann["area"] = areas[i]
+            ann["iscrowd"] = iscrowd[i]
+            ann["id"] = ann_id
+            if "masks" in targets:
+                ann["segmentation"] = coco_mask.encode(masks[i].numpy())
+            if "keypoints" in targets:
+                ann["keypoints"] = keypoints[i]
+                ann["num_keypoints"] = sum(k != 0 for k in keypoints[i][2::3])
+            dataset["annotations"].append(ann)
+            ann_id += 1
+    dataset["categories"] = [{"id": i} for i in sorted(categories)]
+    coco_ds.dataset = dataset
+    coco_ds.createIndex()
+    return coco_ds
+
+
+def get_coco_api_from_dataset(dataset):
+    for _ in range(10):
+        if isinstance(dataset, torchvision.datasets.CocoDetection):
+            break
+        if isinstance(dataset, torch.utils.data.Subset):
+            dataset = dataset.dataset
+    if isinstance(dataset, torchvision.datasets.CocoDetection):
+        return dataset.coco
+    return convert_to_coco_api(dataset)
+
+
+class CocoDetection(torchvision.datasets.CocoDetection):
+    def __init__(self, img_folder, ann_file, transforms):
+        super().__init__(img_folder, ann_file)
+        self._transforms = transforms
+
+    def __getitem__(self, idx):
+        img, target = super().__getitem__(idx)
+        image_id = self.ids[idx]
+        target = dict(image_id=image_id, annotations=target)
+        if self._transforms is not None:
+            img, target = self._transforms(img, target)
+        return img, target
+
+
+def get_coco(root, image_set, transforms, mode="instances"):
+    anno_file_template = "{}_{}2017.json"
+    PATHS = {
+        "train": ("train2017", os.path.join("annotations", anno_file_template.format(mode, "train"))),
+        "val": ("val2017", os.path.join("annotations", anno_file_template.format(mode, "val"))),
+        # "train": ("val2017", os.path.join("annotations", anno_file_template.format(mode, "val")))
+    }
+
+    t = [ConvertCocoPolysToMask()]
+
+    if transforms is not None:
+        t.append(transforms)
+    transforms = T.Compose(t)
+
+    img_folder, ann_file = PATHS[image_set]
+    img_folder = os.path.join(root, img_folder)
+    ann_file = os.path.join(root, ann_file)
+
+    dataset = CocoDetection(img_folder, ann_file, transforms=transforms)
+
+    if image_set == "train":
+        dataset = _coco_remove_images_without_annotations(dataset)
+
+    # dataset = torch.utils.data.Subset(dataset, [i for i in range(500)])
+
+    return dataset
+
+
+def get_coco_kp(root, image_set, transforms):
+    return get_coco(root, image_set, transforms, mode="person_keypoints")

DenseMammogram/detection/engine.py

@@ -0,0 +1,276 @@
+import math
+import sys
+import time
+
+import torch
+import torchvision.models.detection.mask_rcnn
+import detection.utils as utils
+from detection.coco_eval import CocoEvaluator
+from detection.coco_utils import get_coco_api_from_dataset
+from tqdm import tqdm
+import numpy as np
+
+
+sys.path.append("..")
+from utils import AverageMeter
+from advanced_logger import LogPriority
+
+def train_one_epoch(model, optimizer, data_loader, device, epoch, print_freq, scaler=None):
+    model.train()
+    metric_logger = utils.MetricLogger(delimiter="  ")
+    metric_logger.add_meter("lr", utils.SmoothedValue(window_size=1, fmt="{value:.6f}"))
+    header = f"Epoch: [{epoch}]"
+
+    lr_scheduler = None
+    if epoch == 0:
+        warmup_factor = 1.0 / 1000
+        warmup_iters = min(1000, len(data_loader) - 1)
+
+        lr_scheduler = torch.optim.lr_scheduler.LinearLR(
+            optimizer, start_factor=warmup_factor, total_iters=warmup_iters
+        )
+    #for batch_idx,(images, targets) in enumerate(tqdm(data_loader)):
+    for images, targets in metric_logger.log_every(data_loader, print_freq, header):
+        #print(images.shape)
+        images = list(image.to(device) if len(image)>2 else [image[0].to(device),image[1].to(device)] for image in images)
+        #print(len(images))
+        #print(images[0].shape)
+        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
+        with torch.cuda.amp.autocast(enabled=scaler is not None):
+            loss_dict = model(images, targets)
+            losses = sum(loss for loss in loss_dict.values())
+
+        # reduce losses over all GPUs for logging purposes
+        loss_dict_reduced = utils.reduce_dict(loss_dict)
+        losses_reduced = sum(loss for loss in loss_dict_reduced.values())
+
+        loss_value = losses_reduced.item()
+
+        if not math.isfinite(loss_value):
+            print(f"Loss is {loss_value}, stopping training")
+            print(loss_dict_reduced)
+            sys.exit(1)
+
+        optimizer.zero_grad()
+        if scaler is not None:
+            scaler.scale(losses).backward()
+            scaler.step(optimizer)
+            scaler.update()
+        else:
+            losses.backward()
+            optimizer.step()
+
+        if lr_scheduler is not None:
+            lr_scheduler.step()
+        
+        #if(batch_idx%20==0):
+        #    print('epoch {} batch {} : {}'.format(epoch,batch_idx,losses_reduced))
+
+        metric_logger.update(loss=losses_reduced, **loss_dict_reduced)
+        metric_logger.update(lr=optimizer.param_groups[0]["lr"])
+
+    return metric_logger
+
+
+def train_one_epoch_simplified(model, optimizer, data_loader, device, epoch, experimenter,optimizer_backbone=None):
+    
+    model.train()
+    lr_scheduler = None
+    lr_scheduler_backbone = None
+    if epoch == 0:
+        warmup_factor = 1.0 / 1000
+        warmup_iters = min(1000, len(data_loader) - 1)
+
+        lr_scheduler = torch.optim.lr_scheduler.LinearLR(
+            optimizer, start_factor=warmup_factor, total_iters=warmup_iters
+        )
+        if(optimizer_backbone is not None):
+           lr_scheduler_backbone = torch.optim.lr_scheduler.LinearLR(optimizer_backbone, start_factor=warmup_factor, total_iters=warmup_iters)
+           
+
+    loss_meter = AverageMeter()
+
+    for step, (images, targets) in enumerate(tqdm(data_loader)):
+
+        optimizer.zero_grad()
+        if(optimizer_backbone is not None):
+            optimizer_backbone.zero_grad()
+
+        images = list(image.to(device) if len(image)>2 else [image[0].to(device),image[1].to(device)] for image in images)
+        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
+        loss_dict = model(images, targets)
+        losses = sum(loss for loss in loss_dict.values())
+
+
+        if not math.isfinite(losses.item()):
+            print(f"Loss is {losses.item()}, stopping training")
+            print(loss_dict)
+            experimenter.log(f"Loss is {losses.item()}, stopping training")
+            sys.exit(1)
+
+        losses.backward()
+        loss_meter.update(losses.item())
+        optimizer.step()
+        if optimizer_backbone is not None:
+            optimizer_backbone.step()
+        if lr_scheduler is not None:
+            lr_scheduler.step()
+        if lr_scheduler_backbone is not None:
+            lr_scheduler_backbone.step()
+            
+        if (step+1)%10 == 0:
+            experimenter.log('Loss after {} steps: {}'.format(step+1, loss_meter.avg))
+        if epoch == 0 and (step+1)%50 == 0:
+            experimenter.log('LR after {} steps: {}'.format(step+1, optimizer.param_groups[0]['lr']))
+
+def _get_iou_types(model):
+    model_without_ddp = model
+    if isinstance(model, torch.nn.parallel.DistributedDataParallel):
+        model_without_ddp = model.module
+    iou_types = ["bbox"]
+    if isinstance(model_without_ddp, torchvision.models.detection.MaskRCNN):
+        iou_types.append("segm")
+    if isinstance(model_without_ddp, torchvision.models.detection.KeypointRCNN):
+        iou_types.append("keypoints")
+    return iou_types
+
+
+@torch.inference_mode()
+def evaluate(model, data_loader, device):
+    n_threads = torch.get_num_threads()
+    # FIXME remove this and make paste_masks_in_image run on the GPU
+    torch.set_num_threads(1)
+    cpu_device = torch.device("cpu")
+    model.eval()
+    metric_logger = utils.MetricLogger(delimiter="  ")
+    header = "Test:"
+
+    coco = get_coco_api_from_dataset(data_loader.dataset)
+    iou_types = _get_iou_types(model)
+    coco_evaluator = CocoEvaluator(coco, iou_types)
+
+    for images, targets in metric_logger.log_every(data_loader, 100, header):
+        images = list(img.to(device) for img in images)
+
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+        model_time = time.time()
+        outputs = model(images)
+
+        outputs = [{k: v.to(cpu_device) for k, v in t.items()} for t in outputs]
+        model_time = time.time() - model_time
+
+        res = {target["image_id"].item(): output for target, output in zip(targets, outputs)}
+        evaluator_time = time.time()
+        coco_evaluator.update(res)
+        evaluator_time = time.time() - evaluator_time
+        metric_logger.update(model_time=model_time, evaluator_time=evaluator_time)
+
+    # gather the stats from all processes
+    metric_logger.synchronize_between_processes()
+    print("Averaged stats:", metric_logger)
+    coco_evaluator.synchronize_between_processes()
+
+    # accumulate predictions from all images
+    coco_evaluator.accumulate()
+    coco_evaluator.summarize()
+    torch.set_num_threads(n_threads)
+    return coco_evaluator
+
+
+def coco_summ(coco_eval, experimenter):
+    self = coco_eval
+    def _summarize( ap=1, iouThr=None, areaRng='all', maxDets=100 ):
+        p = self.params
+        iStr = ' {:<18} {} @[ IoU={:<9} | area={:>6s} | maxDets={:>3d} ] = {:0.3f}'
+        titleStr = 'Average Precision' if ap == 1 else 'Average Recall'
+        typeStr = '(AP)' if ap==1 else '(AR)'
+        iouStr = '{:0.2f}:{:0.2f}'.format(p.iouThrs[0], p.iouThrs[-1]) \
+            if iouThr is None else '{:0.2f}'.format(iouThr)
+
+        aind = [i for i, aRng in enumerate(p.areaRngLbl) if aRng == areaRng]
+        mind = [i for i, mDet in enumerate(p.maxDets) if mDet == maxDets]
+        if ap == 1:
+            # dimension of precision: [TxRxKxAxM]
+            s = self.eval['precision']
+            # IoU
+            if iouThr is not None:
+                t = np.where(iouThr == p.iouThrs)[0]
+                s = s[t]
+            s = s[:,:,:,aind,mind]
+        else:
+            # dimension of recall: [TxKxAxM]
+            s = self.eval['recall']
+            if iouThr is not None:
+                t = np.where(iouThr == p.iouThrs)[0]
+                s = s[t]
+            s = s[:,:,aind,mind]
+        if len(s[s>-1])==0:
+            mean_s = -1
+        else:
+            mean_s = np.mean(s[s>-1])
+        experimenter.log(iStr.format(titleStr, typeStr, iouStr, areaRng, maxDets, mean_s), priority = LogPriority.MEDIUM)
+        return mean_s
+    def _summarizeDets():
+        stats = np.zeros((12,))
+        stats[0] = _summarize(1)
+        stats[1] = _summarize(1, iouThr=.5, maxDets=self.params.maxDets[2])
+        stats[2] = _summarize(1, iouThr=.75, maxDets=self.params.maxDets[2])
+        stats[3] = _summarize(1, areaRng='small', maxDets=self.params.maxDets[2])
+        stats[4] = _summarize(1, areaRng='medium', maxDets=self.params.maxDets[2])
+        stats[5] = _summarize(1, areaRng='large', maxDets=self.params.maxDets[2])
+        stats[6] = _summarize(0, maxDets=self.params.maxDets[0])
+        stats[7] = _summarize(0, maxDets=self.params.maxDets[1])
+        stats[8] = _summarize(0, maxDets=self.params.maxDets[2])
+        stats[9] = _summarize(0, areaRng='small', maxDets=self.params.maxDets[2])
+        stats[10] = _summarize(0, areaRng='medium', maxDets=self.params.maxDets[2])
+        stats[11] = _summarize(0, areaRng='large', maxDets=self.params.maxDets[2])
+        return stats
+    _summarizeDets()
+
+@torch.inference_mode()
+def evaluate_simplified(model, data_loader, device, experimenter):
+    cpu_device = torch.device("cpu")
+    model.eval()
+    experimenter.log('Evaluating Validation Parameters')
+
+    coco = get_coco_api_from_dataset(data_loader.dataset)
+    iou_types = _get_iou_types(model)
+    coco_evaluator = CocoEvaluator(coco, iou_types)
+
+    for images, targets in data_loader:
+        images = list(img.to(device) for img in images)
+
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+        outputs = model(images)
+        outputs = [{k: v.to(cpu_device) for k, v in t.items()} for t in outputs]
+        res = {target["image_id"].item(): output for target, output in zip(targets, outputs)}
+        coco_evaluator.update(res)
+        
+    # gather the stats from all processes
+    coco_evaluator.synchronize_between_processes()
+
+    # accumulate predictions from all images
+    coco_evaluator.accumulate()
+
+    # Debug and see what all info it has
+    # coco_evaluator.summarize()
+    for iou_type, coco_eval in coco_evaluator.coco_eval.items():
+        print(f"IoU metric: {iou_type}")
+        coco_summ(coco_eval, experimenter)
+
+    return coco_evaluator
+
+def evaluate_loss(model, device, val_loader, experimenter=None):
+    model.train()
+    #experimenter.log('Evaluating Validation Loss')
+    with torch.no_grad():
+         loss_meter = AverageMeter()
+    for images, targets in tqdm(val_loader):
+        images = list(image.to(device) if len(image)>2 else [image[0].to(device),image[1].to(device)] for image in images)
+        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
+        loss_dict = model(images, targets)
+        losses = sum(loss for loss in loss_dict.values())
+        loss_meter.update(losses.item())
+    return loss_meter.avg

DenseMammogram/detection/group_by_aspect_ratio.py

@@ -0,0 +1,196 @@
+import bisect
+import copy
+import math
+from collections import defaultdict
+from itertools import repeat, chain
+
+import numpy as np
+import torch
+import torch.utils.data
+import torchvision
+from PIL import Image
+from torch.utils.data.sampler import BatchSampler, Sampler
+from torch.utils.model_zoo import tqdm
+
+
+def _repeat_to_at_least(iterable, n):
+    repeat_times = math.ceil(n / len(iterable))
+    repeated = chain.from_iterable(repeat(iterable, repeat_times))
+    return list(repeated)
+
+
+class GroupedBatchSampler(BatchSampler):
+    """
+    Wraps another sampler to yield a mini-batch of indices.
+    It enforces that the batch only contain elements from the same group.
+    It also tries to provide mini-batches which follows an ordering which is
+    as close as possible to the ordering from the original sampler.
+    Args:
+        sampler (Sampler): Base sampler.
+        group_ids (list[int]): If the sampler produces indices in range [0, N),
+            `group_ids` must be a list of `N` ints which contains the group id of each sample.
+            The group ids must be a continuous set of integers starting from
+            0, i.e. they must be in the range [0, num_groups).
+        batch_size (int): Size of mini-batch.
+    """
+
+    def __init__(self, sampler, group_ids, batch_size):
+        if not isinstance(sampler, Sampler):
+            raise ValueError(f"sampler should be an instance of torch.utils.data.Sampler, but got sampler={sampler}")
+        self.sampler = sampler
+        self.group_ids = group_ids
+        self.batch_size = batch_size
+
+    def __iter__(self):
+        buffer_per_group = defaultdict(list)
+        samples_per_group = defaultdict(list)
+
+        num_batches = 0
+        for idx in self.sampler:
+            group_id = self.group_ids[idx]
+            buffer_per_group[group_id].append(idx)
+            samples_per_group[group_id].append(idx)
+            if len(buffer_per_group[group_id]) == self.batch_size:
+                yield buffer_per_group[group_id]
+                num_batches += 1
+                del buffer_per_group[group_id]
+            assert len(buffer_per_group[group_id]) < self.batch_size
+
+        # now we have run out of elements that satisfy
+        # the group criteria, let's return the remaining
+        # elements so that the size of the sampler is
+        # deterministic
+        expected_num_batches = len(self)
+        num_remaining = expected_num_batches - num_batches
+        if num_remaining > 0:
+            # for the remaining batches, take first the buffers with largest number
+            # of elements
+            for group_id, _ in sorted(buffer_per_group.items(), key=lambda x: len(x[1]), reverse=True):
+                remaining = self.batch_size - len(buffer_per_group[group_id])
+                samples_from_group_id = _repeat_to_at_least(samples_per_group[group_id], remaining)
+                buffer_per_group[group_id].extend(samples_from_group_id[:remaining])
+                assert len(buffer_per_group[group_id]) == self.batch_size
+                yield buffer_per_group[group_id]
+                num_remaining -= 1
+                if num_remaining == 0:
+                    break
+        assert num_remaining == 0
+
+    def __len__(self):
+        return len(self.sampler) // self.batch_size
+
+
+def _compute_aspect_ratios_slow(dataset, indices=None):
+    print(
+        "Your dataset doesn't support the fast path for "
+        "computing the aspect ratios, so will iterate over "
+        "the full dataset and load every image instead. "
+        "This might take some time..."
+    )
+    if indices is None:
+        indices = range(len(dataset))
+
+    class SubsetSampler(Sampler):
+        def __init__(self, indices):
+            self.indices = indices
+
+        def __iter__(self):
+            return iter(self.indices)
+
+        def __len__(self):
+            return len(self.indices)
+
+    sampler = SubsetSampler(indices)
+    data_loader = torch.utils.data.DataLoader(
+        dataset,
+        batch_size=1,
+        sampler=sampler,
+        num_workers=14,  # you might want to increase it for faster processing
+        collate_fn=lambda x: x[0],
+    )
+    aspect_ratios = []
+    with tqdm(total=len(dataset)) as pbar:
+        for _i, (img, _) in enumerate(data_loader):
+            pbar.update(1)
+            height, width = img.shape[-2:]
+            aspect_ratio = float(width) / float(height)
+            aspect_ratios.append(aspect_ratio)
+    return aspect_ratios
+
+
+def _compute_aspect_ratios_custom_dataset(dataset, indices=None):
+    if indices is None:
+        indices = range(len(dataset))
+    aspect_ratios = []
+    for i in indices:
+        height, width = dataset.get_height_and_width(i)
+        aspect_ratio = float(width) / float(height)
+        aspect_ratios.append(aspect_ratio)
+    return aspect_ratios
+
+
+def _compute_aspect_ratios_coco_dataset(dataset, indices=None):
+    if indices is None:
+        indices = range(len(dataset))
+    aspect_ratios = []
+    for i in indices:
+        img_info = dataset.coco.imgs[dataset.ids[i]]
+        aspect_ratio = float(img_info["width"]) / float(img_info["height"])
+        aspect_ratios.append(aspect_ratio)
+    return aspect_ratios
+
+
+def _compute_aspect_ratios_voc_dataset(dataset, indices=None):
+    if indices is None:
+        indices = range(len(dataset))
+    aspect_ratios = []
+    for i in indices:
+        # this doesn't load the data into memory, because PIL loads it lazily
+        width, height = Image.open(dataset.images[i]).size
+        aspect_ratio = float(width) / float(height)
+        aspect_ratios.append(aspect_ratio)
+    return aspect_ratios
+
+
+def _compute_aspect_ratios_subset_dataset(dataset, indices=None):
+    if indices is None:
+        indices = range(len(dataset))
+
+    ds_indices = [dataset.indices[i] for i in indices]
+    return compute_aspect_ratios(dataset.dataset, ds_indices)
+
+
+def compute_aspect_ratios(dataset, indices=None):
+    if hasattr(dataset, "get_height_and_width"):
+        return _compute_aspect_ratios_custom_dataset(dataset, indices)
+
+    if isinstance(dataset, torchvision.datasets.CocoDetection):
+        return _compute_aspect_ratios_coco_dataset(dataset, indices)
+
+    if isinstance(dataset, torchvision.datasets.VOCDetection):
+        return _compute_aspect_ratios_voc_dataset(dataset, indices)
+
+    if isinstance(dataset, torch.utils.data.Subset):
+        return _compute_aspect_ratios_subset_dataset(dataset, indices)
+
+    # slow path
+    return _compute_aspect_ratios_slow(dataset, indices)
+
+
+def _quantize(x, bins):
+    bins = copy.deepcopy(bins)
+    bins = sorted(bins)
+    quantized = list(map(lambda y: bisect.bisect_right(bins, y), x))
+    return quantized
+
+
+def create_aspect_ratio_groups(dataset, k=0):
+    aspect_ratios = compute_aspect_ratios(dataset)
+    bins = (2 ** np.linspace(-1, 1, 2 * k + 1)).tolist() if k > 0 else [1.0]
+    groups = _quantize(aspect_ratios, bins)
+    # count number of elements per group
+    counts = np.unique(groups, return_counts=True)[1]
+    fbins = [0] + bins + [np.inf]
+    print(f"Using {fbins} as bins for aspect ratio quantization")
+    print(f"Count of instances per bin: {counts}")
+    return groups

DenseMammogram/detection/presets.py

@@ -0,0 +1,47 @@
+import torch
+import detection.transforms as T
+
+
+class DetectionPresetTrain:
+    def __init__(self, data_augmentation, hflip_prob=0.5, mean=(123.0, 117.0, 104.0)):
+        if data_augmentation == "hflip":
+            self.transforms = T.Compose(
+                [
+                    T.RandomHorizontalFlip(p=hflip_prob),
+                    T.PILToTensor(),
+                    T.ConvertImageDtype(torch.float),
+                ]
+            )
+        elif data_augmentation == "ssd":
+            self.transforms = T.Compose(
+                [
+                    T.RandomPhotometricDistort(),
+                    T.RandomZoomOut(fill=list(mean)),
+                    T.RandomIoUCrop(),
+                    T.RandomHorizontalFlip(p=hflip_prob),
+                    T.PILToTensor(),
+                    T.ConvertImageDtype(torch.float),
+                ]
+            )
+        elif data_augmentation == "ssdlite":
+            self.transforms = T.Compose(
+                [
+                    T.RandomIoUCrop(),
+                    T.RandomHorizontalFlip(p=hflip_prob),
+                    T.PILToTensor(),
+                    T.ConvertImageDtype(torch.float),
+                ]
+            )
+        else:
+            raise ValueError(f'Unknown data augmentation policy "{data_augmentation}"')
+
+    def __call__(self, img, target):
+        return self.transforms(img, target)
+
+
+class DetectionPresetEval:
+    def __init__(self):
+        self.transforms = T.ToTensor()
+
+    def __call__(self, img, target):
+        return self.transforms(img, target)

DenseMammogram/detection/train.py

@@ -0,0 +1,269 @@
+r"""PyTorch Detection Training.
+
+To run in a multi-gpu environment, use the distributed launcher::
+
+    python -m torch.distributed.launch --nproc_per_node=$NGPU --use_env \
+        train.py ... --world-size $NGPU
+
+The default hyperparameters are tuned for training on 8 gpus and 2 images per gpu.
+    --lr 0.02 --batch-size 2 --world-size 8
+If you use different number of gpus, the learning rate should be changed to 0.02/8*$NGPU.
+
+On top of that, for training Faster/Mask R-CNN, the default hyperparameters are
+    --epochs 26 --lr-steps 16 22 --aspect-ratio-group-factor 3
+
+Also, if you train Keypoint R-CNN, the default hyperparameters are
+    --epochs 46 --lr-steps 36 43 --aspect-ratio-group-factor 3
+Because the number of images is smaller in the person keypoint subset of COCO,
+the number of epochs should be adapted so that we have the same number of iterations.
+"""
+import datetime
+import os
+import time
+
+import detection.presets
+import torch
+import torch.utils.data
+import torchvision
+import torchvision.models.detection
+import torchvision.models.detection.mask_rcnn
+import detection.utils as utils
+from detection.coco_utils import get_coco, get_coco_kp
+from detection.engine import train_one_epoch, evaluate
+from detection.group_by_aspect_ratio import GroupedBatchSampler, create_aspect_ratio_groups
+
+
+try:
+    from torchvision.prototype import models as PM
+except ImportError:
+    PM = None
+
+
+def get_dataset(name, image_set, transform, data_path):
+    paths = {"coco": (data_path, get_coco, 91), "coco_kp": (data_path, get_coco_kp, 2)}
+    p, ds_fn, num_classes = paths[name]
+
+    ds = ds_fn(p, image_set=image_set, transforms=transform)
+    return ds, num_classes
+
+
+def get_transform(train, args):
+    if train:
+        return presets.DetectionPresetTrain(args.data_augmentation)
+    elif not args.weights:
+        return presets.DetectionPresetEval()
+    else:
+        weights = PM.get_weight(args.weights)
+        return weights.transforms()
+
+
+def get_args_parser(add_help=True):
+    import argparse
+
+    parser = argparse.ArgumentParser(description="PyTorch Detection Training", add_help=add_help)
+
+    parser.add_argument("--data-path", default="/datasets01/COCO/022719/", type=str, help="dataset path")
+    parser.add_argument("--dataset", default="coco", type=str, help="dataset name")
+    parser.add_argument("--model", default="maskrcnn_resnet50_fpn", type=str, help="model name")
+    parser.add_argument("--device", default="cuda", type=str, help="device (Use cuda or cpu Default: cuda)")
+    parser.add_argument(
+        "-b", "--batch-size", default=2, type=int, help="images per gpu, the total batch size is $NGPU x batch_size"
+    )
+    parser.add_argument("--epochs", default=26, type=int, metavar="N", help="number of total epochs to run")
+    parser.add_argument(
+        "-j", "--workers", default=4, type=int, metavar="N", help="number of data loading workers (default: 4)"
+    )
+    parser.add_argument(
+        "--lr",
+        default=0.02,
+        type=float,
+        help="initial learning rate, 0.02 is the default value for training on 8 gpus and 2 images_per_gpu",
+    )
+    parser.add_argument("--momentum", default=0.9, type=float, metavar="M", help="momentum")
+    parser.add_argument(
+        "--wd",
+        "--weight-decay",
+        default=1e-4,
+        type=float,
+        metavar="W",
+        help="weight decay (default: 1e-4)",
+        dest="weight_decay",
+    )
+    parser.add_argument(
+        "--lr-scheduler", default="multisteplr", type=str, help="name of lr scheduler (default: multisteplr)"
+    )
+    parser.add_argument(
+        "--lr-step-size", default=8, type=int, help="decrease lr every step-size epochs (multisteplr scheduler only)"
+    )
+    parser.add_argument(
+        "--lr-steps",
+        default=[16, 22],
+        nargs="+",
+        type=int,
+        help="decrease lr every step-size epochs (multisteplr scheduler only)",
+    )
+    parser.add_argument(
+        "--lr-gamma", default=0.1, type=float, help="decrease lr by a factor of lr-gamma (multisteplr scheduler only)"
+    )
+    parser.add_argument("--print-freq", default=20, type=int, help="print frequency")
+    parser.add_argument("--output-dir", default=".", type=str, help="path to save outputs")
+    parser.add_argument("--resume", default="", type=str, help="path of checkpoint")
+    parser.add_argument("--start_epoch", default=0, type=int, help="start epoch")
+    parser.add_argument("--aspect-ratio-group-factor", default=3, type=int)
+    parser.add_argument("--rpn-score-thresh", default=None, type=float, help="rpn score threshold for faster-rcnn")
+    parser.add_argument(
+        "--trainable-backbone-layers", default=None, type=int, help="number of trainable layers of backbone"
+    )
+    parser.add_argument(
+        "--data-augmentation", default="hflip", type=str, help="data augmentation policy (default: hflip)"
+    )
+    parser.add_argument(
+        "--sync-bn",
+        dest="sync_bn",
+        help="Use sync batch norm",
+        action="store_true",
+    )
+    parser.add_argument(
+        "--test-only",
+        dest="test_only",
+        help="Only test the model",
+        action="store_true",
+    )
+    parser.add_argument(
+        "--pretrained",
+        dest="pretrained",
+        help="Use pre-trained models from the modelzoo",
+        action="store_true",
+    )
+
+    # distributed training parameters
+    parser.add_argument("--world-size", default=1, type=int, help="number of distributed processes")
+    parser.add_argument("--dist-url", default="env://", type=str, help="url used to set up distributed training")
+
+    # Prototype models only
+    parser.add_argument("--weights", default=None, type=str, help="the weights enum name to load")
+
+    # Mixed precision training parameters
+    parser.add_argument("--amp", action="store_true", help="Use torch.cuda.amp for mixed precision training")
+
+    return parser
+
+
+def main(args):
+    if args.weights and PM is None:
+        raise ImportError("The prototype module couldn't be found. Please install the latest torchvision nightly.")
+    if args.output_dir:
+        utils.mkdir(args.output_dir)
+
+    utils.init_distributed_mode(args)
+    print(args)
+
+    device = torch.device(args.device)
+
+    # Data loading code
+    print("Loading data")
+
+    dataset, num_classes = get_dataset(args.dataset, "train", get_transform(True, args), args.data_path)
+    dataset_test, _ = get_dataset(args.dataset, "val", get_transform(False, args), args.data_path)
+
+    print("Creating data loaders")
+    if args.distributed:
+        train_sampler = torch.utils.data.distributed.DistributedSampler(dataset)
+        test_sampler = torch.utils.data.distributed.DistributedSampler(dataset_test)
+    else:
+        train_sampler = torch.utils.data.RandomSampler(dataset)
+        test_sampler = torch.utils.data.SequentialSampler(dataset_test)
+
+    if args.aspect_ratio_group_factor >= 0:
+        group_ids = create_aspect_ratio_groups(dataset, k=args.aspect_ratio_group_factor)
+        train_batch_sampler = GroupedBatchSampler(train_sampler, group_ids, args.batch_size)
+    else:
+        train_batch_sampler = torch.utils.data.BatchSampler(train_sampler, args.batch_size, drop_last=True)
+
+    data_loader = torch.utils.data.DataLoader(
+        dataset, batch_sampler=train_batch_sampler, num_workers=args.workers, collate_fn=utils.collate_fn
+    )
+
+    data_loader_test = torch.utils.data.DataLoader(
+        dataset_test, batch_size=1, sampler=test_sampler, num_workers=args.workers, collate_fn=utils.collate_fn
+    )
+
+    print("Creating model")
+    kwargs = {"trainable_backbone_layers": args.trainable_backbone_layers}
+    if "rcnn" in args.model:
+        if args.rpn_score_thresh is not None:
+            kwargs["rpn_score_thresh"] = args.rpn_score_thresh
+    if not args.weights:
+        model = torchvision.models.detection.__dict__[args.model](
+            pretrained=args.pretrained, num_classes=num_classes, **kwargs
+        )
+    else:
+        model = PM.detection.__dict__[args.model](weights=args.weights, num_classes=num_classes, **kwargs)
+    model.to(device)
+    if args.distributed and args.sync_bn:
+        model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
+
+    model_without_ddp = model
+    if args.distributed:
+        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])
+        model_without_ddp = model.module
+
+    params = [p for p in model.parameters() if p.requires_grad]
+    optimizer = torch.optim.SGD(params, lr=args.lr, momentum=args.momentum, weight_decay=args.weight_decay)
+
+    scaler = torch.cuda.amp.GradScaler() if args.amp else None
+
+    args.lr_scheduler = args.lr_scheduler.lower()
+    if args.lr_scheduler == "multisteplr":
+        lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=args.lr_steps, gamma=args.lr_gamma)
+    elif args.lr_scheduler == "cosineannealinglr":
+        lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs)
+    else:
+        raise RuntimeError(
+            f"Invalid lr scheduler '{args.lr_scheduler}'. Only MultiStepLR and CosineAnnealingLR are supported."
+        )
+
+    if args.resume:
+        checkpoint = torch.load(args.resume, map_location="cpu")
+        model_without_ddp.load_state_dict(checkpoint["model"])
+        optimizer.load_state_dict(checkpoint["optimizer"])
+        lr_scheduler.load_state_dict(checkpoint["lr_scheduler"])
+        args.start_epoch = checkpoint["epoch"] + 1
+        if args.amp:
+            scaler.load_state_dict(checkpoint["scaler"])
+
+    if args.test_only:
+        evaluate(model, data_loader_test, device=device)
+        return
+
+    print("Start training")
+    start_time = time.time()
+    for epoch in range(args.start_epoch, args.epochs):
+        if args.distributed:
+            train_sampler.set_epoch(epoch)
+        train_one_epoch(model, optimizer, data_loader, device, epoch, args.print_freq, scaler)
+        lr_scheduler.step()
+        if args.output_dir:
+            checkpoint = {
+                "model": model_without_ddp.state_dict(),
+                "optimizer": optimizer.state_dict(),
+                "lr_scheduler": lr_scheduler.state_dict(),
+                "args": args,
+                "epoch": epoch,
+            }
+            if args.amp:
+                checkpoint["scaler"] = scaler.state_dict()
+            utils.save_on_master(checkpoint, os.path.join(args.output_dir, f"model_{epoch}.pth"))
+            utils.save_on_master(checkpoint, os.path.join(args.output_dir, "checkpoint.pth"))
+
+        # evaluate after every epoch
+        evaluate(model, data_loader_test, device=device)
+
+    total_time = time.time() - start_time
+    total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+    print(f"Training time {total_time_str}")
+
+
+if __name__ == "__main__":
+    args = get_args_parser().parse_args()
+    main(args)

DenseMammogram/detection/transforms.py

@@ -0,0 +1,283 @@
+from typing import List, Tuple, Dict, Optional
+
+import torch
+import torchvision
+from torch import nn, Tensor
+from torchvision.transforms import functional as F
+from torchvision.transforms import transforms as T
+
+
+def _flip_coco_person_keypoints(kps, width):
+    flip_inds = [0, 2, 1, 4, 3, 6, 5, 8, 7, 10, 9, 12, 11, 14, 13, 16, 15]
+    flipped_data = kps[:, flip_inds]
+    flipped_data[..., 0] = width - flipped_data[..., 0]
+    # Maintain COCO convention that if visibility == 0, then x, y = 0
+    inds = flipped_data[..., 2] == 0
+    flipped_data[inds] = 0
+    return flipped_data
+
+
+class Compose:
+    def __init__(self, transforms):
+        self.transforms = transforms
+
+    def __call__(self, image, target):
+        for t in self.transforms:
+            image, target = t(image, target)
+        return image, target
+
+
+class RandomHorizontalFlip(T.RandomHorizontalFlip):
+    def forward(
+        self, image: Tensor, target: Optional[Dict[str, Tensor]] = None
+    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
+        if torch.rand(1) < self.p:
+            image = F.hflip(image)
+            if target is not None:
+                width, _ = F.get_image_size(image)
+                target["boxes"][:, [0, 2]] = width - target["boxes"][:, [2, 0]]
+                if "masks" in target:
+                    target["masks"] = target["masks"].flip(-1)
+                if "keypoints" in target:
+                    keypoints = target["keypoints"]
+                    keypoints = _flip_coco_person_keypoints(keypoints, width)
+                    target["keypoints"] = keypoints
+        return image, target
+
+
+class ToTensor(nn.Module):
+    def forward(
+        self, image: Tensor, target: Optional[Dict[str, Tensor]] = None
+    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
+        image = F.pil_to_tensor(image)
+        image = F.convert_image_dtype(image)
+        return image, target
+
+
+class PILToTensor(nn.Module):
+    def forward(
+        self, image: Tensor, target: Optional[Dict[str, Tensor]] = None
+    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
+        image = F.pil_to_tensor(image)
+        return image, target
+
+
+class ConvertImageDtype(nn.Module):
+    def __init__(self, dtype: torch.dtype) -> None:
+        super().__init__()
+        self.dtype = dtype
+
+    def forward(
+        self, image: Tensor, target: Optional[Dict[str, Tensor]] = None
+    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
+        image = F.convert_image_dtype(image, self.dtype)
+        return image, target
+
+
+class RandomIoUCrop(nn.Module):
+    def __init__(
+        self,
+        min_scale: float = 0.3,
+        max_scale: float = 1.0,
+        min_aspect_ratio: float = 0.5,
+        max_aspect_ratio: float = 2.0,
+        sampler_options: Optional[List[float]] = None,
+        trials: int = 40,
+    ):
+        super().__init__()
+        # Configuration similar to https://github.com/weiliu89/caffe/blob/ssd/examples/ssd/ssd_coco.py#L89-L174
+        self.min_scale = min_scale
+        self.max_scale = max_scale
+        self.min_aspect_ratio = min_aspect_ratio
+        self.max_aspect_ratio = max_aspect_ratio
+        if sampler_options is None:
+            sampler_options = [0.0, 0.1, 0.3, 0.5, 0.7, 0.9, 1.0]
+        self.options = sampler_options
+        self.trials = trials
+
+    def forward(
+        self, image: Tensor, target: Optional[Dict[str, Tensor]] = None
+    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
+        if target is None:
+            raise ValueError("The targets can't be None for this transform.")
+
+        if isinstance(image, torch.Tensor):
+            if image.ndimension() not in {2, 3}:
+                raise ValueError(f"image should be 2/3 dimensional. Got {image.ndimension()} dimensions.")
+            elif image.ndimension() == 2:
+                image = image.unsqueeze(0)
+
+        orig_w, orig_h = F.get_image_size(image)
+
+        while True:
+            # sample an option
+            idx = int(torch.randint(low=0, high=len(self.options), size=(1,)))
+            min_jaccard_overlap = self.options[idx]
+            if min_jaccard_overlap >= 1.0:  # a value larger than 1 encodes the leave as-is option
+                return image, target
+
+            for _ in range(self.trials):
+                # check the aspect ratio limitations
+                r = self.min_scale + (self.max_scale - self.min_scale) * torch.rand(2)
+                new_w = int(orig_w * r[0])
+                new_h = int(orig_h * r[1])
+                aspect_ratio = new_w / new_h
+                if not (self.min_aspect_ratio <= aspect_ratio <= self.max_aspect_ratio):
+                    continue
+
+                # check for 0 area crops
+                r = torch.rand(2)
+                left = int((orig_w - new_w) * r[0])
+                top = int((orig_h - new_h) * r[1])
+                right = left + new_w
+                bottom = top + new_h
+                if left == right or top == bottom:
+                    continue
+
+                # check for any valid boxes with centers within the crop area
+                cx = 0.5 * (target["boxes"][:, 0] + target["boxes"][:, 2])
+                cy = 0.5 * (target["boxes"][:, 1] + target["boxes"][:, 3])
+                is_within_crop_area = (left < cx) & (cx < right) & (top < cy) & (cy < bottom)
+                if not is_within_crop_area.any():
+                    continue
+
+                # check at least 1 box with jaccard limitations
+                boxes = target["boxes"][is_within_crop_area]
+                ious = torchvision.ops.boxes.box_iou(
+                    boxes, torch.tensor([[left, top, right, bottom]], dtype=boxes.dtype, device=boxes.device)
+                )
+                if ious.max() < min_jaccard_overlap:
+                    continue
+
+                # keep only valid boxes and perform cropping
+                target["boxes"] = boxes
+                target["labels"] = target["labels"][is_within_crop_area]
+                target["boxes"][:, 0::2] -= left
+                target["boxes"][:, 1::2] -= top
+                target["boxes"][:, 0::2].clamp_(min=0, max=new_w)
+                target["boxes"][:, 1::2].clamp_(min=0, max=new_h)
+                image = F.crop(image, top, left, new_h, new_w)
+
+                return image, target
+
+
+class RandomZoomOut(nn.Module):
+    def __init__(
+        self, fill: Optional[List[float]] = None, side_range: Tuple[float, float] = (1.0, 4.0), p: float = 0.5
+    ):
+        super().__init__()
+        if fill is None:
+            fill = [0.0, 0.0, 0.0]
+        self.fill = fill
+        self.side_range = side_range
+        if side_range[0] < 1.0 or side_range[0] > side_range[1]:
+            raise ValueError(f"Invalid canvas side range provided {side_range}.")
+        self.p = p
+
+    @torch.jit.unused
+    def _get_fill_value(self, is_pil):
+        # type: (bool) -> int
+        # We fake the type to make it work on JIT
+        return tuple(int(x) for x in self.fill) if is_pil else 0
+
+    def forward(
+        self, image: Tensor, target: Optional[Dict[str, Tensor]] = None
+    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
+        if isinstance(image, torch.Tensor):
+            if image.ndimension() not in {2, 3}:
+                raise ValueError(f"image should be 2/3 dimensional. Got {image.ndimension()} dimensions.")
+            elif image.ndimension() == 2:
+                image = image.unsqueeze(0)
+
+        if torch.rand(1) < self.p:
+            return image, target
+
+        orig_w, orig_h = F.get_image_size(image)
+
+        r = self.side_range[0] + torch.rand(1) * (self.side_range[1] - self.side_range[0])
+        canvas_width = int(orig_w * r)
+        canvas_height = int(orig_h * r)
+
+        r = torch.rand(2)
+        left = int((canvas_width - orig_w) * r[0])
+        top = int((canvas_height - orig_h) * r[1])
+        right = canvas_width - (left + orig_w)
+        bottom = canvas_height - (top + orig_h)
+
+        if torch.jit.is_scripting():
+            fill = 0
+        else:
+            fill = self._get_fill_value(F._is_pil_image(image))
+
+        image = F.pad(image, [left, top, right, bottom], fill=fill)
+        if isinstance(image, torch.Tensor):
+            v = torch.tensor(self.fill, device=image.device, dtype=image.dtype).view(-1, 1, 1)
+            image[..., :top, :] = image[..., :, :left] = image[..., (top + orig_h) :, :] = image[
+                ..., :, (left + orig_w) :
+            ] = v
+
+        if target is not None:
+            target["boxes"][:, 0::2] += left
+            target["boxes"][:, 1::2] += top
+
+        return image, target
+
+
+class RandomPhotometricDistort(nn.Module):
+    def __init__(
+        self,
+        contrast: Tuple[float] = (0.5, 1.5),
+        saturation: Tuple[float] = (0.5, 1.5),
+        hue: Tuple[float] = (-0.05, 0.05),
+        brightness: Tuple[float] = (0.875, 1.125),
+        p: float = 0.5,
+    ):
+        super().__init__()
+        self._brightness = T.ColorJitter(brightness=brightness)
+        self._contrast = T.ColorJitter(contrast=contrast)
+        self._hue = T.ColorJitter(hue=hue)
+        self._saturation = T.ColorJitter(saturation=saturation)
+        self.p = p
+
+    def forward(
+        self, image: Tensor, target: Optional[Dict[str, Tensor]] = None
+    ) -> Tuple[Tensor, Optional[Dict[str, Tensor]]]:
+        if isinstance(image, torch.Tensor):
+            if image.ndimension() not in {2, 3}:
+                raise ValueError(f"image should be 2/3 dimensional. Got {image.ndimension()} dimensions.")
+            elif image.ndimension() == 2:
+                image = image.unsqueeze(0)
+
+        r = torch.rand(7)
+
+        if r[0] < self.p:
+            image = self._brightness(image)
+
+        contrast_before = r[1] < 0.5
+        if contrast_before:
+            if r[2] < self.p:
+                image = self._contrast(image)
+
+        if r[3] < self.p:
+            image = self._saturation(image)
+
+        if r[4] < self.p:
+            image = self._hue(image)
+
+        if not contrast_before:
+            if r[5] < self.p:
+                image = self._contrast(image)
+
+        if r[6] < self.p:
+            channels = F.get_image_num_channels(image)
+            permutation = torch.randperm(channels)
+
+            is_pil = F._is_pil_image(image)
+            if is_pil:
+                image = F.pil_to_tensor(image)
+                image = F.convert_image_dtype(image)
+            image = image[..., permutation, :, :]
+            if is_pil:
+                image = F.to_pil_image(image)
+
+        return image, target

DenseMammogram/detection/utils.py

@@ -0,0 +1,282 @@
+import datetime
+import errno
+import os
+import time
+from collections import defaultdict, deque
+
+import torch
+import torch.distributed as dist
+
+
+class SmoothedValue:
+    """Track a series of values and provide access to smoothed values over a
+    window or the global series average.
+    """
+
+    def __init__(self, window_size=20, fmt=None):
+        if fmt is None:
+            fmt = "{median:.4f} ({global_avg:.4f})"
+        self.deque = deque(maxlen=window_size)
+        self.total = 0.0
+        self.count = 0
+        self.fmt = fmt
+
+    def update(self, value, n=1):
+        self.deque.append(value)
+        self.count += n
+        self.total += value * n
+
+    def synchronize_between_processes(self):
+        """
+        Warning: does not synchronize the deque!
+        """
+        if not is_dist_avail_and_initialized():
+            return
+        t = torch.tensor([self.count, self.total], dtype=torch.float64, device="cuda")
+        dist.barrier()
+        dist.all_reduce(t)
+        t = t.tolist()
+        self.count = int(t[0])
+        self.total = t[1]
+
+    @property
+    def median(self):
+        d = torch.tensor(list(self.deque))
+        return d.median().item()
+
+    @property
+    def avg(self):
+        d = torch.tensor(list(self.deque), dtype=torch.float32)
+        return d.mean().item()
+
+    @property
+    def global_avg(self):
+        return self.total / self.count
+
+    @property
+    def max(self):
+        return max(self.deque)
+
+    @property
+    def value(self):
+        return self.deque[-1]
+
+    def __str__(self):
+        return self.fmt.format(
+            median=self.median, avg=self.avg, global_avg=self.global_avg, max=self.max, value=self.value
+        )
+
+
+def all_gather(data):
+    """
+    Run all_gather on arbitrary picklable data (not necessarily tensors)
+    Args:
+        data: any picklable object
+    Returns:
+        list[data]: list of data gathered from each rank
+    """
+    world_size = get_world_size()
+    if world_size == 1:
+        return [data]
+    data_list = [None] * world_size
+    dist.all_gather_object(data_list, data)
+    return data_list
+
+
+def reduce_dict(input_dict, average=True):
+    """
+    Args:
+        input_dict (dict): all the values will be reduced
+        average (bool): whether to do average or sum
+    Reduce the values in the dictionary from all processes so that all processes
+    have the averaged results. Returns a dict with the same fields as
+    input_dict, after reduction.
+    """
+    world_size = get_world_size()
+    if world_size < 2:
+        return input_dict
+    with torch.inference_mode():
+        names = []
+        values = []
+        # sort the keys so that they are consistent across processes
+        for k in sorted(input_dict.keys()):
+            names.append(k)
+            values.append(input_dict[k])
+        values = torch.stack(values, dim=0)
+        dist.all_reduce(values)
+        if average:
+            values /= world_size
+        reduced_dict = {k: v for k, v in zip(names, values)}
+    return reduced_dict
+
+
+class MetricLogger:
+    def __init__(self, delimiter="\t"):
+        self.meters = defaultdict(SmoothedValue)
+        self.delimiter = delimiter
+
+    def update(self, **kwargs):
+        for k, v in kwargs.items():
+            if isinstance(v, torch.Tensor):
+                v = v.item()
+            assert isinstance(v, (float, int))
+            self.meters[k].update(v)
+
+    def __getattr__(self, attr):
+        if attr in self.meters:
+            return self.meters[attr]
+        if attr in self.__dict__:
+            return self.__dict__[attr]
+        raise AttributeError(f"'{type(self).__name__}' object has no attribute '{attr}'")
+
+    def __str__(self):
+        loss_str = []
+        for name, meter in self.meters.items():
+            loss_str.append(f"{name}: {str(meter)}")
+        return self.delimiter.join(loss_str)
+
+    def synchronize_between_processes(self):
+        for meter in self.meters.values():
+            meter.synchronize_between_processes()
+
+    def add_meter(self, name, meter):
+        self.meters[name] = meter
+
+    def log_every(self, iterable, print_freq, header=None):
+        i = 0
+        if not header:
+            header = ""
+        start_time = time.time()
+        end = time.time()
+        iter_time = SmoothedValue(fmt="{avg:.4f}")
+        data_time = SmoothedValue(fmt="{avg:.4f}")
+        space_fmt = ":" + str(len(str(len(iterable)))) + "d"
+        if torch.cuda.is_available():
+            log_msg = self.delimiter.join(
+                [
+                    header,
+                    "[{0" + space_fmt + "}/{1}]",
+                    "eta: {eta}",
+                    "{meters}",
+                    "time: {time}",
+                    "data: {data}",
+                    "max mem: {memory:.0f}",
+                ]
+            )
+        else:
+            log_msg = self.delimiter.join(
+                [header, "[{0" + space_fmt + "}/{1}]", "eta: {eta}", "{meters}", "time: {time}", "data: {data}"]
+            )
+        MB = 1024.0 * 1024.0
+        for obj in iterable:
+            data_time.update(time.time() - end)
+            yield obj
+            iter_time.update(time.time() - end)
+            if i % print_freq == 0 or i == len(iterable) - 1:
+                eta_seconds = iter_time.global_avg * (len(iterable) - i)
+                eta_string = str(datetime.timedelta(seconds=int(eta_seconds)))
+                if torch.cuda.is_available():
+                    print(
+                        log_msg.format(
+                            i,
+                            len(iterable),
+                            eta=eta_string,
+                            meters=str(self),
+                            time=str(iter_time),
+                            data=str(data_time),
+                            memory=torch.cuda.max_memory_allocated() / MB,
+                        )
+                    )
+                else:
+                    print(
+                        log_msg.format(
+                            i, len(iterable), eta=eta_string, meters=str(self), time=str(iter_time), data=str(data_time)
+                        )
+                    )
+            i += 1
+            end = time.time()
+        total_time = time.time() - start_time
+        total_time_str = str(datetime.timedelta(seconds=int(total_time)))
+        print(f"{header} Total time: {total_time_str} ({total_time / len(iterable):.4f} s / it)")
+
+
+def collate_fn(batch):
+    return tuple(zip(*batch))
+
+
+def mkdir(path):
+    try:
+        os.makedirs(path)
+    except OSError as e:
+        if e.errno != errno.EEXIST:
+            raise
+
+
+def setup_for_distributed(is_master):
+    """
+    This function disables printing when not in master process
+    """
+    import builtins as __builtin__
+
+    builtin_print = __builtin__.print
+
+    def print(*args, **kwargs):
+        force = kwargs.pop("force", False)
+        if is_master or force:
+            builtin_print(*args, **kwargs)
+
+    __builtin__.print = print
+
+
+def is_dist_avail_and_initialized():
+    if not dist.is_available():
+        return False
+    if not dist.is_initialized():
+        return False
+    return True
+
+
+def get_world_size():
+    if not is_dist_avail_and_initialized():
+        return 1
+    return dist.get_world_size()
+
+
+def get_rank():
+    if not is_dist_avail_and_initialized():
+        return 0
+    return dist.get_rank()
+
+
+def is_main_process():
+    return get_rank() == 0
+
+
+def save_on_master(*args, **kwargs):
+    if is_main_process():
+        torch.save(*args, **kwargs)
+
+
+def init_distributed_mode(args):
+    if "RANK" in os.environ and "WORLD_SIZE" in os.environ:
+        args.rank = int(os.environ["RANK"])
+        args.world_size = int(os.environ["WORLD_SIZE"])
+        args.gpu = int(os.environ["LOCAL_RANK"])
+    elif "SLURM_PROCID" in os.environ:
+        args.rank = int(os.environ["SLURM_PROCID"])
+        args.gpu = args.rank % torch.cuda.device_count()
+    else:
+        print("Not using distributed mode")
+        args.distributed = False
+        return
+
+    args.distributed = True
+
+    torch.cuda.set_device(args.gpu)
+    args.dist_backend = "nccl"
+    print(f"| distributed init (rank {args.rank}): {args.dist_url}", flush=True)
+    torch.distributed.init_process_group(
+        backend=args.dist_backend, init_method=args.dist_url, world_size=args.world_size, rank=args.rank
+    )
+    torch.distributed.barrier()
+    setup_for_distributed(args.rank == 0)

DenseMammogram/ensemble_boxes/__init__.py

@@ -0,0 +1,9 @@
+# coding: utf-8
+__author__ = 'ZFTurbo: https://kaggle.com/zfturbo'
+
+from .ensemble_boxes_wbf import weighted_boxes_fusion
+from .ensemble_boxes_nmw import non_maximum_weighted
+from .ensemble_boxes_nms import nms_method
+from .ensemble_boxes_nms import nms
+from .ensemble_boxes_nms import soft_nms
+from .ensemble_boxes_wbf_3d import weighted_boxes_fusion_3d

DenseMammogram/ensemble_boxes/ensemble_boxes_nms.py

@@ -0,0 +1,249 @@
+# coding: utf-8
+__author__ = 'ZFTurbo: https://kaggle.com/zfturbo'
+
+import numpy as np
+from numba import jit
+
+
+def prepare_boxes(boxes, scores, labels):
+    result_boxes = boxes.copy()
+
+    cond = (result_boxes < 0)
+    cond_sum = cond.astype(np.int32).sum()
+    if cond_sum > 0:
+        print('Warning. Fixed {} boxes coordinates < 0'.format(cond_sum))
+        result_boxes[cond] = 0
+
+    cond = (result_boxes > 1)
+    cond_sum = cond.astype(np.int32).sum()
+    if cond_sum > 0:
+        print('Warning. Fixed {} boxes coordinates > 1. Check that your boxes was normalized at [0, 1]'.format(cond_sum))
+        result_boxes[cond] = 1
+
+    boxes1 = result_boxes.copy()
+    result_boxes[:, 0] = np.min(boxes1[:, [0, 2]], axis=1)
+    result_boxes[:, 2] = np.max(boxes1[:, [0, 2]], axis=1)
+    result_boxes[:, 1] = np.min(boxes1[:, [1, 3]], axis=1)
+    result_boxes[:, 3] = np.max(boxes1[:, [1, 3]], axis=1)
+
+    area = (result_boxes[:, 2] - result_boxes[:, 0]) * (result_boxes[:, 3] - result_boxes[:, 1])
+    cond = (area == 0)
+    cond_sum = cond.astype(np.int32).sum()
+    if cond_sum > 0:
+        print('Warning. Removed {} boxes with zero area!'.format(cond_sum))
+        result_boxes = result_boxes[area > 0]
+        scores = scores[area > 0]
+        labels = labels[area > 0]
+
+    return result_boxes, scores, labels
+
+
+def cpu_soft_nms_float(dets, sc, Nt, sigma, thresh, method):
+    """
+    Based on: https://github.com/DocF/Soft-NMS/blob/master/soft_nms.py
+    It's different from original soft-NMS because we have float coordinates on range [0; 1]
+
+    :param dets:   boxes format [x1, y1, x2, y2]
+    :param sc:     scores for boxes
+    :param Nt:     required iou 
+    :param sigma:  
+    :param thresh: 
+    :param method: 1 - linear soft-NMS, 2 - gaussian soft-NMS, 3 - standard NMS
+    :return: index of boxes to keep
+    """
+
+    # indexes concatenate boxes with the last column
+    N = dets.shape[0]
+    indexes = np.array([np.arange(N)])
+    dets = np.concatenate((dets, indexes.T), axis=1)
+
+    # the order of boxes coordinate is [y1, x1, y2, x2]
+    y1 = dets[:, 1]
+    x1 = dets[:, 0]
+    y2 = dets[:, 3]
+    x2 = dets[:, 2]
+    scores = sc
+    areas = (x2 - x1) * (y2 - y1)
+
+    for i in range(N):
+        # intermediate parameters for later parameters exchange
+        tBD = dets[i, :].copy()
+        tscore = scores[i].copy()
+        tarea = areas[i].copy()
+        pos = i + 1
+
+        #
+        if i != N - 1:
+            maxscore = np.max(scores[pos:], axis=0)
+            maxpos = np.argmax(scores[pos:], axis=0)
+        else:
+            maxscore = scores[-1]
+            maxpos = 0
+        if tscore < maxscore:
+            dets[i, :] = dets[maxpos + i + 1, :]
+            dets[maxpos + i + 1, :] = tBD
+            tBD = dets[i, :]
+
+            scores[i] = scores[maxpos + i + 1]
+            scores[maxpos + i + 1] = tscore
+            tscore = scores[i]
+
+            areas[i] = areas[maxpos + i + 1]
+            areas[maxpos + i + 1] = tarea
+            tarea = areas[i]
+
+        # IoU calculate
+        xx1 = np.maximum(dets[i, 1], dets[pos:, 1])
+        yy1 = np.maximum(dets[i, 0], dets[pos:, 0])
+        xx2 = np.minimum(dets[i, 3], dets[pos:, 3])
+        yy2 = np.minimum(dets[i, 2], dets[pos:, 2])
+
+        w = np.maximum(0.0, xx2 - xx1)
+        h = np.maximum(0.0, yy2 - yy1)
+        inter = w * h
+        ovr = inter / (areas[i] + areas[pos:] - inter)
+
+        # Three methods: 1.linear 2.gaussian 3.original NMS
+        if method == 1:  # linear
+            weight = np.ones(ovr.shape)
+            weight[ovr > Nt] = weight[ovr > Nt] - ovr[ovr > Nt]
+        elif method == 2:  # gaussian
+            weight = np.exp(-(ovr * ovr) / sigma)
+        else:  # original NMS
+            weight = np.ones(ovr.shape)
+            weight[ovr > Nt] = 0
+
+        scores[pos:] = weight * scores[pos:]
+
+    # select the boxes and keep the corresponding indexes
+    inds = dets[:, 4][scores > thresh]
+    keep = inds.astype(int)
+    return keep
+
+
+@jit(nopython=True)
+def nms_float_fast(dets, scores, thresh):
+    """
+    # It's different from original nms because we have float coordinates on range [0; 1]
+    :param dets: numpy array of boxes with shape: (N, 5). Order: x1, y1, x2, y2, score. All variables in range [0; 1]
+    :param thresh: IoU value for boxes
+    :return: index of boxes to keep
+    """
+    x1 = dets[:, 0]
+    y1 = dets[:, 1]
+    x2 = dets[:, 2]
+    y2 = dets[:, 3]
+
+    areas = (x2 - x1) * (y2 - y1)
+    order = scores.argsort()[::-1]
+
+    keep = []
+    while order.size > 0:
+        i = order[0]
+        keep.append(i)
+        xx1 = np.maximum(x1[i], x1[order[1:]])
+        yy1 = np.maximum(y1[i], y1[order[1:]])
+        xx2 = np.minimum(x2[i], x2[order[1:]])
+        yy2 = np.minimum(y2[i], y2[order[1:]])
+
+        w = np.maximum(0.0, xx2 - xx1)
+        h = np.maximum(0.0, yy2 - yy1)
+        inter = w * h
+        ovr = inter / (areas[i] + areas[order[1:]] - inter)
+        inds = np.where(ovr <= thresh)[0]
+        order = order[inds + 1]
+
+    return keep
+
+
+def nms_method(boxes, scores, labels, method=3, iou_thr=0.5, sigma=0.5, thresh=0.001, weights=None):
+    """
+    :param boxes: list of boxes predictions from each model, each box is 4 numbers. 
+    It has 3 dimensions (models_number, model_preds, 4)
+    Order of boxes: x1, y1, x2, y2. We expect float normalized coordinates [0; 1] 
+    :param scores: list of scores for each model 
+    :param labels: list of labels for each model
+    :param method: 1 - linear soft-NMS, 2 - gaussian soft-NMS, 3 - standard NMS
+    :param iou_thr: IoU value for boxes to be a match 
+    :param sigma: Sigma value for SoftNMS
+    :param thresh: threshold for boxes to keep (important for SoftNMS)
+    :param weights: list of weights for each model. Default: None, which means weight == 1 for each model
+
+    :return: boxes: boxes coordinates (Order of boxes: x1, y1, x2, y2). 
+    :return: scores: confidence scores
+    :return: labels: boxes labels
+    """
+
+    # If weights are specified
+    if weights is not None:
+        if len(boxes) != len(weights):
+            print('Incorrect number of weights: {}. Must be: {}. Skip it'.format(len(weights), len(boxes)))
+        else:
+            weights = np.array(weights)
+            for i in range(len(weights)):
+                scores[i] = (np.array(scores[i]) * weights[i]) / weights.sum()
+
+    # We concatenate everything
+    boxes = np.concatenate(boxes)
+    scores = np.concatenate(scores)
+    labels = np.concatenate(labels)
+
+    # Fix coordinates and removed zero area boxes
+    boxes, scores, labels = prepare_boxes(boxes, scores, labels)
+
+    # Run NMS independently for each label
+    unique_labels = np.unique(labels)
+    final_boxes = []
+    final_scores = []
+    final_labels = []
+    for l in unique_labels:
+        condition = (labels == l)
+        boxes_by_label = boxes[condition]
+        scores_by_label = scores[condition]
+        labels_by_label = np.array([l] * len(boxes_by_label))
+
+        if method != 3:
+            keep = cpu_soft_nms_float(boxes_by_label.copy(), scores_by_label.copy(), Nt=iou_thr, sigma=sigma, thresh=thresh, method=method)
+        else:
+            # Use faster function
+            keep = nms_float_fast(boxes_by_label, scores_by_label, thresh=iou_thr)
+
+        final_boxes.append(boxes_by_label[keep])
+        final_scores.append(scores_by_label[keep])
+        final_labels.append(labels_by_label[keep])
+    final_boxes = np.concatenate(final_boxes)
+    final_scores = np.concatenate(final_scores)
+    final_labels = np.concatenate(final_labels)
+
+    return final_boxes, final_scores, final_labels
+
+
+def nms(boxes, scores, labels, iou_thr=0.5, weights=None):
+    """
+    Short call for standard NMS 
+    
+    :param boxes: 
+    :param scores: 
+    :param labels: 
+    :param iou_thr: 
+    :param weights: 
+    :return: 
+    """
+    return nms_method(boxes, scores, labels, method=3, iou_thr=iou_thr, weights=weights)
+
+
+def soft_nms(boxes, scores, labels, method=2, iou_thr=0.5, sigma=0.5, thresh=0.001, weights=None):
+    """
+    Short call for Soft-NMS
+     
+    :param boxes: 
+    :param scores: 
+    :param labels: 
+    :param method: 
+    :param iou_thr: 
+    :param sigma: 
+    :param thresh: 
+    :param weights: 
+    :return: 
+    """
+    return nms_method(boxes, scores, labels, method=method, iou_thr=iou_thr, sigma=sigma, thresh=thresh, weights=weights)

DenseMammogram/ensemble_boxes/ensemble_boxes_nmw.py

@@ -0,0 +1,202 @@
+# coding: utf-8
+__author__ = 'ZFTurbo: https://kaggle.com/zfturbo'
+
+"""
+Method described in:
+CAD: Scale Invariant Framework for Real-Time Object Detection
+http://openaccess.thecvf.com/content_ICCV_2017_workshops/papers/w14/Zhou_CAD_Scale_Invariant_ICCV_2017_paper.pdf
+"""
+
+import warnings
+import numpy as np
+from numba import jit
+
+
+@jit(nopython=True)
+def bb_intersection_over_union(A, B):
+    xA = max(A[0], B[0])
+    yA = max(A[1], B[1])
+    xB = min(A[2], B[2])
+    yB = min(A[3], B[3])
+
+    # compute the area of intersection rectangle
+    interArea = max(0, xB - xA) * max(0, yB - yA)
+
+    if interArea == 0:
+        return 0.0
+
+    # compute the area of both the prediction and ground-truth rectangles
+    boxAArea = (A[2] - A[0]) * (A[3] - A[1])
+    boxBArea = (B[2] - B[0]) * (B[3] - B[1])
+
+    iou = interArea / float(boxAArea + boxBArea - interArea)
+    return iou
+
+
+def prefilter_boxes(boxes, scores, labels, weights, thr):
+    # Create dict with boxes stored by its label
+    new_boxes = dict()
+    for t in range(len(boxes)):
+
+        if len(boxes[t]) != len(scores[t]):
+            print('Error. Length of boxes arrays not equal to length of scores array: {} != {}'.format(len(boxes[t]),
+                                                                                                       len(scores[t])))
+            exit()
+
+        if len(boxes[t]) != len(labels[t]):
+            print('Error. Length of boxes arrays not equal to length of labels array: {} != {}'.format(len(boxes[t]),
+                                                                                                       len(labels[t])))
+            exit()
+
+        for j in range(len(boxes[t])):
+            score = scores[t][j]
+            if score < thr:
+                continue
+            label = int(labels[t][j])
+            box_part = boxes[t][j]
+            x1 = float(box_part[0])
+            y1 = float(box_part[1])
+            x2 = float(box_part[2])
+            y2 = float(box_part[3])
+
+            # Box data checks
+            if x2 < x1:
+                warnings.warn('X2 < X1 value in box. Swap them.')
+                x1, x2 = x2, x1
+            if y2 < y1:
+                warnings.warn('Y2 < Y1 value in box. Swap them.')
+                y1, y2 = y2, y1
+            if x1 < 0:
+                warnings.warn('X1 < 0 in box. Set it to 0.')
+                x1 = 0
+            if x1 > 1:
+                warnings.warn('X1 > 1 in box. Set it to 1. Check that you normalize boxes in [0, 1] range.')
+                x1 = 1
+            if x2 < 0:
+                warnings.warn('X2 < 0 in box. Set it to 0.')
+                x2 = 0
+            if x2 > 1:
+                warnings.warn('X2 > 1 in box. Set it to 1. Check that you normalize boxes in [0, 1] range.')
+                x2 = 1
+            if y1 < 0:
+                warnings.warn('Y1 < 0 in box. Set it to 0.')
+                y1 = 0
+            if y1 > 1:
+                warnings.warn('Y1 > 1 in box. Set it to 1. Check that you normalize boxes in [0, 1] range.')
+                y1 = 1
+            if y2 < 0:
+                warnings.warn('Y2 < 0 in box. Set it to 0.')
+                y2 = 0
+            if y2 > 1:
+                warnings.warn('Y2 > 1 in box. Set it to 1. Check that you normalize boxes in [0, 1] range.')
+                y2 = 1
+            if (x2 - x1) * (y2 - y1) == 0.0:
+                warnings.warn("Zero area box skipped: {}.".format(box_part))
+                continue
+
+            b = [int(label), float(score) * weights[t], x1, y1, x2, y2]
+            if label not in new_boxes:
+                new_boxes[label] = []
+            new_boxes[label].append(b)
+
+    # Sort each list in dict by score and transform it to numpy array
+    for k in new_boxes:
+        current_boxes = np.array(new_boxes[k])
+        new_boxes[k] = current_boxes[current_boxes[:, 1].argsort()[::-1]]
+
+    return new_boxes
+
+
+def get_weighted_box(boxes):
+    """
+    Create weighted box for set of boxes
+    :param boxes: set of boxes to fuse 
+    :return: weighted box
+    """
+
+    box = np.zeros(6, dtype=np.float32)
+    best_box = boxes[0]
+    conf = 0
+    for b in boxes:
+        iou = bb_intersection_over_union(b[2:], best_box[2:])
+        weight = b[1] * iou
+        box[2:] += (weight * b[2:])
+        conf += weight
+    box[0] = best_box[0]
+    box[1] = best_box[1]
+    box[2:] /= conf
+    return box
+
+
+def find_matching_box(boxes_list, new_box, match_iou):
+    best_iou = match_iou
+    best_index = -1
+    for i in range(len(boxes_list)):
+        box = boxes_list[i]
+        if box[0] != new_box[0]:
+            continue
+        iou = bb_intersection_over_union(box[2:], new_box[2:])
+        if iou > best_iou:
+            best_index = i
+            best_iou = iou
+
+    return best_index, best_iou
+
+
+def non_maximum_weighted(boxes_list, scores_list, labels_list, weights=None, iou_thr=0.55, skip_box_thr=0.0):
+    '''
+    :param boxes_list: list of boxes predictions from each model, each box is 4 numbers. 
+    It has 3 dimensions (models_number, model_preds, 4)
+    Order of boxes: x1, y1, x2, y2. We expect float normalized coordinates [0; 1]
+    :param scores_list: list of scores for each model 
+    :param labels_list: list of labels for each model
+    :param weights: list of weights for each model. Default: None, which means weight == 1 for each model
+    :param iou_thr: IoU value for boxes to be a match
+    :param skip_box_thr: exclude boxes with score lower than this variable  
+    
+    :return: boxes: boxes coordinates (Order of boxes: x1, y1, x2, y2). 
+    :return: scores: confidence scores
+    :return: labels: boxes labels
+    '''
+
+    if weights is None:
+        weights = np.ones(len(boxes_list))
+    if len(weights) != len(boxes_list):
+        print('Warning: incorrect number of weights {}. Must be: {}. Set weights equal to 1.'.format(len(weights), len(boxes_list)))
+        weights = np.ones(len(boxes_list))
+    weights = np.array(weights) / max(weights)
+    # for i in range(len(weights)):
+    #     scores_list[i] = (np.array(scores_list[i]) * weights[i])
+
+    filtered_boxes = prefilter_boxes(boxes_list, scores_list, labels_list, weights, skip_box_thr)
+    if len(filtered_boxes) == 0:
+        return np.zeros((0, 4)), np.zeros((0,)), np.zeros((0,))
+
+    overall_boxes = []
+    for label in filtered_boxes:
+        boxes = filtered_boxes[label]
+        new_boxes = []
+        main_boxes = []
+
+        # Clusterize boxes
+        for j in range(0, len(boxes)):
+            index, best_iou = find_matching_box(main_boxes, boxes[j], iou_thr)
+            if index != -1:
+                new_boxes[index].append(boxes[j].copy())
+            else:
+                new_boxes.append([boxes[j].copy()])
+                main_boxes.append(boxes[j].copy())
+
+        weighted_boxes = []
+        for j in range(0, len(new_boxes)):
+            box = get_weighted_box(new_boxes[j])
+            weighted_boxes.append(box.copy())
+
+        overall_boxes.append(np.array(weighted_boxes))
+
+    overall_boxes = np.concatenate(overall_boxes, axis=0)
+    overall_boxes = overall_boxes[overall_boxes[:, 1].argsort()[::-1]]
+    boxes = overall_boxes[:, 2:]
+    scores = overall_boxes[:, 1]
+    labels = overall_boxes[:, 0]
+    return boxes, scores, labels

DenseMammogram/ensemble_boxes/ensemble_boxes_wbf.py

@@ -0,0 +1,269 @@
+# coding: utf-8
+__author__ = 'ZFTurbo: https://kaggle.com/zfturbo'
+
+
+import warnings
+import numpy as np
+from numba import jit
+import time
+
+@jit(nopython=True)
+def bb_intersection_over_union(A, B) -> float:
+    xA = max(A[0], B[0])
+    yA = max(A[1], B[1])
+    xB = min(A[2], B[2])
+    yB = min(A[3], B[3])
+
+    # compute the area of intersection rectangle
+    interArea = max(0, xB - xA) * max(0, yB - yA)
+
+    if interArea == 0:
+        return 0.0
+
+    # compute the area of both the prediction and ground-truth rectangles
+    boxAArea = (A[2] - A[0]) * (A[3] - A[1])
+    boxBArea = (B[2] - B[0]) * (B[3] - B[1])
+
+    iou = interArea / float(boxAArea + boxBArea - interArea)
+    return iou
+
+
+def prefilter_boxes(boxes, scores, labels, weights, thr):
+    # Create dict with boxes stored by its label
+    new_boxes = dict()
+
+    for t in range(len(boxes)):
+
+        if len(boxes[t]) != len(scores[t]):
+            print('Error. Length of boxes arrays not equal to length of scores array: {} != {}'.format(len(boxes[t]), len(scores[t])))
+            exit()
+
+        if len(boxes[t]) != len(labels[t]):
+            print('Error. Length of boxes arrays not equal to length of labels array: {} != {}'.format(len(boxes[t]), len(labels[t])))
+            exit()
+
+        for j in range(len(boxes[t])):
+            score = scores[t][j]
+            if score < thr:
+                continue
+            label = int(labels[t][j])
+            box_part = boxes[t][j]
+            x1 = float(box_part[0])
+            y1 = float(box_part[1])
+            x2 = float(box_part[2])
+            y2 = float(box_part[3])
+
+            # Box data checks
+            if x2 < x1:
+                warnings.warn('X2 < X1 value in box. Swap them.')
+                x1, x2 = x2, x1
+            if y2 < y1:
+                warnings.warn('Y2 < Y1 value in box. Swap them.')
+                y1, y2 = y2, y1
+            if x1 < 0:
+                warnings.warn('X1 < 0 in box. Set it to 0.')
+                x1 = 0
+            if x1 > 1:
+                warnings.warn('X1 > 1 in box. Set it to 1. Check that you normalize boxes in [0, 1] range.')
+                x1 = 1
+            if x2 < 0:
+                warnings.warn('X2 < 0 in box. Set it to 0.')
+                x2 = 0
+            if x2 > 1:
+                warnings.warn('X2 > 1 in box. Set it to 1. Check that you normalize boxes in [0, 1] range.')
+                x2 = 1
+            if y1 < 0:
+                warnings.warn('Y1 < 0 in box. Set it to 0.')
+                y1 = 0
+            if y1 > 1:
+                warnings.warn('Y1 > 1 in box. Set it to 1. Check that you normalize boxes in [0, 1] range.')
+                y1 = 1
+            if y2 < 0:
+                warnings.warn('Y2 < 0 in box. Set it to 0.')
+                y2 = 0
+            if y2 > 1:
+                warnings.warn('Y2 > 1 in box. Set it to 1. Check that you normalize boxes in [0, 1] range.')
+                y2 = 1
+            if (x2 - x1) * (y2 - y1) == 0.0:
+                warnings.warn("Zero area box skipped: {}.".format(box_part))
+                continue
+
+            # [label, score, weight, model index, x1, y1, x2, y2]
+            b = [int(label), float(score) * weights[t], weights[t], t, x1, y1, x2, y2]
+            if label not in new_boxes:
+                new_boxes[label] = []
+            new_boxes[label].append(b)
+
+    # Sort each list in dict by score and transform it to numpy array
+    for k in new_boxes:
+        current_boxes = np.array(new_boxes[k])
+        new_boxes[k] = current_boxes[current_boxes[:, 1].argsort()[::-1]]
+
+    return new_boxes
+
+
+def get_weighted_box(boxes, conf_type='avg'):
+    """
+    Create weighted box for set of boxes
+    :param boxes: set of boxes to fuse
+    :param conf_type: type of confidence one of 'avg' or 'max'
+    :return: weighted box (label, score, weight, x1, y1, x2, y2)
+    """
+
+    box = np.zeros(8, dtype=np.float32)
+    conf = 0
+    conf_list = []
+    w = 0
+    for b in boxes:
+        box[4:] += (b[1] * b[4:])
+        conf += b[1]
+        conf_list.append(b[1])
+        w += b[2]
+    box[0] = boxes[0][0]
+    if conf_type == 'avg':
+        box[1] = conf / len(boxes)
+    elif conf_type == 'max':
+        box[1] = np.array(conf_list).max()
+    elif conf_type in ['box_and_model_avg', 'absent_model_aware_avg']:
+        box[1] = conf / len(boxes)
+    box[2] = w
+    box[3] = -1 # model index field is retained for consistensy but is not used.
+    box[4:] /= conf
+    return box
+
+
+def find_matching_box(boxes_list, new_box, match_iou):
+    best_iou = match_iou
+    best_index = -1
+    for i in range(len(boxes_list)):
+        box = boxes_list[i]
+        if box[0] != new_box[0]:
+            continue
+        iou = bb_intersection_over_union(box[4:], new_box[4:])
+        if iou > best_iou:
+            best_index = i
+            best_iou = iou
+
+    return best_index, best_iou
+
+
+def find_matching_box_quickly(boxes_list, new_box, match_iou):
+    """ Reimplementation of find_matching_box with numpy instead of loops. Gives significant speed up for larger arrays
+        (~100x). This was previously the bottleneck since the function is called for every entry in the array.
+    """
+    def bb_iou_array(boxes, new_box):
+        # bb interesection over union
+        xA = np.maximum(boxes[:, 0], new_box[0])
+        yA = np.maximum(boxes[:, 1], new_box[1])
+        xB = np.minimum(boxes[:, 2], new_box[2])
+        yB = np.minimum(boxes[:, 3], new_box[3])
+
+        interArea = np.maximum(xB - xA, 0) * np.maximum(yB - yA, 0)
+
+        # compute the area of both the prediction and ground-truth rectangles
+        boxAArea = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1])
+        boxBArea = (new_box[2] - new_box[0]) * (new_box[3] - new_box[1])
+
+        iou = interArea / (boxAArea + boxBArea - interArea)
+
+        return iou
+
+    if boxes_list.shape[0] == 0:
+        return -1, match_iou
+
+    # boxes = np.array(boxes_list)
+    boxes = boxes_list
+
+    ious = bb_iou_array(boxes[:, 4:], new_box[4:])
+
+    ious[boxes[:, 0] != new_box[0]] = -1
+
+    best_idx = np.argmax(ious)
+    best_iou = ious[best_idx]
+
+    if best_iou <= match_iou:
+        best_iou = match_iou
+        best_idx = -1
+
+    return best_idx, best_iou
+
+
+def weighted_boxes_fusion(boxes_list, scores_list, labels_list, weights=None, iou_thr=0.55, skip_box_thr=0.0, conf_type='avg', allows_overflow=False):
+    '''
+    :param boxes_list: list of boxes predictions from each model, each box is 4 numbers.
+    It has 3 dimensions (models_number, model_preds, 4)
+    Order of boxes: x1, y1, x2, y2. We expect float normalized coordinates [0; 1]
+    :param scores_list: list of scores for each model
+    :param labels_list: list of labels for each model
+    :param weights: list of weights for each model. Default: None, which means weight == 1 for each model
+    :param iou_thr: IoU value for boxes to be a match
+    :param skip_box_thr: exclude boxes with score lower than this variable
+    :param conf_type: how to calculate confidence in weighted boxes. 'avg': average value, 'max': maximum value, 'box_and_model_avg': box and model wise hybrid weighted average, 'absent_model_aware_avg': weighted average that takes into account the absent model.
+    :param allows_overflow: false if we want confidence score not exceed 1.0
+
+    :return: boxes: boxes coordinates (Order of boxes: x1, y1, x2, y2).
+    :return: scores: confidence scores
+    :return: labels: boxes labels
+    '''
+
+    if weights is None:
+        weights = np.ones(len(boxes_list))
+    if len(weights) != len(boxes_list):
+        print('Warning: incorrect number of weights {}. Must be: {}. Set weights equal to 1.'.format(len(weights), len(boxes_list)))
+        weights = np.ones(len(boxes_list))
+    weights = np.array(weights)
+
+    if conf_type not in ['avg', 'max', 'box_and_model_avg', 'absent_model_aware_avg']:
+        print('Unknown conf_type: {}. Must be "avg", "max" or "box_and_model_avg", or "absent_model_aware_avg"'.format(conf_type))
+        exit()
+
+    filtered_boxes = prefilter_boxes(boxes_list, scores_list, labels_list, weights, skip_box_thr)
+    if len(filtered_boxes) == 0:
+        return np.zeros((0, 4)), np.zeros((0,)), np.zeros((0,))
+
+    overall_boxes = []
+    for label in filtered_boxes:
+        boxes = filtered_boxes[label]
+        new_boxes = []
+        weighted_boxes = np.empty((0,8))
+        # Clusterize boxes
+        for j in range(0, len(boxes)):
+            index, best_iou = find_matching_box_quickly(weighted_boxes, boxes[j], iou_thr)
+
+            if index != -1:
+                new_boxes[index].append(boxes[j])
+                weighted_boxes[index] = get_weighted_box(new_boxes[index], conf_type)
+            else:
+                new_boxes.append([boxes[j].copy()])
+                weighted_boxes = np.vstack((weighted_boxes, boxes[j].copy()))
+        # Rescale confidence based on number of models and boxes
+        for i in range(len(new_boxes)):
+            clustered_boxes = np.array(new_boxes[i])
+            if conf_type == 'box_and_model_avg':
+                # weighted average for boxes
+                weighted_boxes[i, 1] = weighted_boxes[i, 1] * len(clustered_boxes) / weighted_boxes[i, 2]
+                # identify unique model index by model index column
+                _, idx = np.unique(clustered_boxes[:, 3], return_index=True)
+                # rescale by unique model weights
+                weighted_boxes[i, 1] = weighted_boxes[i, 1] *  clustered_boxes[idx, 2].sum() / weights.sum()
+            elif conf_type == 'absent_model_aware_avg':
+                # get unique model index in the cluster
+                models = np.unique(clustered_boxes[:, 3]).astype(int)
+                # create a mask to get unused model weights
+                mask = np.ones(len(weights), dtype=bool)
+                mask[models] = False
+                # absent model aware weighted average
+                weighted_boxes[i, 1] = weighted_boxes[i, 1] * len(clustered_boxes) / (weighted_boxes[i, 2] + weights[mask].sum())
+            elif conf_type == 'max':
+                weighted_boxes[i, 1] = weighted_boxes[i, 1] / weights.max()
+            elif not allows_overflow:
+                weighted_boxes[i, 1] = weighted_boxes[i, 1] * min(len(weights), len(clustered_boxes)) / weights.sum()
+            else:
+                weighted_boxes[i, 1] = weighted_boxes[i, 1] * len(clustered_boxes) / weights.sum()
+        overall_boxes.append(weighted_boxes)
+    overall_boxes = np.concatenate(overall_boxes, axis=0)
+    overall_boxes = overall_boxes[overall_boxes[:, 1].argsort()[::-1]]
+    boxes = overall_boxes[:, 4:]
+    scores = overall_boxes[:, 1]
+    labels = overall_boxes[:, 0]
+    return boxes, scores, labels

DenseMammogram/ensemble_boxes/ensemble_boxes_wbf_3d.py

@@ -0,0 +1,222 @@
+# coding: utf-8
+__author__ = 'ZFTurbo: https://kaggle.com/zfturbo'
+
+
+import warnings
+import numpy as np
+from numba import jit
+
+
+@jit(nopython=True)
+def bb_intersection_over_union_3d(A, B) -> float:
+    xA = max(A[0], B[0])
+    yA = max(A[1], B[1])
+    zA = max(A[2], B[2])
+    xB = min(A[3], B[3])
+    yB = min(A[4], B[4])
+    zB = min(A[5], B[5])
+
+    interVol = max(0, xB - xA) * max(0, yB - yA) * max(0, zB - zA)
+    if interVol == 0:
+        return 0.0
+
+    # compute the volume of both the prediction and ground-truth rectangular boxes
+    boxAVol = (A[3] - A[0]) * (A[4] - A[1]) * (A[5] - A[2])
+    boxBVol = (B[3] - B[0]) * (B[4] - B[1]) * (B[5] - B[2])
+
+    iou = interVol / float(boxAVol + boxBVol - interVol)
+    return iou
+
+
+def prefilter_boxes(boxes, scores, labels, weights, thr):
+    # Create dict with boxes stored by its label
+    new_boxes = dict()
+
+    for t in range(len(boxes)):
+
+        if len(boxes[t]) != len(scores[t]):
+            print('Error. Length of boxes arrays not equal to length of scores array: {} != {}'.format(len(boxes[t]), len(scores[t])))
+            exit()
+
+        if len(boxes[t]) != len(labels[t]):
+            print('Error. Length of boxes arrays not equal to length of labels array: {} != {}'.format(len(boxes[t]), len(labels[t])))
+            exit()
+
+        for j in range(len(boxes[t])):
+            score = scores[t][j]
+            if score < thr:
+                continue
+            label = int(labels[t][j])
+            box_part = boxes[t][j]
+            x1 = float(box_part[0])
+            y1 = float(box_part[1])
+            z1 = float(box_part[2])
+            x2 = float(box_part[3])
+            y2 = float(box_part[4])
+            z2 = float(box_part[5])
+
+            # Box data checks
+            if x2 < x1:
+                warnings.warn('X2 < X1 value in box. Swap them.')
+                x1, x2 = x2, x1
+            if y2 < y1:
+                warnings.warn('Y2 < Y1 value in box. Swap them.')
+                y1, y2 = y2, y1
+            if z2 < z1:
+                warnings.warn('Z2 < Z1 value in box. Swap them.')
+                z1, z2 = z2, z1
+            if x1 < 0:
+                warnings.warn('X1 < 0 in box. Set it to 0.')
+                x1 = 0
+            if x1 > 1:
+                warnings.warn('X1 > 1 in box. Set it to 1. Check that you normalize boxes in [0, 1] range.')
+                x1 = 1
+            if x2 < 0:
+                warnings.warn('X2 < 0 in box. Set it to 0.')
+                x2 = 0
+            if x2 > 1:
+                warnings.warn('X2 > 1 in box. Set it to 1. Check that you normalize boxes in [0, 1] range.')
+                x2 = 1
+            if y1 < 0:
+                warnings.warn('Y1 < 0 in box. Set it to 0.')
+                y1 = 0
+            if y1 > 1:
+                warnings.warn('Y1 > 1 in box. Set it to 1. Check that you normalize boxes in [0, 1] range.')
+                y1 = 1
+            if y2 < 0:
+                warnings.warn('Y2 < 0 in box. Set it to 0.')
+                y2 = 0
+            if y2 > 1:
+                warnings.warn('Y2 > 1 in box. Set it to 1. Check that you normalize boxes in [0, 1] range.')
+                y2 = 1
+            if z1 < 0:
+                warnings.warn('Z1 < 0 in box. Set it to 0.')
+                z1 = 0
+            if z1 > 1:
+                warnings.warn('Z1 > 1 in box. Set it to 1. Check that you normalize boxes in [0, 1] range.')
+                z1 = 1
+            if z2 < 0:
+                warnings.warn('Z2 < 0 in box. Set it to 0.')
+                z2 = 0
+            if z2 > 1:
+                warnings.warn('Z2 > 1 in box. Set it to 1. Check that you normalize boxes in [0, 1] range.')
+                z2 = 1
+            if (x2 - x1) * (y2 - y1) * (z2 - z1) == 0.0:
+                warnings.warn("Zero volume box skipped: {}.".format(box_part))
+                continue
+
+            b = [int(label), float(score) * weights[t], x1, y1, z1, x2, y2, z2]
+            if label not in new_boxes:
+                new_boxes[label] = []
+            new_boxes[label].append(b)
+
+    # Sort each list in dict by score and transform it to numpy array
+    for k in new_boxes:
+        current_boxes = np.array(new_boxes[k])
+        new_boxes[k] = current_boxes[current_boxes[:, 1].argsort()[::-1]]
+
+    return new_boxes
+
+
+def get_weighted_box(boxes, conf_type='avg'):
+    """
+    Create weighted box for set of boxes
+    :param boxes: set of boxes to fuse
+    :param conf_type: type of confidence one of 'avg' or 'max'
+    :return: weighted box
+    """
+
+    box = np.zeros(8, dtype=np.float32)
+    conf = 0
+    conf_list = []
+    for b in boxes:
+        box[2:] += (b[1] * b[2:])
+        conf += b[1]
+        conf_list.append(b[1])
+    box[0] = boxes[0][0]
+    if conf_type == 'avg':
+        box[1] = conf / len(boxes)
+    elif conf_type == 'max':
+        box[1] = np.array(conf_list).max()
+    box[2:] /= conf
+    return box
+
+
+def find_matching_box(boxes_list, new_box, match_iou):
+    best_iou = match_iou
+    best_index = -1
+    for i in range(len(boxes_list)):
+        box = boxes_list[i]
+        if box[0] != new_box[0]:
+            continue
+        iou = bb_intersection_over_union_3d(box[2:], new_box[2:])
+        if iou > best_iou:
+            best_index = i
+            best_iou = iou
+
+    return best_index, best_iou
+
+
+def weighted_boxes_fusion_3d(boxes_list, scores_list, labels_list, weights=None, iou_thr=0.55, skip_box_thr=0.0, conf_type='avg', allows_overflow=False):
+    '''
+    :param boxes_list: list of boxes predictions from each model, each box is 6 numbers.
+    It has 3 dimensions (models_number, model_preds, 6)
+    Order of boxes: x1, y1, z1, x2, y2 z2. We expect float normalized coordinates [0; 1]
+    :param scores_list: list of scores for each model
+    :param labels_list: list of labels for each model
+    :param weights: list of weights for each model. Default: None, which means weight == 1 for each model
+    :param iou_thr: IoU value for boxes to be a match
+    :param skip_box_thr: exclude boxes with score lower than this variable
+    :param conf_type: how to calculate confidence in weighted boxes. 'avg': average value, 'max': maximum value
+    :param allows_overflow: false if we want confidence score not exceed 1.0
+
+    :return: boxes: boxes coordinates (Order of boxes: x1, y1, z1, x2, y2, z2).
+    :return: scores: confidence scores
+    :return: labels: boxes labels
+    '''
+
+    if weights is None:
+        weights = np.ones(len(boxes_list))
+    if len(weights) != len(boxes_list):
+        print('Warning: incorrect number of weights {}. Must be: {}. Set weights equal to 1.'.format(len(weights), len(boxes_list)))
+        weights = np.ones(len(boxes_list))
+    weights = np.array(weights)
+
+    if conf_type not in ['avg', 'max']:
+        print('Error. Unknown conf_type: {}. Must be "avg" or "max". Use "avg"'.format(conf_type))
+        conf_type = 'avg'
+
+    filtered_boxes = prefilter_boxes(boxes_list, scores_list, labels_list, weights, skip_box_thr)
+    if len(filtered_boxes) == 0:
+        return np.zeros((0, 6)), np.zeros((0,)), np.zeros((0,))
+
+    overall_boxes = []
+    for label in filtered_boxes:
+        boxes = filtered_boxes[label]
+        new_boxes = []
+        weighted_boxes = []
+
+        # Clusterize boxes
+        for j in range(0, len(boxes)):
+            index, best_iou = find_matching_box(weighted_boxes, boxes[j], iou_thr)
+            if index != -1:
+                new_boxes[index].append(boxes[j])
+                weighted_boxes[index] = get_weighted_box(new_boxes[index], conf_type)
+            else:
+                new_boxes.append([boxes[j].copy()])
+                weighted_boxes.append(boxes[j].copy())
+
+        # Rescale confidence based on number of models and boxes
+        for i in range(len(new_boxes)):
+            if not allows_overflow:
+                weighted_boxes[i][1] = weighted_boxes[i][1] * min(weights.sum(), len(new_boxes[i])) / weights.sum()
+            else:
+                weighted_boxes[i][1] = weighted_boxes[i][1] * len(new_boxes[i]) / weights.sum()
+        overall_boxes.append(np.array(weighted_boxes))
+
+    overall_boxes = np.concatenate(overall_boxes, axis=0)
+    overall_boxes = overall_boxes[overall_boxes[:, 1].argsort()[::-1]]
+    boxes = overall_boxes[:, 2:]
+    scores = overall_boxes[:, 1]
+    labels = overall_boxes[:, 0]
+    return boxes, scores, labels

DenseMammogram/experimenter.py

@@ -0,0 +1,213 @@
+# Experimenter Class is responsible for mainly four things:
+# 1. Configuration - Done
+# 2. Logging using the AdvancedLogger class - Almost Done
+# 3. Model Handling, including loading and saving models - Done(Upgrades Left)
+# 4. Running Different Variants Paralelly/Sequentially of experiments
+# 5. Combining frcnn training followed by bilateral training and final froc calculation - Done
+# 6. Version Control
+
+from advanced_config import AdvancedConfig
+from advanced_logger import AdvancedLogger, LogPriority
+import os
+from os.path import join
+from plot_froc import plot_froc
+from train_frcnn import main as TRAIN_FRCNN
+from train_bilateral import main as TRAIN_BILATERAL
+import torch
+from model_utils import generate_predictions, generate_predictions_bilateral
+import argparse
+from dataloaders import get_dict
+from utils import create_backup
+from torch.utils.tensorboard import SummaryWriter
+
+class Experimenter:
+
+    def __init__(self, cfg_file, BASE_DIR = 'experiments'):
+        self.cfg_file = cfg_file        
+        
+        self.con = AdvancedConfig(cfg_file)
+        self.config = self.con.config
+        self.exp_dir = join(BASE_DIR,self.config['EXP_NAME'])
+        os.makedirs(self.exp_dir, exist_ok=True)
+        self.con.save(join(self.exp_dir,'config.cfg'))
+    
+        self.logger = AdvancedLogger(self.exp_dir)
+        self.logger.log('Experiment:',self.config['EXP_NAME'],priority = LogPriority.STATS)
+        self.logger.log('Experiment Description:', self.config['EXP_DESC'], priority = LogPriority.STATS)
+        self.logger.log('Config File:',self.cfg_file, priority = LogPriority.STATS)
+        self.logger.log('Experiment started', priority = LogPriority.LOW)
+        self.losses = dict()
+        self.frocs = dict()
+
+        self.writer = SummaryWriter(join(self.exp_dir,'tensor_logs'))
+
+        create_backup(backup_dir=join(self.exp_dir,'scripts'))
+
+    def log(self, *args, **kwargs):
+        self.logger.log(*args, **kwargs)
+
+
+    def init_losses(self,mode):
+        if mode == 'FRCNN' or mode == 'FRCNN_BILATERAL':
+            self.losses['frcnn_loss'] = []
+            self.frocs['frcnn_froc'] = []
+        elif mode == 'BILATERAL' or mode == 'FRCNN_BILATERAL':
+            self.losses['bilateral_loss'] = []
+            self.frocs['bilateral_froc'] = []
+
+    def start_epoch(self):
+        self.curr_epoch += 1
+        self.logger.log('Epoch:',self.curr_epoch, priority = LogPriority.MEDIUM)
+
+    def end_epoch(self, loss, model = None, device = None):
+        if self.curr_mode == 'FRCNN':
+            self.losses['frcnn_loss'].append(loss)
+            self.best_loss = min(self.losses['frcnn_loss'])
+            if self.config['EVAL_METHOD'] == 'FROC':
+                exp_name = self.config['EXP_NAME']
+                _, val_path, _ = self.init_paths()
+                generate_predictions(model,device,val_path,f'preds_frcnn_{exp_name}')
+                from froc_by_pranjal import get_froc_points
+                senses, _ = get_froc_points(f'preds_frcnn_{exp_name}', root_fol= join(self.config['DATA_DIR'],self.config['AIIMS_DATA'], self.config['AIIMS_VAL_SPLIT']), fps_req = [0.2])
+                self.frocs['frcnn_froc'].append(senses[0])
+                self.best_froc = max(self.frocs['frcnn_froc'])
+                self.logger.log(f'Val FROC: {senses[0]}', LogPriority.MEDIUM)
+                self.logger.log(f'Best FROC: {self.best_froc}')
+        elif self.curr_mode == 'BILATERAL':
+            self.losses['bilateral_loss'].append(loss)
+            self.best_loss = min(self.losses['bilateral_loss'])    
+            if self.config['EVAL_METHOD'] == 'FROC':
+                exp_name = self.config['EXP_NAME']
+                _, val_path, _ = self.init_paths()
+                data_dir = self.config['DATA_DIR']
+                print('Generating')
+                generate_predictions_bilateral(model,device,val_path,get_dict(data_dir,self.abs_path(self.config['AIIMS_CORRS_LIST'])),preds_folder = f'preds_bilateral_{exp_name}')
+                print('Generation Done')
+                from froc_by_pranjal import get_froc_points
+                senses, _ = get_froc_points(f'preds_bilateral_{exp_name}', root_fol= join(self.config['DATA_DIR'],self.config['AIIMS_DATA'], self.config['AIIMS_VAL_SPLIT']), fps_req = [0.1])
+                print('Reading Sens from',f'preds_bilateral_{exp_name}', join(self.config['DATA_DIR'],self.config['AIIMS_DATA'], self.config['AIIMS_VAL_SPLIT']),)
+                
+                self.frocs['bilateral_froc'].append(senses[0])
+                self.best_froc = max(self.frocs['bilateral_froc'])
+                self.logger.log(f'Val FROC: {senses[0]}', priority = LogPriority.MEDIUM)
+                self.logger.log(f'Best FROC: {self.best_froc}')
+  
+        self.writer.add_scalar(f"{self.curr_mode}/Loss/Valid", loss, self.curr_epoch)
+    
+    
+    
+    def save_model(self, model):
+        if self.curr_mode == 'FRCNN':
+            self.logger.log('Saving FRCNN Model', priority = LogPriority.LOW)
+            model_file = join(self.exp_dir,'frcnn_models',f'frcnn_model.pth')
+            if self.config['EVAL_METHOD']:
+                SAVE = self.best_froc == self.frocs['frcnn_froc'][-1]
+            else:
+                SAVE = self.best_loss == self.losses['frcnn_loss'][-1]
+        elif self.curr_mode == 'BILATERAL':
+            self.logger.log('Saving Bilateral Model', priority = LogPriority.LOW)
+            model_file = join(self.exp_dir,'bilateral_models',f'bilateral_model.pth')
+            if self.config['EVAL_METHOD'] == 'FROC':
+                SAVE = self.best_froc == self.frocs['bilateral_froc'][-1]
+            else:
+                SAVE = self.best_loss == self.losses['bilateral_loss'][-1]
+        os.makedirs(os.path.split(model_file)[0], exist_ok=True)
+        if SAVE:
+            torch.save(model.state_dict(), model_file)
+        
+        torch.save(model.state_dict(), f'{model_file[:-4]}_{self.curr_epoch}.pth')
+
+    def init_paths(self,):
+        train_path = join(self.config['DATA_DIR'], self.config['AIIMS_DATA'], self.config['AIIMS_TRAIN_SPLIT'])
+        val_path = join(self.config['DATA_DIR'],   self.config['AIIMS_DATA'], self.config['AIIMS_VAL_SPLIT'])
+        test_path = join(self.config['DATA_DIR'],  self.config['AIIMS_DATA'], self.config['AIIMS_TEST_SPLIT'])
+        return train_path, val_path, test_path
+
+    def abs_path(self, path):
+        return join(self.config['DATA_DIR'], path)
+
+    # Impure Function, upadtes the model with best state dicts 
+    def generate_predictions(self,model, device):
+        self.logger.log('Generating Predictions')
+        self.logger.flush()
+        exp_name = self.config['EXP_NAME']
+        train_path, val_path, test_path = self.init_paths()
+
+        # Load the best val_loss model's state dicts
+        if self.curr_mode == 'FRCNN':
+            model_file = join(self.exp_dir,'frcnn_models','frcnn_model.pth')
+        elif self.curr_mode == 'BILATERAL':
+            model_file = join(self.exp_dir,'bilateral_models','bilateral_model.pth')
+        model.load_state_dict(torch.load(model_file))
+
+        if self.curr_mode == 'FRCNN':
+            generate_predictions(model,device,train_path,f'preds_frcnn_{exp_name}')
+            generate_predictions(model,device,val_path,f'preds_frcnn_{exp_name}')
+            generate_predictions(model,device,test_path,f'preds_frcnn_{exp_name}')
+        elif self.curr_mode == 'BILATERAL':
+            data_dir = self.config['DATA_DIR']
+            generate_predictions_bilateral(model,device,train_path,get_dict(data_dir,self.abs_path(self.config['AIIMS_CORRS_LIST'])),'aiims',f'preds_bilateral_{exp_name}')
+            generate_predictions_bilateral(model,device,val_path,get_dict(data_dir,self.abs_path(self.config['AIIMS_CORRS_LIST'])),'aiims',f'preds_bilateral_{exp_name}')
+            generate_predictions_bilateral(model,device,test_path,get_dict(data_dir,self.abs_path(self.config['AIIMS_CORRS_LIST'])),'aiims',f'preds_bilateral_{exp_name}')
+            test_path = join(self.config['DATA_DIR'], self.config['AIIMS_DATA'], self.config['AIIMS_TEST_SPLIT'])
+
+    def run_experiment(self):
+
+        # First Determine the mode of running the experiment
+        mode = self.config['MODE']
+        self.init_losses(mode)
+        self.curr_mode = 'FRCNN'
+        self.curr_epoch = -1
+        self.best_loss = 999999
+        self.best_froc = 0
+        if mode == 'FRCNN':
+            TRAIN_FRCNN(self.config['FRCNN'], self)
+        elif mode == 'BILATERAL':
+            self.curr_mode = 'BILATERAL'
+            TRAIN_BILATERAL(self.config['BILATERAL'], self)
+        elif mode == 'FRCNN_BILATERAL':
+            TRAIN_FRCNN(self.config['FRCNN'], self)
+            self.curr_mode = 'BILATERAL'
+            self.curr_epoch = -1
+            self.best_loss = 999999
+            # Note the path to frcnn model must be the same as that dictated by experiment
+            self.config['BILATERAL']['FRCNN_MODEL_PATH'] = join(self.exp_dir,'frcnn_models','frcnn_model.pth')
+            TRAIN_BILATERAL(self.config['BILATERAL'], self)
+        
+        self.logger.log(f'Best Loss: {self.best_loss}', priority= LogPriority.STATS)
+        self.logger.log('Experiment Training and Generation Ended', priority = LogPriority.MEDIUM)
+
+        # Now evaluate the results
+
+        frcnn_file = join(self.exp_dir, 'senses_fps_frcnn.txt')
+        bilateral_file = join(self.exp_dir, 'senses_fps_bilateral.txt')
+        from froc_by_pranjal import get_froc_points
+        exp_name = self.config['EXP_NAME']
+        if mode == 'FRCNN' or mode == 'FRCNN_BILATERAL':
+            senses, fps = get_froc_points(f'preds_frcnn_{exp_name}', root_fol= join(self.config['DATA_DIR'],self.config['AIIMS_DATA'], self.config['AIIMS_TEST_SPLIT']), save_to = frcnn_file)
+            self.logger.log('FRCNN RESULTS', priority = LogPriority.STATS)
+            for s,f in zip(senses, fps):
+                self.logger.log(f'Sensitivty at {f}: {s}', priority = LogPriority.STATS)
+        if mode == 'BILATERAL' or mode == 'FRCNN_BILATERAL':
+            senses, fps = get_froc_points(f'preds_bilateral_{exp_name}', root_fol= join(self.config['DATA_DIR'],self.config['AIIMS_DATA'], self.config['AIIMS_TEST_SPLIT']), save_to = bilateral_file)
+            self.logger.log('BILATERAL RESULTS', priority = LogPriority.STATS)
+            for s,f in zip(senses, fps):
+                self.logger.log(f'Sensitivty at {f}: {s}', priority = LogPriority.STATS)
+
+
+        # Now draw the graphs.... If FRCNN and BILATERAL both done, draw them on one graph
+        # Else draw single graphs only
+        if mode == 'FRCNN':
+            plot_froc({frcnn_file : 'FRCNN'}, join(self.exp_dir,'plot.png'), TITLE = 'FRCNN FROC')
+        elif mode == 'BILATERAL':
+            plot_froc({bilateral_file : 'BILATERAL'}, join(self.exp_dir,'plot.png'), TITLE = 'BILATERAL FROC')
+        elif mode == 'FRCNN_BILATERAL':
+            plot_froc({frcnn_file : 'FRCNN', bilateral_file : 'BILATERAL'}, join(self.exp_dir,'plot.png'), TITLE = 'FRCNN vs BILATERAL FROC')
+        self.logger.flush()
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--cfg_file', type=str, default='configs/AIIMS_C1.cfg')
+    args = parser.parse_args()
+    exp = Experimenter(args.cfg_file)
+    exp.run_experiment()

DenseMammogram/froc_by_pranjal.py

@@ -0,0 +1,236 @@
+import os
+import glob
+import sys
+from os.path import join
+
+
+'''
+    Note: Anywhere empty boxes means [] and not [[]]
+'''
+
+
+def remove_true_positives(gts, preds):
+
+    def true_positive(gt, pred):
+        # If center of pred is inside the gt, it is a true positive
+        c_pred = ((pred[0]+pred[2])/2., (pred[1]+pred[3])/2.)
+        if (c_pred[0] >= gt[0] and c_pred[0] <= gt[2] and
+                c_pred[1] >= gt[1] and c_pred[1] <= gt[3]):
+            return True
+        return False
+
+    tps = 0
+    fns = 0
+
+    for gt in gts:
+        # First check if any true positive exists
+        # If more than one exists, do not include it in next set of preds
+        add_tp = False
+        new_preds = []
+        for pred in preds:
+            if true_positive(gt, pred):
+                add_tp = True
+            else:
+                new_preds.append(pred)
+        preds = new_preds
+        if add_tp:
+            tps += 1
+        else:
+            fns += 1
+    return preds, tps, fns
+
+
+
+def calc_metric_single(gts, preds, threshold,):
+    '''
+        Returns fp, tp, tn, fn
+    '''
+    preds = list(filter(lambda x: x[0] >= threshold, preds))
+    preds = [pred[1:] for pred in preds]  # Remove the scores
+
+    if len(gts) == 0:
+        return len(preds), 0, 1 if len(preds) == 0 else 0, 0
+    preds, tps, fns = remove_true_positives(gts, preds)
+    # All remaining will have to fps
+    fps = len(preds)
+    return fps, tps, 0, fns
+
+
+def calc_metrics_at_thresh(im_dict, threshold):
+    '''
+        Returns fp, tp, tn, fn
+    '''
+    fps, tps, tns, fns = 0, 0, 0, 0
+    for key in im_dict:
+        fp,tp,tn,fn = calc_metric_single(im_dict[key]['gt'],
+                           im_dict[key]['preds'], threshold)
+        fps+=fp
+        tps+=tp
+        tns+=tn
+        fns+=fn
+
+    return fps, tps, tns, fns
+
+from joblib import Parallel, delayed
+
+def calc_metrics(inp):
+    im_dict, tr = inp
+    out = dict()
+    for t in tr:
+        fp, tp, tn, fn = calc_metrics_at_thresh(im_dict, t)
+        out[t] = [fp, tp, tn, fn]
+    return out
+    
+
+def calc_froc_from_dict(im_dict, fps_req = [0.025,0.05,0.1,0.15,0.2,0.3], save_to = None):
+
+    num_images = len(im_dict)
+
+    gap = 0.005
+    n = int(1/gap)
+    thresholds = [i * gap for i in range(n)]
+    fps = [0 for _ in range(n)]
+    tps = [0 for _ in range(n)]
+    tns = [0 for _ in range(n)]
+    fns = [0 for _ in range(n)]
+
+
+    for i,t in enumerate(thresholds):
+        fps[i], tps[i], tns[i], fns[i] = calc_metrics_at_thresh(im_dict, t)
+
+
+    # Now calculate the sensitivities
+    senses = []
+    for t,f in zip(tps, fns):
+        try: senses.append(t/(t+f))
+        except: senses.append(0.)
+
+    if save_to is not None:
+        f = open(save_to, 'w')
+        for fp,s in zip(fps, senses):
+            f.write(f'{fp/num_images} {s}\n')
+        f.close()
+
+    senses_req = []
+    for fp_req in fps_req:
+        for i,f in enumerate(fps):
+            if f/num_images < fp_req:
+                if fp_req == 0.1:
+                    print(fps[i], tps[i], tns[i], fns[i])
+                    prec = tps[i]/(tps[i] + fps[i])
+                    recall = tps[i]/(tps[i] + fns[i])
+                    f1 = 2*prec*recall/(prec+recall)
+                    spec = tns[i]/ (tns[i] + fps[i])
+                    print(f'Specificity: {spec}')
+                    print(f'Precision: {prec}')
+                    print(f'Recall: {recall}')
+                    print(f'F1: {f1}')
+                senses_req.append(senses[i-1])
+                break
+    return senses_req, fps_req
+
+
+
+
+def file_to_bbox(file_name):
+    try:
+        content = open(file_name, 'r').readlines()
+        st = 0
+        if len(content) == 0:
+            # Empty File Should Return []
+            return []
+        if content[0].split()[0].isalpha():
+            st = 1
+        return [[float(x) for x in line.split()[st:]] for line in content]
+    except FileNotFoundError:
+        print(f'No Corresponding Box Found for file {file_name}, using [] as preds')
+        return []
+    except Exception as e:
+        print('Some Error',e)
+        return []
+
+def generate_image_dict(preds_folder_name='preds_42',
+                        root_fol='/home/pranjal/densebreeast_datasets/AIIMS_C1',
+                        mal_path=None, ben_path=None, gt_path=None,
+                        mal_img_path = None, ben_img_path = None
+                        ):
+
+    mal_path = join(root_fol, mal_path) if mal_path else join(
+        root_fol, 'mal', preds_folder_name)
+    ben_path = join(root_fol, ben_path) if ben_path else join(
+        root_fol, 'ben', preds_folder_name)
+    mal_img_path = join(root_fol, mal_img_path) if mal_img_path else join(
+        root_fol, 'mal', 'images')
+    ben_img_path = join(root_fol, ben_img_path) if ben_img_path else join(
+        root_fol, 'ben', 'images')
+    gt_path = join(root_fol, gt_path) if gt_path else join(
+        root_fol, 'mal', 'gt')
+
+
+    '''
+        image_dict structure:
+            'image_name(without txt/png)' : {'gt' : [[...]], 'preds' : [[]]}
+    '''
+    image_dict = dict()
+
+    # GT Might be sightly different from images, therefore we will index gts based on
+    # the images folder instead.
+    for file in os.listdir(mal_img_path):
+        if not file.endswith('.png'):
+            continue
+        file = file[:-4] + '.txt'
+        file = join(gt_path, file)
+        key = os.path.split(file)[-1][:-4]
+        image_dict[key] = dict()
+        image_dict[key]['gt'] = file_to_bbox(file)
+        image_dict[key]['preds'] = []
+
+    for file in glob.glob(join(mal_path, '*.txt')):
+        key = os.path.split(file)[-1][:-4]
+        assert key in image_dict
+        image_dict[key]['preds'] = file_to_bbox(file)
+
+    for file in os.listdir(ben_img_path):
+        if not file.endswith('.png'):
+            continue
+
+        file = file[:-4] + '.txt'
+        file = join(ben_path, file)
+        key = os.path.split(file)[-1][:-4]
+        if key == 'Calc-Test_P_00353_LEFT_CC' or key == 'Calc-Training_P_00600_LEFT_CC': # Corrupt Files in Dataset
+            continue
+        if key in image_dict:
+            print(key)
+        # assert key not in image_dict
+        if key in image_dict:
+            print(f'Unexpected Error. {key} exists in multiple splits')
+            continue
+        image_dict[key] = dict()
+        image_dict[key]['preds'] = file_to_bbox(file)
+        image_dict[key]['gt'] = []
+    return image_dict
+
+
+def pretty_print_fps(senses,fps):
+    for s,f in zip(senses,fps):
+        print(f'Sensitivty at {f}: {s}')
+
+def get_froc_points(preds_image_folder, root_fol, fps_req = [0.025,0.05,0.1,0.15,0.2,0.3], save_to = None):
+    im_dict = generate_image_dict(preds_image_folder, root_fol = root_fol)
+    # print(im_dict)
+    print(len(im_dict))
+    senses, fps = calc_froc_from_dict(im_dict, fps_req, save_to = save_to)
+    return senses, fps
+
+if __name__ == '__main__':
+    seed = '42' if len(sys.argv)== 1 else sys.argv[1]
+
+    root_fol = '../bilateral_new/MammoDatasets/AIIMS_highres_reliable/test_2'
+
+    if len(sys.argv) <= 2:
+        save_to = None
+    else:
+        save_to = sys.argv[2]
+    senses, fps = get_froc_points(f'preds_{seed}',root_fol, save_to = save_to)
+
+    pretty_print_fps(senses, fps)

DenseMammogram/geenerate_aiims.py

@@ -0,0 +1,61 @@
+import os
+import torch
+from os.path import join
+from model_utils import generate_predictions, generate_predictions_bilateral
+from models import get_FRCNN_model, Bilateral_model
+from froc_by_pranjal import get_froc_points
+from auc_by_pranjal import get_auc_score
+
+####### PARAMETERS TO ADJUST #######
+exp_name = 'BILATERAL'
+OUT_FILE = 'aiims_full_test_results/bil_complete.txt'
+BILATERAL = True
+dataset_path = 'AIIMS_highres_reliable/test_2'
+####################################
+
+
+
+
+if os.path.split(OUT_FILE)[0]:
+    os.makedirs(os.path.split(OUT_FILE)[0], exist_ok=True)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+frcnn_model =  get_FRCNN_model().to(device)
+
+if BILATERAL:
+    model = Bilateral_model(frcnn_model).to(device)
+    MODEL_PATH = f'experiments/{exp_name}/bilateral_models/bilateral_model.pth'
+    model.load_state_dict(torch.load(MODEL_PATH))
+else:
+    model = frcnn_model
+    MODEL_PATH = f'experiments/{exp_name}/frcnn_models/frcnn_model.pth'
+    model.load_state_dict(torch.load(MODEL_PATH))
+
+
+test_path = join('../bilateral_new', 'MammoDatasets',dataset_path)
+
+
+def get_aiims_dict(test_path, corr_file):    
+    extract_file = lambda x: x[x.find('test_2/')+7:]
+    corr_dict = {extract_file(line.split()[0].replace('"','')):extract_file(line.split()[1].replace('"','')) for line in open(corr_file).readlines()}
+    corr_dict = {join(test_path,k):join(test_path,v) for k,v in corr_dict.items()}
+    return corr_dict
+
+if BILATERAL:
+    pred_dir = f'preds_bilateral_{exp_name}'
+    generate_predictions_bilateral(model,device,test_path, get_aiims_dict(test_path, '../bilateral_new/corr_lists/aiims_corr_list_with_val_full_test.txt'),'aiims',pred_dir)
+else:
+    pred_dir = f'preds_frcnn_{exp_name}'
+    generate_predictions(model, device, test_path, preds_folder = pred_dir)
+
+
+file = open(OUT_FILE, 'a')
+file.writelines(f'{exp_name}  FROC Score:\n')
+senses, fps = get_froc_points(pred_dir, root_fol= test_path, fps_req = [0.025,0.05,0.1,0.15,0.2,0.3,1.0,1.5])
+for s,f in zip(senses, fps):
+    print(f'Sensitivty at {f}: {s}')
+    file.writelines(f'Sensitivty at {f}: {s}\n')
+file.close()
+
+print('AUC Score:',get_auc_score(pred_dir, test_path, retAcc = True, acc_thresh = 1.))
+

DenseMammogram/geenerate_ddsm_preds.py

@@ -0,0 +1,61 @@
+import os
+import torch
+from os.path import join
+from model_utils import generate_predictions, generate_predictions_bilateral
+from models import get_FRCNN_model, Bilateral_model
+from froc_by_pranjal import get_froc_points
+from auc_by_pranjal import get_auc_score
+
+####### PARAMETERS TO ADJUST #######
+exp_name = 'frcnn_16'
+OUT_FILE = 'ddsm_results/ddsm_dset.txt'
+BILATERAL = False
+dataset_path = 'ddsm_data_no_proc_2100_nocrop/val'
+####################################
+
+
+
+
+if os.path.split(OUT_FILE)[0]:
+    os.makedirs(os.path.split(OUT_FILE)[0], exist_ok=True)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+frcnn_model =  get_FRCNN_model().to(device)
+
+if BILATERAL:
+    model = Bilateral_model(frcnn_model).to(device)
+    MODEL_PATH = f'experiments/{exp_name}/bilateral_models/bilateral_model.pth'
+    model.load_state_dict(torch.load(MODEL_PATH))
+else:
+    model = frcnn_model
+    MODEL_PATH = f'experiments/{exp_name}/frcnn_models/frcnn_model.pth'
+    model.load_state_dict(torch.load(MODEL_PATH))
+
+
+test_path = join('../bilateral_new', 'MammoDatasets',dataset_path)
+
+
+def get_ddsm_dict(test_path, corr_file):    
+    extract_file = lambda x: x[x.find('val/')+4:]
+    corr_dict = {extract_file(line.split()[0].replace('"','')):extract_file(line.split()[1].replace('"','')) for line in open(corr_file).readlines()}
+    corr_dict = {join(test_path,k):join(test_path,v) for k,v in corr_dict.items()}
+    return corr_dict
+
+if BILATERAL:
+    pred_dir = f'preds_bilateral_{exp_name}'
+    generate_predictions_bilateral(model,device,test_path, get_ddsm_dict(test_path, '../bilateral_new/corr_lists/ddsm_corr_list_with_val.txt'),'ddsm',pred_dir)
+else:
+    pred_dir = f'preds_frcnn_{exp_name}'
+    generate_predictions(model, device, test_path, preds_folder = pred_dir)
+
+
+file = open(OUT_FILE, 'a')
+file.writelines(f'{exp_name}  FROC Score:\n')
+senses, fps = get_froc_points(pred_dir, root_fol= test_path, fps_req = [0.025,0.05,0.1,0.15,0.2,0.3,1.0,1.5])
+for s,f in zip(senses, fps):
+    print(f'Sensitivty at {f}: {s}')
+    file.writelines(f'Sensitivty at {f}: {s}\n')
+file.close()
+
+print('AUC Score:',get_auc_score(pred_dir, test_path, retAcc = True, acc_thresh = 1.))
+

DenseMammogram/geenerate_inbreast_preds.py

@@ -0,0 +1,57 @@
+import os
+import torch
+from os.path import join
+from model_utils import generate_predictions, generate_predictions_bilateral
+from models import get_FRCNN_model, Bilateral_model
+from froc_by_pranjal import get_froc_points
+
+####### PARAMETERS TO ADJUST #######
+exp_name = 'AIIMS_C3'
+OUT_FILE = 'ib_results/c3_frcnn.txt'
+BILATERAL = False
+dataset_path = 'INBREAST_C3/test'
+####################################
+
+
+
+
+if os.path.split(OUT_FILE)[0]:
+    os.makedirs(os.path.split(OUT_FILE)[0], exist_ok=True)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+frcnn_model =  get_FRCNN_model().to(device)
+
+if BILATERAL:
+    model = Bilateral_model(frcnn_model).to(device)
+    MODEL_PATH = f'experiments/{exp_name}/bilateral_models/bilateral_model.pth'
+    model.load_state_dict(torch.load(MODEL_PATH))
+else:
+    model = frcnn_model
+    MODEL_PATH = f'experiments/{exp_name}/frcnn_models/frcnn_model.pth'
+    model.load_state_dict(torch.load(MODEL_PATH))
+
+
+test_path = join('../bilateral_new', 'MammoDatasets',dataset_path)
+
+
+def get_inbreast_dict(test_path, corr_file):    
+    extract_file = lambda x: x[x.find('test/')+5:]
+    corr_dict = {extract_file(line.split()[0]):extract_file(line.split()[1]) for line in open(corr_file).readlines()}
+    corr_dict = {join(test_path,k):join(test_path,v) for k,v in corr_dict.items()}
+    return corr_dict
+
+if BILATERAL:
+    pred_dir = f'preds_bilateral_{exp_name}'
+    generate_predictions_bilateral(model,device,test_path, get_inbreast_dict(test_path, '../bilateral_new/corr_lists/Inbreast_final_correspondence_list.txt'),'inbreast',pred_dir)
+else:
+    pred_dir = f'preds_frcnn_{exp_name}'
+    generate_predictions(model, device, test_path, preds_folder = pred_dir)
+
+
+file = open(OUT_FILE, 'a')
+file.writelines(f'{exp_name}  FROC Score:\n')
+senses, fps = get_froc_points(pred_dir, root_fol= test_path)
+for s,f in zip(senses, fps):
+    file.writelines(f'Sensitivty at {f}: {s}\n')
+file.close()
+

DenseMammogram/geenerate_irch.py

@@ -0,0 +1,62 @@
+import os
+import torch
+from os.path import join
+from model_utils import generate_predictions, generate_predictions_bilateral
+from models import get_FRCNN_model, Bilateral_model
+from froc_by_pranjal import get_froc_points
+from auc_by_pranjal import get_auc_score
+
+####### PARAMETERS TO ADJUST #######
+exp_name = 'BILATERAL'
+OUT_FILE = 'irchvalres/bil_final.txt'
+BILATERAL = True
+dataset_path = 'IRCHVal'
+####################################
+
+
+
+
+if os.path.split(OUT_FILE)[0]:
+    os.makedirs(os.path.split(OUT_FILE)[0], exist_ok=True)
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+frcnn_model =  get_FRCNN_model().to(device)
+
+if BILATERAL:
+    model = Bilateral_model(frcnn_model).to(device)
+    MODEL_PATH = f'experiments/{exp_name}/bilateral_models/bilateral_model.pth'
+    model.load_state_dict(torch.load(MODEL_PATH))
+else:
+    model = frcnn_model
+    MODEL_PATH = f'experiments/{exp_name}/frcnn_models/frcnn_model.pth'
+    model.load_state_dict(torch.load(MODEL_PATH))
+
+
+test_path = join('../bilateral_new', 'MammoDatasets',dataset_path)
+
+
+def get_aiims_dict(test_path, corr_file):    
+    extract_file = lambda x: x
+    corr_dict = {extract_file(line.split('" "')[0].strip().replace('"','')):extract_file(line.split('" "')[1].strip().replace('"','')) for line in open(corr_file).readlines()}
+    corr_dict = {join(test_path,k):join(test_path,v) for k,v in corr_dict.items()}
+    print(list(corr_dict.keys())[:20])
+    return corr_dict
+
+if BILATERAL:
+    pred_dir = f'preds_bilateral_{exp_name}'
+    generate_predictions_bilateral(model,device,test_path, get_aiims_dict(test_path, '../bilateral_new/corr_lists/irch_val.txt'),'irch',pred_dir)
+else:
+    pred_dir = f'preds_frcnn_{exp_name}'
+    generate_predictions(model, device, test_path, preds_folder = pred_dir)
+
+
+file = open(OUT_FILE, 'a')
+file.writelines(f'{exp_name}  FROC Score:\n')
+senses, fps = get_froc_points(pred_dir, root_fol= test_path, fps_req = [0.025,0.05,0.1,0.15,0.2,0.3,1.0,1.5])
+for s,f in zip(senses, fps):
+    print(f'Sensitivty at {f}: {s}')
+    file.writelines(f'Sensitivty at {f}: {s}\n')
+file.close()
+
+print('AUC Score:',get_auc_score(pred_dir, test_path, retAcc = True, acc_thresh = 1.))
+

DenseMammogram/merge_predictions.py

@@ -0,0 +1,152 @@
+import os
+import glob
+import torch
+from os.path import join
+import numpy as np
+from froc_by_pranjal import file_to_bbox, calc_froc_from_dict, pretty_print_fps
+import sys
+from ensemble_boxes import *
+import json
+import pickle
+
+
+
+get_file_id = lambda x: x.split('_')[1]
+get_acr_cat = lambda x: '0' if x not in acr_cat else acr_cat[x]
+cat_to_idx = {'a':1,'b':2,'c':3,'d':4}
+
+
+def get_image_dict(dataset_paths, labels = ['mal','ben'], allowed = [], USE_ACR = False, acr_cat = None, mp_dict = None):
+    image_dict = dict()
+    if allowed == []:
+        allowed = [i for i in range(len(dataset_paths))]
+    for label in labels:
+        images = list(set.intersection(*map(set, [os.listdir(dset.format(label)) for dset in dataset_paths])))
+        for image in images:
+            if USE_ACR:
+                acr = get_acr_cat(get_file_id(image))
+            # print(acr, image)
+            key = image[:-4]
+            gts = []
+            preds = []
+            for i,dset in enumerate(dataset_paths):
+                if i not in allowed:
+                    continue
+                if USE_ACR:
+                    if dset.find('AIIMS_C')!=-1:
+                        if acr == '0': continue
+                        if dset.find(f'AIIMS_C{cat_to_idx[acr]}') == -1:
+                            continue
+                        # Now choose dset to be the acr category one
+                        dset = dset.replace('/test',f'/test_{acr}')
+                # print('ds',dset)
+                pred_file = join(dset.format(label), key+'.txt')
+                gt_file = join(os.path.split(dset.format(label))[0],'gt', key+'.txt')
+                if label == 'mal':
+                    gts.append(file_to_bbox(gt_file))
+                else:
+                    gts.append([])
+
+                # TODO: Note this
+                flag = False
+                for mp in mp_dict:
+                    if dataset_paths[i].find(mp) != -1:
+                        preds.append(mp_dict[mp](file_to_bbox(pred_file)))
+                        flag = True
+                        break
+                if not flag:
+                    preds.append(file_to_bbox(pred_file))
+
+            # Ensure all gts are same
+            gt = gts[0]
+            for g in gts[1:]:
+                assert g == gt
+                gt = g
+            
+            # Flatten Preds
+            preds = [np.array(p) for p in preds]
+            preds = [np.array([[0.,0.,0.,0.,0.]]) if pred.shape==(0,) else pred for pred in preds]
+            preds = [np.vstack((p, np.zeros((100 - len(p), 5)))) for p in preds]
+            image_dict[key] = dict()
+            image_dict[key]['gt'] = gts[0]
+            image_dict[key]['preds'] = preds
+    return image_dict
+
+
+def apply_merge(image_dict, METHOD = 'wbf', weights = None, conf_type = None):
+    FACTOR = 5000
+    fusion_func = weighted_boxes_fusion if METHOD == 'wbf' else non_maximum_weighted
+    for key in image_dict:
+        preds = np.array(image_dict[key]['preds'])
+        if len(preds) != 0:
+            boxes_list = [pred[:,1:]/FACTOR for pred in preds]
+            scores_list = [pred[:,0] for pred in preds]
+            labels = [[0. for _ in range(len(p))] for p in preds]
+            if weights is None:
+                weights = [1 for _ in range(len(preds))]
+            if METHOD == 'wbf' and conf_type is not None:
+                boxes,scores,_ = fusion_func(boxes_list, scores_list, labels, weights = weights,iou_thr  = 0.5, conf_type = conf_type)
+            else:
+                boxes,scores,_ = fusion_func(boxes_list, scores_list, labels, weights = weights,iou_thr  = 0.5,)
+            preds_t = [[scores[i],FACTOR*boxes[i][0],FACTOR*boxes[i][1],FACTOR*boxes[i][2],FACTOR*boxes[i][3]] for i in range(len(boxes))]
+            image_dict[key]['preds'] = preds_t
+    return image_dict
+
+def manipulate_preds(preds):
+    return preds
+
+
+
+def manipulate_preds_4(preds):
+    return preds
+
+tot = 0
+def manipulate_preds_t1(preds): #return manipulate_preds(preds)
+    preds  = list(filter(lambda x: x[0]>0.6,preds))
+
+    return preds
+
+def manipulate_preds_t2(preds): return manipulate_preds_t1(preds)
+
+
+if __name__ == '__main__':
+    USE_ACR = False
+    dataset_paths = [
+        'MammoDatasets/AIIMS_C1/test/{0}/preds_frcnn_AIIMS_C1',
+        'MammoDatasets/AIIMS_C2/test/{0}/preds_frcnn_AIIMS_C2',
+        'MammoDatasets/AIIMS_C3/test/{0}/preds_frcnn_AIIMS_C3',
+        'MammoDatasets/AIIMS_C4/test/{0}/preds_frcnn_AIIMS_C4',
+        'MammoDatasets/AIIMS_highres_reliable/test/{0}/preds_bilateral_BILATERAL',
+        'MammoDatasets/AIIMS_highres_reliable/test/{0}/preds_frcnn_16',
+    ]
+
+
+    st = int(sys.argv[1])
+    end = len(dataset_paths) - int(sys.argv[2])
+    allowed = [i for i in range(st,end)]
+    allowed = [0,1,2,3,4,5]
+
+    OUT_FILE = 'contrast_frcnn.txt'
+    if OUT_FILE is not None:
+        fol = os.path.split(OUT_FILE)[0]
+        if fol != '':
+            os.makedirs(fol, exist_ok=True)
+    
+    acr_cat = json.load(open('aiims_categories.json','r'))
+    print(allowed)
+
+    mp_dict = {
+        'preds_frcnn_AIIMS_C3': manipulate_preds,
+        'preds_frcnn_AIIMS_C4': manipulate_preds_4,
+        'AIIMS_T2': manipulate_preds_t2,
+        'AIIMS_T1': manipulate_preds_t1,
+    }
+
+    image_dict = get_image_dict(dataset_paths, allowed = allowed, USE_ACR = USE_ACR, acr_cat = acr_cat, mp_dict = mp_dict)
+    
+    image_dict = apply_merge(image_dict, METHOD = 'nms') # or wbf
+
+    if OUT_FILE:
+        pickle.dump(image_dict, open(OUT_FILE.replace('.txt','.pkl'),'wb'))
+        senses, fps = calc_froc_from_dict(image_dict, fps_req = [0.025,0.05,0.1,0.15,0.2,0.3,1.],save_to=OUT_FILE)
+    pretty_print_fps(senses, fps)

DenseMammogram/model_utils.py

@@ -0,0 +1,83 @@
+import os
+import torchvision.transforms as T
+import cv2
+from tqdm import tqdm
+import detection.transforms as transforms
+from dataloaders import get_direction
+
+def generate_predictions_bilateral(model,device,testpath_,cor_dict,dset='aiims',preds_folder='preds_new'):
+    transform = T.Compose([T.ToPILImage(),T.ToTensor()])
+    model.eval()
+    for label in ['mal','ben']:
+        testpath = os.path.join(testpath_,label)
+        # testpath = os.path.join(dataset_path,'Training', 'train',label)
+        testimg = os.path.join(testpath, 'images')
+
+        #preds_folder  = 'preds_new'
+        os.makedirs(os.path.join(testpath, preds_folder),exist_ok=True)
+
+        if not os.path.exists(os.path.join(testpath,preds_folder)):
+            os.makedirs(os.path.join(testpath+preds_folder),exist_ok = True)
+
+        for file in tqdm(os.listdir(testimg)):
+            img1 = cv2.imread(os.path.join(testimg,file))
+            img1 = transform(img1)
+            # if False:
+            if(os.path.join(testimg,file) in cor_dict and os.path.isfile(cor_dict[os.path.join(testimg,file)])):
+                print('Using Bilateral')
+                img2 = cv2.imread(cor_dict[os.path.join(testimg,file)])
+                img2 = transform(img2)
+                if(get_direction(dset,file)==1):
+                    img1,_ = transforms.RandomHorizontalFlip(1.0)(img1)
+                    
+                    images = [img1.to(device),img2.to(device)]
+                    output = model([images])[0]
+                    img1,output = transforms.RandomHorizontalFlip(1.0)(img1,output)
+                else:
+                    img2,_ = transforms.RandomHorizontalFlip(1.0)(img2)
+                    
+                    images = [img1.to(device),img2.to(device)]
+                    output = model([images])[0]
+            else:
+                print('Using FRCNN')
+                output = model.frcnn([img1.to(device)])[0]
+            #output = model.frcnn([img1.to(device)])[0]
+            boxes = output['boxes']
+            scores = output['scores']
+            labels = output['labels']
+            f = open(os.path.join(testpath,preds_folder,file[:-4]+'.txt'),'w')
+            for i in range(len(boxes)):
+                box = boxes[i].detach().cpu().numpy()
+                #f.write('{} {} {} {} {} {}\n'.format(scores[i].item(),labels[i].item(),box[0],box[1],box[2],box[3]))
+                f.write('{} {} {} {} {}\n'.format(scores[i].item(),box[0],box[1],box[2],box[3]))
+
+
+def generate_predictions(model,device,testpath_,preds_folder='preds_frcnn'):
+    transform = T.Compose([T.ToPILImage(),T.ToTensor()])
+    model.eval()
+    for label in ['mal','ben']:
+        testpath = os.path.join(testpath_,label)
+        # testpath = os.path.join(dataset_path,'Training', 'train',label)
+        testimg = os.path.join(testpath, 'images')
+
+        #preds_folder  = 'preds_new'
+        os.makedirs(os.path.join(testpath, preds_folder),exist_ok=True)
+
+        if not os.path.exists(os.path.join(testpath,preds_folder)):
+            os.makedirs(os.path.join(testpath+preds_folder),exist_ok = True)
+
+        for file in tqdm(os.listdir(testimg)):
+            im = cv2.imread(os.path.join(testimg,file))
+            if file == 'Mass-Training_P_00444_LEFT_CC.png':
+                print('Test this')
+                continue
+            im = transform(im)
+
+            output = model([im.to(device)])[0]
+            boxes = output['boxes'] #/ FAC
+            scores = output['scores']
+            labels = output['labels']
+            f = open(os.path.join(testpath,preds_folder,file[:-4]+'.txt'),'w')
+            for i in range(len(boxes)):
+                box = boxes[i].detach().cpu().numpy()
+                f.write('{} {} {} {} {}\n'.format(scores[i].item(),box[0],box[1],box[2],box[3]))

DenseMammogram/models.py

@@ -0,0 +1,201 @@
+from typing import List, OrderedDict, Tuple
+import warnings
+import numpy as np
+import pandas as pd
+import cv2
+import os
+from torch.nn.modules.conv import Conv2d
+from torch.utils.data.dataset import ConcatDataset
+from tqdm import tqdm
+import argparse
+from torch.utils.data import Dataset,DataLoader
+import torch
+import torch.nn as nn
+from torchvision import models
+import detection.transforms as transforms
+import torchvision.transforms as T
+import detection.utils as utils
+import torch.nn.functional as F
+import shutil
+import json
+from detection.engine import train_one_epoch, evaluate
+from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
+import torch.multiprocessing
+import copy
+from torchvision.ops import MultiScaleRoIAlign
+from torchvision.models.detection.roi_heads import RoIHeads
+
+
+
+
+# First we will create the FRCNN model
+def get_FRCNN_model(num_classes=1):
+    model = models.detection.fasterrcnn_resnet50_fpn(pretrained=True,trainable_backbone_layers=3,min_size=1800,max_size=3600,image_std=(1.0,1.0,1.0),box_score_thresh=0.001)
+    # get number of input features for the classifier
+    in_features = model.roi_heads.box_predictor.cls_score.in_features
+    # replace the pre-trained head with a new one
+    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes+1) 
+    return model
+
+# Some utility heads for Bilateral Model
+
+class RoIpool(nn.Module):
+
+    def __init__(self,pool):
+        super().__init__()
+        self.box_roi_pool1 = copy.deepcopy(pool)
+        self.box_roi_pool2 = copy.deepcopy(pool)
+        
+
+    def forward(self,features,proposals,image_shapes):
+        x = self.box_roi_pool1(features[0],proposals,image_shapes)
+        y = self.box_roi_pool2(features[1],proposals,image_shapes)
+        z = torch.cat((x,y),dim=1)
+        return z
+
+class TwoMLPHead(nn.Module):
+    """
+    Standard heads for FPN-based models
+    Args:
+        in_channels (int): number of input channels
+        representation_size (int): size of the intermediate representation
+    """
+
+    def __init__(self, in_channels=None, representation_size=None):
+        super().__init__()
+
+        self.fc6 = nn.Linear(in_channels, representation_size)
+        self.fc7 = nn.Linear(representation_size, representation_size)
+
+    def forward(self, x):
+        x = x.flatten(start_dim=1)
+
+        x = F.relu(self.fc6(x))
+        x = F.relu(self.fc7(x))
+        return x
+
+# Next the bilateral model
+
+class Bilateral_model(nn.Module):
+
+    def __init__(self,frcnn_model):
+        super().__init__()
+        self.frcnn = frcnn_model
+        self.transform = copy.deepcopy(frcnn_model.transform)
+        self.backbone1 = copy.deepcopy(frcnn_model.backbone)
+        self.backbone2 = copy.deepcopy(frcnn_model.backbone)
+        self.rpn = copy.deepcopy(frcnn_model.rpn)
+        for param in self.rpn.parameters():
+            param.requires_grad = False
+        for param in self.backbone1.parameters():
+            param.requires_grad = False
+        for param in self.backbone2.parameters():
+           param.requires_grad = False
+        box_roi_pool = RoIpool(frcnn_model.roi_heads.box_roi_pool)
+        box_head = TwoMLPHead(512*7*7,1024)
+        box_predictor = copy.deepcopy(frcnn_model.roi_heads.box_predictor)
+        box_score_thresh=0.001
+        box_nms_thresh=0.5
+        box_detections_per_img=100
+        box_fg_iou_thresh=0.5
+        box_bg_iou_thresh=0.5
+        box_batch_size_per_image=512
+        box_positive_fraction=0.25
+        bbox_reg_weights=None
+        self.roi_heads =  RoIHeads(
+            # Box
+            box_roi_pool,
+            box_head,
+            box_predictor,
+            box_fg_iou_thresh,
+            box_bg_iou_thresh,
+            box_batch_size_per_image,
+            box_positive_fraction,
+            bbox_reg_weights,
+            box_score_thresh,
+            box_nms_thresh,
+            box_detections_per_img,
+        )
+    
+    @torch.jit.unused
+    def eager_outputs(self, losses, detections):
+        if self.training:
+            return losses
+
+        return detections
+    
+    
+    def forward(self, images, targets=None):
+        """
+        Args:
+            images (list[Tensor(tuples)]): images to be processed
+            targets (list[Dict[str, Tensor]]): ground-truth boxes present in the image (optional)
+        Returns:
+            result (list[BoxList] or dict[Tensor]): the output from the model.
+                During training, it returns a dict[Tensor] which contains the losses.
+                During testing, it returns list[BoxList] contains additional fields
+                like `scores`, `labels` and `mask` (for Mask R-CNN models).
+        """
+        if self.training and targets is None:
+            raise ValueError("In training mode, targets should be passed")
+        if self.training:
+            assert targets is not None
+            for target in targets:
+                boxes = target["boxes"]
+                if isinstance(boxes, torch.Tensor):
+                    if len(boxes.shape) != 2 or boxes.shape[-1] != 4:
+                        raise ValueError(f"Expected target boxes to be a tensor of shape [N, 4], got {boxes.shape}.")
+                else:
+                    raise ValueError(f"Expected target boxes to be of type Tensor, got {type(boxes)}.")
+
+        original_image_sizes: List[Tuple[int, int]] = []
+        for img in images:
+            val = img[0].shape[-2:]
+            assert len(val) == 2
+            original_image_sizes.append((val[0], val[1]))
+        images1 = [img[0] for img in images]
+        images2 = [img[1] for img in images]
+        targets2 = copy.deepcopy(targets)
+        #print(images1.shape)
+        #print(images2.shape)
+        images1, targets = self.transform(images1, targets)
+        images2, targets2 = self.transform(images2, targets2)
+        
+        # Check for degenerate boxes
+        # TODO: Move this to a function
+        if targets is not None:
+            for target_idx, target in enumerate(targets):
+                boxes = target["boxes"]
+                degenerate_boxes = boxes[:, 2:] <= boxes[:, :2]
+                if degenerate_boxes.any():
+                    # print the first degenerate box
+                    bb_idx = torch.where(degenerate_boxes.any(dim=1))[0][0]
+                    degen_bb: List[float] = boxes[bb_idx].tolist()
+                    raise ValueError(
+                        "All bounding boxes should have positive height and width."
+                        f" Found invalid box {degen_bb} for target at index {target_idx}."
+                    )
+
+        features1 = self.backbone1(images1.tensors)
+        features2 = self.backbone2(images2.tensors)
+        #print(self.backbone1.out_channels)
+        if isinstance(features1, torch.Tensor):
+            features1 = OrderedDict([("0", features1)])
+        if isinstance(features2, torch.Tensor):
+            features2 = OrderedDict([("0", features2)])
+        proposals, proposal_losses = self.rpn(images1, features1, targets)
+        features = {0:features1,1:features2}
+        detections, detector_losses = self.roi_heads(features, proposals, images1.image_sizes, targets)
+        detections = self.transform.postprocess(detections, images1.image_sizes, original_image_sizes)  # type: ignore[operator]
+
+        losses = {}
+        losses.update(detector_losses)
+        losses.update(proposal_losses)
+
+        if torch.jit.is_scripting():
+            if not self._has_warned:
+                warnings.warn("RCNN always returns a (Losses, Detections) tuple in scripting")
+                self._has_warned = True
+            return losses, detections
+        else:
+            return self.eager_outputs(losses, detections)

DenseMammogram/plot_froc.py

@@ -0,0 +1,43 @@
+import matplotlib.pyplot as plt
+import numpy as np
+
+####### PARAMETERS TO ADJUST #######
+
+# Specify the files generated from merge_nms and plot corresponding graphs
+base_fol = 'normal_test'
+input_files = {
+    f'thresh_uni.txt' : 'Thresh + Uni',
+    f'thresh_nouni.txt' : 'Thresh + NoUni',
+}
+save_file = 'uni_vs_nouni.png'
+# TITLE = 'Thresh + Contrast + Bilateral vs Contrast + Bilateral FROC Comparison (Normal Test)'
+TITLE = 'Uni vs NoUni FROC Comparison (Normal Test)'
+
+SHOW = False
+CLIP_FPI = 1.2
+MIN_CLIP_FPI = 0.0
+####################################
+
+def plot_froc(input_files, save_file, TITLE = 'FRCNN vs BILATERAL FROC', SHOW = False, CLIP_FPI = 1.2):
+    for file in input_files:
+        lines = open(file).readlines()
+        x = np.array([float(line.split()[0]) for line in lines])
+        y = np.array([float(line.split()[1]) for line in lines])
+        y = y[x<CLIP_FPI]
+        x = x[x<CLIP_FPI]
+        y = y[MIN_CLIP_FPI<x]
+        x = x[MIN_CLIP_FPI<x]
+        plt.plot(x, y, label = input_files[file])
+        plt.legend()
+
+    plt.title(TITLE)
+    plt.xlabel('Average False Positive Per Image')
+    plt.ylabel('Sensetivity')
+
+    if SHOW:
+        plt.show()
+    plt.savefig(save_file)
+    plt.clf()
+
+if __name__ == '__main__':
+    plot_froc(input_files, save_file, TITLE, SHOW, CLIP_FPI)

DenseMammogram/requirements.txt

@@ -0,0 +1,11 @@
+torch==1.10.2
+tqdm==4.62.3
+torchvision==0.11.3
+scipy==1.7.3
+scikit-learn==1.0.2
+PyYAML==6.0
+Pillow==8.4.0
+pandas==1.4.0
+matplotlib==3.5.1
+numpy
+easydict==1.9

DenseMammogram/train_bilateral.py

@@ -0,0 +1,47 @@
+import torch
+import math
+from advanced_logger import LogPriority
+from dataloaders import get_bilateral_dataloaders
+from models import get_FRCNN_model, Bilateral_model
+from detection.engine import evaluate_loss, train_one_epoch_simplified
+
+def main(cfg, experimenter):
+
+    LR = cfg['LR']
+    WEIGHT_DECAY = cfg['WEIGHT_DECAY']
+    NUM_EPOCHS = cfg['NUM_EPOCHS']
+    BATCH_SIZE = cfg['BATCH_SIZE']
+
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    frcnn_model =  get_FRCNN_model().to(device)
+    frcnn_model.load_state_dict(torch.load(cfg['FRCNN_MODEL_PATH']))
+
+    model = Bilateral_model(frcnn_model).to(device)
+
+    train_loader, val_loader = get_bilateral_dataloaders(experimenter.config, batch_size = BATCH_SIZE, data_dir = experimenter.config['DATA_DIR'])
+    
+    if cfg["OPTIM"] == "SGD":
+        optimizer = torch.optim.SGD(filter(lambda p: p.requires_grad , model.roi_heads.parameters()),lr=LR,momentum=0.9,weight_decay=WEIGHT_DECAY)
+    elif cfg["OPTIM"] == "ADAM":
+        optimizer = torch.optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr = LR, weight_decay = WEIGHT_DECAY)
+    elif cfg["OPTIM"] == "ADAGRAD":
+        optimizer = torch.optim.Adagrad(filter(lambda p: p.requires_grad, model.roi_heads.parameters()), lr = LR, weight_decay = WEIGHT_DECAY)
+    for epoch in range(NUM_EPOCHS):
+        experimenter.start_epoch()
+        train_one_epoch_simplified(model, optimizer, train_loader, device, epoch, experimenter = experimenter,optimizer_backbone=None)
+        loss = evaluate_loss(model, device, val_loader, experimenter = experimenter)
+        experimenter.log('Validation Loss: {}'.format(loss), priority = LogPriority.MEDIUM)
+        
+        experimenter.end_epoch(loss, model, device)
+        experimenter.save_model(model)
+    experimenter.generate_predictions(model, device)
+
+
+if __name__ == '__main__':
+    from experimenter import Experimenter
+    import os
+    os.environ['CUDA_VISIBLE_DEVICES'] = '4'
+    cfg_file = 'configs/default.cfg'
+    experimenter = Experimenter(cfg_file)
+    main(experimenter.config['BILATERAL'], experimenter)

DenseMammogram/train_frcnn.py

@@ -0,0 +1,34 @@
+import torch
+import math
+from advanced_logger import LogPriority
+from dataloaders import get_FRCNN_dataloaders
+from models import get_FRCNN_model
+from detection.engine import evaluate_loss, evaluate_simplified, train_one_epoch_simplified, evaluate_simplified
+
+def main(cfg, experimenter):
+
+    LR = cfg['LR']
+    WEIGHT_DECAY = cfg['WEIGHT_DECAY']
+    NUM_EPOCHS = cfg['NUM_EPOCHS']
+    BATCH_SIZE = cfg['BATCH_SIZE']
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model =  get_FRCNN_model().to(device)
+    train_loader, val_loader = get_FRCNN_dataloaders(experimenter.config, batch_size=BATCH_SIZE, data_dir = experimenter.config['DATA_DIR'])
+    optimizer = torch.optim.SGD(model.parameters(),lr=LR,momentum=0.9,weight_decay=WEIGHT_DECAY)
+
+    for epoch in range(NUM_EPOCHS):
+        experimenter.start_epoch()
+        train_one_epoch_simplified(model, optimizer, train_loader, device, epoch, experimenter = experimenter)
+        evaluate_simplified(model, val_loader, device=device, experimenter = experimenter)
+        loss = evaluate_loss(model, device, val_loader, experimenter = experimenter)
+        experimenter.log('Validation Loss: {}'.format(loss), priority = LogPriority.MEDIUM)
+        experimenter.end_epoch(loss, model = model, device = device)
+        experimenter.save_model(model)
+    experimenter.generate_predictions(model, device)
+
+if __name__ == '__main__':
+    from experimenter import Experimenter
+    cfg_file = 'configs/AIIMS_C1.cfg'
+    experimenter = Experimenter(cfg_file)
+    main(experimenter.config['FRCNN'], experimenter)

DenseMammogram/utils.py

@@ -0,0 +1,41 @@
+import shutil
+import os
+from os.path import join
+
+
+class AverageMeter:
+    """Computes and stores the average and current value"""
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val, n=1):
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count
+
+def create_backup(folders = None, files = None, backup_dir = 'experiments'):
+    if folders is None:
+        folders = ['.', 'corr_lists','detection']
+    if files is None:
+        files = ['.py', '.txt', '.json','.cfg']
+
+    for folder in folders:
+        if not os.path.isdir(folder):
+            continue
+        for file in os.listdir(folder):
+            if file.endswith(tuple(files)):
+                if folder != '.':
+                    src = join(folder, file)
+                    dest = join(backup_dir, folder, file)
+                else:
+                    src = file
+                    dest = join(backup_dir, file)
+                os.makedirs(os.path.split(dest)[0], exist_ok=True)
+                shutil.copy(src, dest)

app.py

@@ -0,0 +1,117 @@
+import gradio as gr
+from model import predict
+import cv2
+
+
+with gr.Blocks() as demo:
+    with gr.Column():
+        title = "<h1 style='margin-bottom: -10px; text-align: center'>Deep Learning for Detection of iso-dense, obscure masses in mammographically dense breasts</h1>"
+        # gr.HTML(title)
+        gr.Markdown(
+            "<h1 style='text-align: center; margin-bottom: 1rem'>"
+            + title
+            + "</h1>"
+        )
+
+        description = "<p style='font-size: 14px; margin: 5px; font-weight: w300; text-align: center'> <a href='' style='text-decoration:none' target='_blank'>Krithika Rangarajan<sup>*</sup>, </a> <a href='https://github.com/Pranjal2041' style='text-decoration:none' target='_blank'>Pranjal Aggarwal<sup>*</sup>, </a> <a href='' style='text-decoration:none' target='_blank'>Dhruv Kumar Gupta, </a> <a href='' style='text-decoration:none' target='_blank'>Rohan Dhanakshirur, </a> <a href='' style='text-decoration:none' target='_blank'>Akhil Baby, </a> <a href='' style='text-decoration:none' target='_blank'>Chandan Pal, </a> <a href='' style='text-decoration:none' target='_blank'>Arun Kumar Gupta, </a> <a href='' style='text-decoration:none' target='_blank'>Smriti Hari, </a> <a href='' style='text-decoration:none' target='_blank'>Subhashis Banerjee, </a> <a href='' style='text-decoration:none' target='_blank'>Chetan Arora, </a> </p>" \
+            + "<p style='font-size: 16px; margin: 5px; font-weight: w600; text-align: center'>  <a href='https://link.springer.com/article/10.1007/s00330-023-09717-7' target='_blank'>Publication</a> | <a href='https://github.com/Pranjal2041/DenseMammogram' target='_blank'>Website</a> | <a href='https://github.com/Pranjal2041/DenseMammogram' target='_blank'>Github Repo</a></p>" \
+            + "<p style='text-align: center; margin: 5px; font-size: 14px; font-weight: w300;'>  \
+                Deep learning suffers from some problems similar to human radiologists, such as poor sensitivity to detection of isodense, obscure masses or cancers in dense breasts. Traditional radiology teaching can be incorporated into the deep learning approach to tackle these problems in the network. Our method suggests collaborative network design, and incorporates core radiology principles resulting in SOTA results. You can use this demo to run inference by providing bilateral mammogram images. To get started, you can try one of the preset examples. \
+                </p>" \
+            + "<p style='text-align: center; font-size: 14px; margin: 5px; font-weight: w300;'> [Note: Inference on CPU may take upto 2 minutes. On a GPU, inference time is approximately 1s.]</p>"
+        # gr.HTML(description)
+        gr.Markdown(description)
+
+        # head_html = gr.HTML('''
+        #     <h1>
+        #          Deep Learning for Detection of iso-dense, obscure masses in mammographically dense breasts
+        #     </h1>
+        #     <p style='text-align: center;'>
+        #         Give bilateral mammograms(both left and right sides), and let our model find the cancers!
+        #     </p>
+            
+        #     <p style='text-align: center;'>
+        #         This is an official demo for our paper: 
+        #          `Deep Learning for Detection of iso-dense, obscure masses in mammographically dense breasts`.
+        #          Check out the paper and code for details! 
+        #     </p>
+        # ''')
+
+        # gr.Markdown(
+        # """
+        #     [![report](https://img.shields.io/badge/arxiv-report-red)](https://arxiv.org/abs/) [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/) 
+        # """)
+
+        def generate_preds(img1, img2):
+            print(img1, img2)
+            print(img1, img2)
+            img_out1 = predict(img1, img2)
+            if img_out1.shape[1] < img_out1.shape[2]:
+                ratio = img_out1.shape[2] / 800
+            else:
+                ratio = img_out1.shape[1] / 800
+            img_out1 = cv2.resize(img_out1, (0,0), fx=1 / ratio, fy=1 / ratio) 
+            img_out2 = predict(img2, img1, baseIsLeft = False) 
+            if img_out2.shape[1] < img_out2.shape[2]:
+                ratio = img_out2.shape[2] / 800
+            else:
+                ratio = img_out2.shape[1] / 800
+            img_out2 = cv2.resize(img_out2, (0,0), fx= 1 / ratio, fy= 1 / ratio) 
+
+            cv2.imwrite('img_out1.jpg', img_out1)
+            cv2.imwrite('img_out2.jpg', img_out2)
+
+
+            return 'img_out1.jpg', 'img_out2.jpg'
+
+        with gr.Column():
+            with gr.Row(variant = 'panel'):
+                
+                with gr.Column(variant = 'panel'):
+                    img1 = gr.Image(type="filepath", label="Left Image" )
+                    img2 = gr.Image(type="filepath", label="Right Image")
+                    # with gr.Row():
+                    #     sub_btn = gr.Button("Predict!", variant="primary")
+
+                with gr.Column(variant = 'panel'):
+                    # img_out1 = gr.inputs.Image(type="file", label="Output Left Image")
+                    # img_out2 = gr.inputs.Image(type="file", label="Output for Right Image")
+                    img_out1 = gr.Image(type="filepath", label="Output for Left Image", shape = None)
+                    img_out1.style(height=250 * 2)
+
+                with gr.Column(variant = 'panel'):
+                    img_out2 = gr.Image(type="filepath", label="Output for Right Image", shape = None)
+                    img_out2.style(height=250 * 2)
+
+            with gr.Row():
+                sub_btn = gr.Button("Predict!", variant="primary")
+
+        gr.Examples([[f'sample_images/img{idx}_l.jpg', f'sample_images/img{idx}_r.jpg'] for idx in range(1,6)], inputs = [img1, img2])
+
+        sub_btn.click(fn = lambda x,y: generate_preds(x,y), inputs = [img1, img2], outputs = [img_out1, img_out2])
+
+        # sub_btn.click(fn = lambda x: gr.update(visible = True), inputs = [sub_btn], outputs = [img_out1, img_out2])
+
+        # gr.Examples(
+
+        # )    
+
+
+        # interface.render()
+        # Object Detection Interface
+
+# def generate_predictions(img1, img2):
+#     return img1
+
+# interface = gr.Interface(
+#     fn=generate_predictions,
+#     inputs=[gr.inputs.Image(type="pil", label="Left Image"), gr.inputs.Image(type="pil", label="Right Image")],
+#     outputs=[gr.outputs.Image(type="pil", label="Output Image")],
+#     title="Object Detection",
+#     description="This model is trained on DenseMammogram dataset. It can detect objects in images. Try it out!",
+#     allow_flagging = False
+# ).launch(share = True, show_api=False)
+
+
+if __name__ == '__main__':
+    demo.launch(share = True, show_api=False)

model.py

@@ -0,0 +1,57 @@
+import sys
+sys.path.append('DenseMammogram')
+
+import torch
+
+from models import get_FRCNN_model, Bilateral_model
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+frcnn_model = get_FRCNN_model().to(device)
+bilat_model = Bilateral_model(frcnn_model).to(device)
+
+FRCNN_PATH = 'pretrained_models/frcnn/frcnn_models/frcnn_model.pth'
+BILAR_PATH = 'pretrained_models/BILATERAL/bilateral_models/bilateral_model.pth'
+
+frcnn_model.load_state_dict(torch.load(FRCNN_PATH, map_location=device))
+bilat_model.load_state_dict(torch.load(BILAR_PATH, map_location=device))
+
+import os
+import torchvision.transforms as T
+import cv2
+from tqdm import tqdm
+import detection.transforms as transforms
+from dataloaders import get_direction
+
+def predict(left_file, right_file, threshold = 0.80, baseIsLeft = True):
+    model = bilat_model
+    with torch.no_grad():
+        transform = T.Compose([T.ToPILImage(),T.ToTensor()])
+        model.eval()
+        # First is left, then right
+        img1 = cv2.imread(left_file)
+        img1 = transform(img1)
+        img2 = cv2.imread(right_file)
+        img2 = transform(img2)
+
+        if baseIsLeft:
+            img1,_ = transforms.RandomHorizontalFlip(1.0)(img1)
+        else:
+            img2,_ = transforms.RandomHorizontalFlip(1.0)(img2)
+
+
+        images = [img1.to(device),img2.to(device)]
+        output = model([images])[0]
+        if baseIsLeft:
+            img1,output = transforms.RandomHorizontalFlip(1.0)(img1,output)
+        
+        image =  cv2.imread(left_file)
+        for b,s,l in zip(output['boxes'], output['scores'], output['labels']):
+            # Convert img1 tensor to numpy array
+            if l == 1 and s > threshold:
+                # Draw the bounding boxes
+                b = b.detach().cpu().numpy().astype(int)   
+                # return image, b           
+                cv2.rectangle(image, (b[0], b[1]), (b[2], b[3]), (0, 255, 0), 2)
+                # Print the % probability just above the box
+                cv2.putText(image, 'Cancer: '+str(round(round(s.item(), 2) * 100, 1)) + '%', (b[0], b[1] - 40), cv2.FONT_HERSHEY_SIMPLEX, 3.6, (36,255,12), 6)
+        return image 

pretrained_models/AIIMS_C1/frcnn_models/frcnn_model.pth

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

pretrained_models/AIIMS_C2/frcnn_models/frcnn_model.pth

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

pretrained_models/AIIMS_C3/frcnn_models/frcnn_model.pth

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

pretrained_models/AIIMS_C4/frcnn_models/frcnn_model.pth

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

pretrained_models/AIIMS_T1/frcnn_models/frcnn_model.pth

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

pretrained_models/AIIMS_T2/frcnn_models/frcnn_model.pth

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

pretrained_models/BILATERAL/bilateral_models/bilateral_model.pth

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

pretrained_models/frcnn/frcnn_models/frcnn_model.pth

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

requirements.txt

@@ -0,0 +1,12 @@
+gradio
+torch==1.10.2
+tqdm==4.62.3
+torchvision==0.11.3
+scipy==1.7.3
+scikit-learn==1.0.2
+PyYAML==6.0
+Pillow==8.4.0
+pandas==1.4.0
+matplotlib==3.5.1
+numpy
+easydict==1.9