├── .gitignore ├── ClsKL ├── README.md ├── __init__.py ├── test │ ├── test_kl.py │ └── test_kl.sh ├── train │ ├── __init__.py │ ├── train_kl.py │ └── train_kl.sh ├── utils │ ├── __init__.py │ ├── eval_eng.py │ ├── eval_util.py │ ├── grad_cam.py │ ├── knee_sets.py │ ├── layer_util.py │ ├── loader.py │ ├── loss_util.py │ ├── model.py │ ├── torch_util.py │ └── train_eng.py └── vis │ ├── conf_matrix.py │ ├── knee_tsne.py │ └── plot_util.py ├── DetJoint ├── README.md ├── eval │ └── model_eval.py ├── preprocess │ ├── extract_yolo_det.py │ ├── save_det_h5.py │ └── split_train_val_test.py ├── test │ ├── select_overlay.py │ ├── test_knee.py │ └── test_knee.sh ├── train │ ├── cal_mean_var.py │ ├── cluster_anchor_box.py │ ├── train_knee.py │ └── train_knee.sh └── yolo_v2 │ ├── __init__.py │ ├── cfgs │ ├── __init__.py │ └── config_knee.py │ ├── darknet.py │ ├── datasets │ ├── __init__.py │ └── knee.py │ ├── knee_utils.py │ ├── loss_utils.py │ ├── make.sh │ ├── proj_utils │ ├── __init__.py │ ├── generic_utils.py │ ├── local_utils.py │ ├── model_utils.py │ ├── plot_utils.py │ └── torch_utils.py │ ├── test_yolo.py │ ├── train_yolo.py │ └── utils │ ├── __init__.py │ ├── bbox.pyx │ ├── build.py │ ├── im_transform.py │ ├── network.py │ ├── nms │ ├── .gitignore │ ├── __init__.py │ ├── cpu_nms.pyx │ ├── gpu_nms.hpp │ ├── gpu_nms.pyx │ ├── nms_kernel.cu │ └── py_cpu_nms.py │ ├── nms_wrapper.py │ ├── timer.py │ ├── yolo.py │ └── yolo.pyx ├── LICENSE ├── README.md └── klg_pipeline.png /.gitignore: -------------------------------------------------------------------------------- 1 | # Folders 2 | Preprocess/ 3 | data/ 4 | Data/ 5 | DATA/ 6 | Docs/ 7 | docs/ 8 | Doc/ 9 | doc/ 10 | .env/ 11 | env/ 12 | ENV/ 13 | __pycache__/ 14 | 15 | 16 | # Files 17 | data 18 | doc 19 | .remote-sync.json 20 | 21 | # Files of certain surfix 22 | *.h5 23 | *.pkl 24 | *.py[cod] 25 | *.so 26 | *.log 27 | *~ 28 | .python-version 29 | -------------------------------------------------------------------------------- /ClsKL/README.md: -------------------------------------------------------------------------------- 1 | # Knee KL Grading 2 | -------------------------------------------------------------------------------- /ClsKL/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PingjunChen/GradingKneeOA/dc19c2462111c3c5b27066e738fc11ddf57a6407/ClsKL/__init__.py -------------------------------------------------------------------------------- /ClsKL/test/test_kl.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | import os, sys, pdb 4 | import argparse 5 | import torch 6 | from torchvision import models 7 | import numpy as np 8 | 9 | FILE_PATH = os.path.abspath(__file__) 10 | PRJ_PATH = os.path.dirname(os.path.dirname(FILE_PATH)) 11 | sys.path.append(PRJ_PATH) 12 | 13 | from utils.loader import data_load 14 | from utils.eval_eng import eval_model, gen_vis_loc, gen_grad_cam 15 | 16 | 17 | def set_args(): 18 | parser = argparse.ArgumentParser(description='Pytorch Fine-Tune KL grading testing') 19 | parser.add_argument('--cuda-id', type=int, default=0) 20 | parser.add_argument('--batch-size', type=int, default=16) 21 | parser.add_argument('--data-dir', type=str, default='../../data/ClsKLData/kneeKL224') 22 | parser.add_argument('--model_dir', type=str, default='') 23 | parser.add_argument('--best_model_name', type=str, default='') 24 | # parser.add_argument('--save_dir', type=str, default='../../data/ClsKLData/models/CAMs/') 25 | parser.add_argument('--phase', type=str, default="test", choices=["test", "auto_test"]) 26 | 27 | args = parser.parse_args() 28 | return args 29 | 30 | 31 | if __name__ == "__main__": 32 | args = set_args() 33 | os.environ["CUDA_VISIBLE_DEVICES"] = str(args.cuda_id) 34 | import torch.backends.cudnn as cudnn 35 | cudnn.benchmark = True 36 | 37 | args.best_model_path = os.path.join(args.model_dir, args.best_model_name) 38 | assert os.path.exists(args.best_model_path), "Model doesnot exist" 39 | 40 | dset_loaders, dset_size, num_class = data_load(args) 41 | 42 | # Evaluate model 43 | print('---Evaluate model : {}--'.format(args.phase)) 44 | eval_model(args, args.phase, dset_loaders, dset_size) 45 | 46 | # Generate saliency visulization 47 | # gen_vis_loc(args, phase, dset_loaders, dset_size, args.save_dir) 48 | # gen_grad_cam(args, phase, dset_loaders, dset_size, args.save_dir) 49 | -------------------------------------------------------------------------------- /ClsKL/test/test_kl.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | python test_kl.py --cuda-id 7 \ 4 | --batch-size 16 \ 5 | --phase "test" \ 6 | --model_dir "../../data/ClsKLData/models/model_matrix/wm04/" \ 7 | --best_model_name "0.697-0.344.pth" 8 | -------------------------------------------------------------------------------- /ClsKL/train/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PingjunChen/GradingKneeOA/dc19c2462111c3c5b27066e738fc11ddf57a6407/ClsKL/train/__init__.py -------------------------------------------------------------------------------- /ClsKL/train/train_kl.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | import os, sys, pdb 4 | import argparse 5 | import torch 6 | import numpy as np 7 | 8 | FILE_PATH = os.path.abspath(__file__) 9 | PRJ_PATH = os.path.dirname(os.path.dirname(FILE_PATH)) 10 | sys.path.append(PRJ_PATH) 11 | 12 | from utils.loader import data_load 13 | # from utils.model import resnet 14 | from utils.model import cls_model 15 | from utils.train_eng import train_model 16 | from utils.eval_eng import eval_model 17 | 18 | 19 | def set_args(): 20 | parser = argparse.ArgumentParser(description='Pytorch Fine-Tune Resnet Training') 21 | parser.add_argument('--net_type', type=str, default="vgg") 22 | parser.add_argument('--depth', type=str, default="19") 23 | parser.add_argument('--lr', type=float, default=5.0e-4) 24 | parser.add_argument('--lr_decay_epoch', type=int, default=5) 25 | parser.add_argument('--num_epoch', type=int, default=12) 26 | parser.add_argument('--batch_size', type=int, default=32) 27 | parser.add_argument('--weight_decay', type=float, default=5.0e-4) 28 | parser.add_argument('--data_dir', type=str, default='../../data/ClsKLData/kneeKL224') 29 | parser.add_argument('--model_dir', type=str, default='../../data/ClsKLData/models/cmpLoss') 30 | parser.add_argument('--pretrained', type=bool, default=True) 31 | parser.add_argument('--cuda_id', type=int, default=0) 32 | parser.add_argument('--optim', type=str, default='SGD') 33 | parser.add_argument('--wloss', type=int, default=1) 34 | parser.add_argument('--session', type=int, default=0) 35 | 36 | args = parser.parse_args() 37 | return args 38 | 39 | 40 | if __name__ == '__main__': 41 | print('--Phase 0: Argument settings--') 42 | args = set_args() 43 | os.environ["CUDA_VISIBLE_DEVICES"] = str(args.cuda_id) 44 | import torch.backends.cudnn as cudnn 45 | cudnn.benchmark = True 46 | 47 | args.best_model_name = '{}-{}-{}-{}'.format(args.net_type, args.depth, args.optim, args.wloss) 48 | # import pdb; pdb.set_trace() 49 | 50 | print('--Phase 1: Data prepration--') 51 | dset_loaders, dset_size, num_class = data_load(args) 52 | args.num_class = num_class 53 | 54 | print('--Phase 2: Model setup--') 55 | model = cls_model(args) 56 | model.cuda() 57 | 58 | print('--Phase 3: Model training--') 59 | train_model(args, model, dset_loaders, dset_size) 60 | 61 | # print('--Phase 4: Evaluate model--') 62 | # phase='val' 63 | # eval_model(args, phase, dset_loaders, dset_size) 64 | -------------------------------------------------------------------------------- /ClsKL/train/train_kl.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | python train_kl.py --batch_size 24 \ 4 | --cuda_id 6 \ 5 | --net_type vgg \ 6 | --depth 19 \ 7 | --num_epoch 16 \ 8 | --session 3 9 | -------------------------------------------------------------------------------- /ClsKL/utils/__init__.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | -------------------------------------------------------------------------------- /ClsKL/utils/eval_eng.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | import os, sys, pdb 4 | import torch 5 | from torch.autograd import Variable 6 | import numpy as np 7 | import itertools 8 | import deepdish as dd 9 | import cv2 10 | from sklearn.metrics import confusion_matrix 11 | import matplotlib.pyplot as plt 12 | from matplotlib.backends.backend_pdf import PdfPages 13 | 14 | from .eval_util import ordinal_mse 15 | from .layer_util import extract_gap_layer, extract_vgg_fea_layer 16 | from .layer_util import gen_cam_visual 17 | from .grad_cam import GradCam, show_cam_on_image 18 | 19 | 20 | def eval_test(args, model, dset_loaders, dset_size, phase="test"): 21 | labels_all = [] * dset_size[phase] 22 | preds_all = [] * dset_size[phase] 23 | 24 | for data in dset_loaders[phase]: 25 | inputs, labels, _ = data 26 | inputs = Variable(inputs.cuda()) 27 | outputs = model(inputs) 28 | _, preds = torch.max(outputs.data, 1) 29 | 30 | labels_np = labels.numpy() 31 | labels = labels_np.tolist() 32 | labels_all.extend(labels) 33 | preds_cpu = preds.cpu() 34 | preds_np = preds_cpu.numpy() 35 | preds = preds_np.tolist() 36 | preds_all.extend(preds) 37 | 38 | conf_matrix = confusion_matrix(labels_all, preds_all) 39 | acc = 1.0*np.trace(conf_matrix)/np.sum(conf_matrix) 40 | mse = ordinal_mse(conf_matrix) 41 | 42 | return acc, mse 43 | 44 | 45 | def gen_vis_loc(args, phase, dset_loaders, dset_size, save_dir): 46 | model = torch.load(args.best_model_path) 47 | model.cuda() 48 | model.eval() 49 | 50 | count, ttl_num = 0, dset_size[phase] 51 | for data in dset_loaders[phase]: 52 | inputs, labels, paths = data 53 | count += len(paths) 54 | print("Processing {}/{}".format(count, ttl_num)) 55 | inputs = Variable(inputs.cuda()) 56 | 57 | # Prediction & CAMs 58 | preds, cams = gen_cam_visual(model, inputs) 59 | 60 | # Inputs & Labels 61 | labels = labels.numpy() 62 | pixel_mean, pixel_std = 0.66133188, 0.21229856 63 | inputs = inputs.permute(0, 2, 3, 1) 64 | inputs = inputs.data.cpu().numpy() 65 | inputs = (inputs * pixel_std) + pixel_mean 66 | inputs = np.clip(inputs, 0.0, 1.0) 67 | 68 | # inputs / cam_list / labels / preds 69 | for img, label, cam, pred, path in zip(inputs, labels, cams, preds, paths): 70 | alpha = 0.6 71 | beta = 1 - alpha 72 | cam = (cam * 255.0).astype(np.uint8) 73 | cam_rgb = cv2.applyColorMap(cam, cv2.COLORMAP_HSV) 74 | img = (img * 255.0).astype(np.uint8) 75 | img_cam = cv2.addWeighted(img, alpha, cam_rgb, beta, 0) 76 | 77 | # Plotting 78 | l_title = "Input: (Grade " + str(label) + ")" 79 | r_title = "Heatmap: (Grade " + str(pred) + ")" 80 | suptitle = path + ":" + str(label==pred) 81 | #subplot.plot_lr(img, img_cam, l_title, r_title, suptitle) 82 | 83 | # Saving plotting 84 | save_folder = os.path.join(save_dir, str(label)) 85 | if not os.path.exists(save_folder): 86 | os.mkdir(save_folder) 87 | 88 | save_path = os.path.join(save_folder, os.path.splitext(path)[0] + ".png") 89 | # save_path = os.path.join(save_folder, os.path.splitext(path)[0] + ".pdf") 90 | # with PdfPages(save_path) as pdf: 91 | fig, axes = plt.subplots(nrows=1, ncols=2, figsize=(10, 4)) 92 | axes[0].imshow(img) 93 | axes[0].set_title(l_title) 94 | axes[1].imshow(img_cam) 95 | axes[1].set_title(r_title) 96 | plt.suptitle(suptitle, fontsize=16) 97 | plt.savefig(save_path) 98 | plt.close('all') 99 | 100 | 101 | def eval_model(args, phase, dset_loaders, dset_size): 102 | model = torch.load(args.best_model_path) 103 | model.cuda() 104 | model.eval() 105 | 106 | labels_all = [] * dset_size[phase] 107 | preds_all = [] * dset_size[phase] 108 | feas_all = [] 109 | 110 | for data in dset_loaders[phase]: 111 | inputs, labels, paths = data 112 | inputs = Variable(inputs.cuda()) 113 | 114 | # forward 115 | outputs = model(inputs) 116 | _, preds = torch.max(outputs.data, 1) 117 | 118 | # retrieve gap layer 119 | # gaps = extract_gap_layer(model, inputs) 120 | vgg_feas = extract_vgg_fea_layer(model, inputs) 121 | feas_all.append(vgg_feas.data.cpu().numpy()) 122 | 123 | labels_np = labels.numpy() 124 | labels = labels_np.tolist() 125 | labels_all.extend(labels) 126 | preds_cpu = preds.cpu() 127 | preds_np = preds_cpu.numpy() 128 | preds = preds_np.tolist() 129 | preds_all.extend(preds) 130 | 131 | conf_matrix = confusion_matrix(labels_all, preds_all) 132 | print("In {}: confusion matrix is:\n {}".format(phase, conf_matrix)) 133 | acc = 1.0*np.trace(conf_matrix)/np.sum(conf_matrix) 134 | print('True/Total: {}/{}'.format(np.trace(conf_matrix), np.sum(conf_matrix))) 135 | # print('MSE: {:.4f}'.format(ordinal_mse(conf_matrix, poly_num=2))) 136 | print('Acc: {:.3f} ABE: {:.3f}'.format(acc, ordinal_mse(conf_matrix, poly_num=1))) 137 | 138 | # plt.imshow(conf_matrix, interpolation='nearest', cmap=plt.cm.Blues) 139 | # plt.title('Acc: {:.3f} MAE: {:.3f}'.format(acc, ordinal_mse(conf_matrix, poly_num=1)), fontsize=18) 140 | # classes = [0, 1, 2, 3, 4] 141 | # tick_marks = np.arange(len(classes)) 142 | # plt.xticks(tick_marks, classes, rotation=45, fontsize=14) 143 | # plt.yticks(tick_marks, classes, fontsize=14) 144 | # for i, j in itertools.product(range(conf_matrix.shape[0]), range(conf_matrix.shape[1])): 145 | # plt.text(j, i, format(conf_matrix[i, j], 'd'), 146 | # horizontalalignment="center", 147 | # color="white" if conf_matrix[i, j] > conf_matrix.max() / 2. else "black") 148 | # plt.tight_layout() 149 | # plt.savefig('cm04.svg') 150 | 151 | 152 | 153 | def gen_grad_cam(args, phase, dset_loaders, dset_size, save_dir): 154 | model = torch.load(args.best_model_path) 155 | model.cuda() 156 | model.eval() 157 | 158 | alpha = 0.6 159 | beta = 1 - alpha 160 | pixel_mean, pixel_std = 0.66133188, 0.21229856 161 | 162 | grad_cam = GradCam(model, target_layer_names = ["35"], use_cuda=1) 163 | count, ttl_num = 0, dset_size[phase] 164 | for data in dset_loaders[phase]: 165 | inputs, labels, paths = data 166 | count += len(paths) 167 | print("Processing {}/{}".format(count, ttl_num)) 168 | inputs = Variable(inputs.cuda()) 169 | 170 | for input, label, path in zip(inputs, labels, paths): 171 | input.unsqueeze_(0) 172 | target_index = label.tolist() 173 | mask = grad_cam(input, target_index) 174 | 175 | input = input.permute(0, 2, 3, 1) 176 | img = input.data.cpu().numpy() 177 | img = (np.squeeze(img) * pixel_std) + pixel_mean 178 | 179 | cam = (mask * 255.0).astype(np.uint8) 180 | cam_rgb = cv2.applyColorMap(cam, cv2.COLORMAP_HSV) 181 | 182 | img = (img * 255.0).astype(np.uint8) 183 | img_cam = cv2.addWeighted(img, alpha, cam_rgb, beta, 0) 184 | 185 | save_path = os.path.join(save_dir, str(target_index), path) 186 | cv2.imwrite(save_path, img_cam) 187 | -------------------------------------------------------------------------------- /ClsKL/utils/eval_util.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | import os, sys 4 | import numpy as np 5 | 6 | def build_cof_mat(num, poly_num=2): 7 | assert num > 2, "num should bigger than 2" 8 | cof_mat = np.zeros((num, num), dtype=np.float32) 9 | for i in range(0, num): 10 | for j in range(0, num): 11 | cof_mat[i, j] = np.power(np.absolute(i-j), poly_num) 12 | 13 | return cof_mat 14 | 15 | 16 | def ordinal_mse(confusion_matrix, poly_num=2): 17 | category_num = confusion_matrix.shape[0] 18 | cof_mat = build_cof_mat(category_num, poly_num=poly_num) 19 | 20 | err_mat = np.multiply(confusion_matrix, cof_mat) 21 | mse_val = np.sum(err_mat) * 1.0 / np.sum(confusion_matrix) 22 | 23 | return mse_val 24 | -------------------------------------------------------------------------------- /ClsKL/utils/grad_cam.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | import os, sys, pdb 4 | 5 | import torch 6 | from torch.autograd import Variable 7 | from torch.autograd import Function 8 | from torchvision import models 9 | from torchvision import utils 10 | import numpy as np 11 | import cv2 12 | 13 | 14 | class FeatureExtractor(): 15 | """ Class for extracting activations and 16 | registering gradients from targetted intermediate layers """ 17 | def __init__(self, model, target_layers): 18 | self.model = model 19 | self.target_layers = target_layers 20 | self.gradients = [] 21 | 22 | def save_gradient(self, grad): 23 | self.gradients.append(grad) 24 | 25 | def __call__(self, x): 26 | outputs = [] 27 | self.gradients = [] 28 | for name, module in self.model._modules.items(): 29 | x = module(x) 30 | if name in self.target_layers: 31 | x.register_hook(self.save_gradient) 32 | outputs += [x] 33 | return outputs, x 34 | 35 | class ModelOutputs(): 36 | """ Class for making a forward pass, and getting: 37 | 1. The network output. 38 | 2. Activations from intermeddiate targetted layers. 39 | 3. Gradients from intermeddiate targetted layers. """ 40 | def __init__(self, model, target_layers): 41 | self.model = model 42 | self.feature_extractor = FeatureExtractor(self.model.features, target_layers) 43 | 44 | def get_gradients(self): 45 | return self.feature_extractor.gradients 46 | 47 | def __call__(self, x): 48 | target_activations, output = self.feature_extractor(x) 49 | output = output.view(output.size(0), -1) 50 | output = self.model.classifier(output) 51 | return target_activations, output 52 | 53 | 54 | def show_cam_on_image(img, mask): 55 | heatmap = cv2.applyColorMap(np.uint8(255*mask), cv2.COLORMAP_JET) 56 | heatmap = np.float32(heatmap) / 255 57 | cam = heatmap + np.float32(img) 58 | cam = cam / np.max(cam) 59 | cv2.imwrite("cam.jpg", np.uint8(255 * cam)) 60 | 61 | 62 | class GradCam: 63 | def __init__(self, model, target_layer_names, use_cuda): 64 | self.model = model 65 | # self.model.eval() 66 | self.cuda = use_cuda 67 | 68 | self.extractor = ModelOutputs(self.model, target_layer_names) 69 | 70 | def forward(self, input): 71 | return self.model(input) 72 | 73 | def __call__(self, input, index = None): 74 | features, output = self.extractor(input) 75 | 76 | if index == None: 77 | index = np.argmax(output.cpu().data.numpy()) 78 | 79 | one_hot = np.zeros((1, output.size()[-1]), dtype = np.float32) 80 | one_hot[0][index] = 1 81 | one_hot = Variable(torch.from_numpy(one_hot), requires_grad = True) 82 | one_hot = torch.sum(one_hot.cuda(0) * output) 83 | 84 | self.model.features.zero_grad() 85 | self.model.classifier.zero_grad() 86 | one_hot.backward() 87 | 88 | grads_val = self.extractor.get_gradients()[-1].cpu().data.numpy() 89 | 90 | target = features[-1] 91 | target = target.cpu().data.numpy()[0, :] 92 | 93 | weights = np.mean(grads_val, axis = (2, 3))[0, :] 94 | cam = np.ones(target.shape[1 : ], dtype = np.float32) 95 | 96 | for i, w in enumerate(weights): 97 | cam += w * target[i, :, :] 98 | 99 | cam = np.maximum(cam, 0) 100 | cam = cv2.resize(cam, (224, 224)) 101 | cam = cam - np.min(cam) 102 | cam = cam / np.max(cam) 103 | 104 | return cam 105 | -------------------------------------------------------------------------------- /ClsKL/utils/knee_sets.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | import os, sys, pdb 4 | import torch.utils.data as data 5 | 6 | from PIL import Image 7 | import os 8 | import os.path 9 | 10 | IMG_EXTENSIONS = ['.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm'] 11 | 12 | 13 | def is_image_file(filename): 14 | """Checks if a file is an image. 15 | Args: 16 | filename (string): path to a file 17 | Returns: 18 | bool: True if the filename ends with a known image extension 19 | """ 20 | filename_lower = filename.lower() 21 | return any(filename_lower.endswith(ext) for ext in IMG_EXTENSIONS) 22 | 23 | 24 | def find_classes(dir): 25 | classes = [d for d in os.listdir(dir) if os.path.isdir(os.path.join(dir, d))] 26 | classes.sort() 27 | class_to_idx = {classes[i]: i for i in range(len(classes))} 28 | return classes, class_to_idx 29 | 30 | 31 | def make_dataset(dir, class_to_idx): 32 | images = [] 33 | dir = os.path.expanduser(dir) 34 | for target in sorted(os.listdir(dir)): 35 | d = os.path.join(dir, target) 36 | if not os.path.isdir(d): 37 | continue 38 | 39 | for root, _, fnames in sorted(os.walk(d)): 40 | for fname in sorted(fnames): 41 | if is_image_file(fname): 42 | path = os.path.join(root, fname) 43 | item = (path, class_to_idx[target]) 44 | images.append(item) 45 | 46 | return images 47 | 48 | def pil_loader(path): 49 | # open path as file to avoid ResourceWarning (https://github.com/python-pillow/Pillow/issues/835) 50 | with open(path, 'rb') as f: 51 | img = Image.open(f) 52 | return img.convert('RGB') 53 | 54 | def accimage_loader(path): 55 | import accimage 56 | try: 57 | return accimage.Image(path) 58 | except IOError: 59 | # Potentially a decoding problem, fall back to PIL.Image 60 | return pil_loader(path) 61 | 62 | def default_loader(path): 63 | from torchvision import get_image_backend 64 | if get_image_backend() == 'accimage': 65 | return accimage_loader(path) 66 | else: 67 | return pil_loader(path) 68 | 69 | class ImageFolder(data.Dataset): 70 | """A generic data loader where the images are arranged in this way: :: 71 | root/dog/xxx.png 72 | root/dog/xxy.png 73 | root/dog/xxz.png 74 | root/cat/123.png 75 | root/cat/nsdf3.png 76 | root/cat/asd932_.png 77 | Args: 78 | root (string): Root directory path. 79 | transform (callable, optional): A function/transform that takes in an PIL image 80 | and returns a transformed version. E.g, ``transforms.RandomCrop`` 81 | target_transform (callable, optional): A function/transform that takes in the 82 | target and transforms it. 83 | loader (callable, optional): A function to load an image given its path. 84 | Attributes: 85 | classes (list): List of the class names. 86 | class_to_idx (dict): Dict with items (class_name, class_index). 87 | imgs (list): List of (image path, class_index) tuples 88 | """ 89 | 90 | def __init__(self, root, transform=None, target_transform=None, 91 | loader=default_loader): 92 | classes, class_to_idx = find_classes(root) 93 | imgs = make_dataset(root, class_to_idx) 94 | if len(imgs) == 0: 95 | raise(RuntimeError("Found 0 images in subfolders of: " + root + "\n" 96 | "Supported image extensions are: " + ",".join(IMG_EXTENSIONS))) 97 | 98 | self.root = root 99 | self.imgs = imgs 100 | self.classes = classes 101 | self.class_to_idx = class_to_idx 102 | self.transform = transform 103 | self.target_transform = target_transform 104 | self.loader = loader 105 | 106 | def __getitem__(self, index): 107 | """ 108 | Args: 109 | index (int): Index 110 | Returns: 111 | tuple: (image, target) where target is class_index of the target class. 112 | """ 113 | path, target = self.imgs[index] 114 | img = self.loader(path) 115 | if self.transform is not None: 116 | img = self.transform(img) 117 | if self.target_transform is not None: 118 | target = self.target_transform(target) 119 | 120 | filename = os.path.basename(path) 121 | 122 | return img, target, filename 123 | 124 | def __len__(self): 125 | return len(self.imgs) 126 | -------------------------------------------------------------------------------- /ClsKL/utils/layer_util.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | import os, sys, pdb 4 | 5 | import torch 6 | import torch.nn as nn 7 | import numpy as np 8 | from skimage import io, transform 9 | 10 | 11 | def extract_vgg_fea_layer(model, inputs): 12 | x = model.features(inputs) 13 | x = x.view(x.size(0), -1) 14 | fea_extractor = nn.Sequential(*list(model.classifier.children())[:-1]) 15 | vgg_feas = fea_extractor(x) 16 | 17 | return vgg_feas 18 | 19 | def extract_gap_layer(model, inputs): 20 | x = model.conv1(inputs) 21 | x = model.bn1(x) 22 | x = model.relu(x) 23 | x = model.maxpool(x) 24 | 25 | x = model.layer1(x) 26 | x = model.layer2(x) 27 | x = model.layer3(x) 28 | x = model.layer4(x) 29 | 30 | gap = model.avgpool(x) 31 | x = x.view(gap.size(0), -1) 32 | 33 | return x 34 | 35 | def gen_cam_visual(model, inputs): 36 | x = model.conv1(inputs) 37 | x = model.bn1(x) 38 | x = model.relu(x) 39 | x = model.maxpool(x) 40 | 41 | x = model.layer1(x) 42 | x = model.layer2(x) 43 | x = model.layer3(x) 44 | fea77 = model.layer4(x) 45 | gap512 = model.avgpool(fea77) 46 | gap512 = gap512.view(gap512.size(0), -1) 47 | pred_prob = model.fc(gap512) 48 | 49 | # Get prediction result 50 | _, pred_ind = torch.max(pred_prob, dim=1) 51 | preds = pred_ind.data 52 | fc_weight = model.fc.weight.data 53 | class_w = fc_weight[preds, :] 54 | 55 | class_w = class_w.cpu().numpy() 56 | fea77 = fea77.data.cpu().numpy() 57 | 58 | cams = [] 59 | for ind in range(fea77.shape[0]): 60 | cur_w = class_w[ind:ind+1, :] 61 | cur_fea = fea77[ind].reshape(fea77.shape[1], -1) 62 | fea_map = np.matmul(cur_w, cur_fea).reshape(fea77.shape[2], fea77.shape[3]) 63 | fea_map = (fea_map - np.amin(fea_map)) * 1.0 / (np.amax(fea_map) - np.amin(fea_map)) 64 | fea_map = transform.resize(fea_map, (inputs.shape[2], inputs.shape[3])) 65 | 66 | cams.append(fea_map) 67 | 68 | preds = preds.cpu().numpy() 69 | cams = np.array(cams) 70 | 71 | return preds, cams 72 | -------------------------------------------------------------------------------- /ClsKL/utils/loader.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | import os, sys, pdb 4 | import torch 5 | from torchvision import transforms 6 | 7 | from .knee_sets import ImageFolder 8 | 9 | 10 | def data_load(args): 11 | pixel_mean, pixel_std = 0.66133188, 0.21229856 12 | phases = ['train', 'val', 'test', 'auto_test'] 13 | # phases = ['train', 'val', 'test', 'auto_test'] 14 | data_transform = { 15 | 'train': transforms.Compose([ 16 | transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.3, hue=0.3), 17 | transforms.ToTensor(), 18 | transforms.Normalize([pixel_mean]*3, [pixel_std]*3) 19 | ]), 20 | 'val': transforms.Compose([ 21 | transforms.ToTensor(), 22 | transforms.Normalize([pixel_mean]*3, [pixel_std]*3) 23 | ]), 24 | 'test': transforms.Compose([ 25 | transforms.ToTensor(), 26 | transforms.Normalize([pixel_mean]*3, [pixel_std]*3) 27 | ]), 28 | 'auto_test': transforms.Compose([ 29 | transforms.ToTensor(), 30 | transforms.Normalize([pixel_mean]*3, [pixel_std]*3) 31 | ]), 32 | 'most_test': transforms.Compose([ 33 | transforms.ToTensor(), 34 | transforms.Normalize([pixel_mean]*3, [pixel_std]*3) 35 | ]), 36 | 'most_auto_test': transforms.Compose([ 37 | transforms.ToTensor(), 38 | transforms.Normalize([pixel_mean]*3, [pixel_std]*3) 39 | ]) 40 | } 41 | 42 | dsets = {x: ImageFolder(os.path.join(args.data_dir, x), data_transform[x]) for x in phases} 43 | dset_loaders = {x: torch.utils.data.DataLoader(dsets[x], batch_size=args.batch_size, 44 | shuffle=(x=='train'), num_workers=4) for x in phases} 45 | dset_classes = dsets['train'].classes 46 | dset_size = {x: len(dsets[x]) for x in phases} 47 | num_class = len(dset_classes) 48 | 49 | return dset_loaders, dset_size, num_class 50 | -------------------------------------------------------------------------------- /ClsKL/utils/loss_util.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | import os, sys, pdb 4 | import torch 5 | import numpy as np 6 | 7 | def weighted_loss(outputs, labels, args): 8 | softmax_op = torch.nn.Softmax(1) 9 | prob_pred = softmax_op(outputs) 10 | 11 | def set_weights(): 12 | # # weight matrix 01 (wm01) 13 | # init_weights = np.array([[1, 2, 3, 4, 5], 14 | # [2, 1, 2, 3, 4], 15 | # [3, 2, 1, 2, 3], 16 | # [4, 3, 2, 1, 2], 17 | # [5, 4, 3, 2, 1]], dtype=np.float) 18 | 19 | # weight matrix 02 (wm02) 20 | init_weights = np.array([[1, 3, 5, 7, 9], 21 | [3, 1, 3, 5, 7], 22 | [5, 3, 1, 3, 5], 23 | [7, 5, 3, 1, 3], 24 | [9, 7, 5, 3, 1]], dtype=np.float) 25 | 26 | # # weight matrix 03 (wm03) 27 | # init_weights = np.array([[1, 4, 7, 10, 13], 28 | # [4, 1, 4, 7, 10], 29 | # [7, 4, 1, 4, 7], 30 | # [10, 7, 4, 1, 4], 31 | # [13, 10, 7, 4, 1]], dtype=np.float) 32 | 33 | # # weight matrix 04 (wm04) 34 | # init_weights = np.array([[1, 3, 6, 7, 9], 35 | # [4, 1, 4, 5, 7], 36 | # [6, 4, 1, 3, 5], 37 | # [7, 5, 3, 1, 3], 38 | # [9, 7, 5, 3, 1]], dtype=np.float) 39 | 40 | adjusted_weights = init_weights + 1.0 41 | np.fill_diagonal(adjusted_weights, 0) 42 | 43 | return adjusted_weights 44 | cls_weights = set_weights() 45 | 46 | batch_num, class_num = outputs.size() 47 | class_hot = np.zeros([batch_num, class_num], dtype=np.float32) 48 | labels_np = labels.data.cpu().numpy() 49 | for ind in range(batch_num): 50 | class_hot[ind, :] = cls_weights[labels_np[ind], :] 51 | class_hot = torch.from_numpy(class_hot) 52 | class_hot = torch.autograd.Variable(class_hot).cuda() 53 | 54 | loss = torch.sum((prob_pred * class_hot)**2) / batch_num 55 | # loss = torch.mean(prob_pred * class_hot) 56 | 57 | return loss 58 | -------------------------------------------------------------------------------- /ClsKL/utils/model.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | import os, sys, pdb 4 | import torch 5 | import torch.nn as nn 6 | from torchvision import models 7 | 8 | 9 | def cls_model(args): 10 | if args.net_type == "resnet": 11 | if args.depth == "18": 12 | model = models.resnet18(pretrained = args.pretrained) 13 | elif args.depth == "34": 14 | model = models.resnet34(pretrained = args.pretrained) 15 | elif args.depth == "50": 16 | model = models.resnet50(pretrained = args.pretrained) 17 | elif args.depth == "101": 18 | model = models.resnet101(pretrained = args.pretrained) 19 | elif args.depth == "152": 20 | model = models.resnet152(pretrained = args.pretrained) 21 | else: 22 | return None 23 | 24 | in_features = model.fc.in_features 25 | model.fc = torch.nn.Linear(in_features, args.num_class) 26 | 27 | elif args.net_type == "vgg": 28 | if args.depth == "16": 29 | model = models.vgg16(pretrained = args.pretrained) 30 | elif args.depth == "19": 31 | model = models.vgg19(pretrained = args.pretrained) 32 | elif args.depth == "16bn": 33 | model = models.vgg16_bn(pretrained = args.pretrained) 34 | elif args.depth == "19bn": 35 | model = models.vgg19_bn(pretrained = args.pretrained) 36 | else: 37 | return None 38 | 39 | num_ftrs = model.classifier[6].in_features 40 | feature_model = list(model.classifier.children()) 41 | feature_model.pop() 42 | feature_model.append(nn.Linear(num_ftrs, args.num_class)) 43 | model.classifier = nn.Sequential(*feature_model) 44 | 45 | elif args.net_type == "densenet": 46 | if args.depth == "121": 47 | model = models.densenet121(pretrained = args.pretrained) 48 | elif args.depth == "169": 49 | model = models.densenet169(pretrained = args.pretrained) 50 | elif args.depth == "201": 51 | model = models.densenet201(pretrained = args.pretrained) 52 | else: 53 | return None 54 | 55 | in_features = model.classifier.in_features 56 | model.classifier = torch.nn.Linear(in_features, args.num_class) 57 | 58 | elif args.net_type == "inception": 59 | if args.depth == "v3": 60 | model = models.inception_v3(pretrained = args.pretrained) 61 | else: 62 | return None 63 | 64 | in_features = model.fc.in_features 65 | model.fc = torch.nn.Linear(in_features, args.num_class) 66 | 67 | else: 68 | return None 69 | 70 | return model 71 | -------------------------------------------------------------------------------- /ClsKL/utils/torch_util.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | import os, sys, pdb 4 | 5 | class LRScheduler(): 6 | def __init__(self, init_lr=1.0e-4, lr_decay_epoch=10): 7 | self.init_lr = init_lr 8 | self.lr_decay_epoch = lr_decay_epoch 9 | 10 | def __call__(self, optimizer, epoch): 11 | '''Decay learning rate by a factor of 0.1 every lr_decay_epoch epochs.''' 12 | lr = self.init_lr * (0.8 ** (epoch // self.lr_decay_epoch)) 13 | lr = max(lr, 1e-8) 14 | if epoch % self.lr_decay_epoch == 0: 15 | print ('LR is set to {}'.format(lr)) 16 | 17 | for param_group in optimizer.param_groups: 18 | param_group['lr'] = lr 19 | 20 | return optimizer 21 | -------------------------------------------------------------------------------- /ClsKL/utils/train_eng.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | import os, sys, pdb 4 | import shutil 5 | import torch 6 | from torch import nn, optim 7 | from torch.autograd import Variable 8 | import time 9 | 10 | from torch_util import LRScheduler 11 | from loss_util import weighted_loss 12 | from eval_eng import eval_test 13 | 14 | 15 | def train_model(args, model, dset_loaders, dset_size): 16 | best_model, best_acc, best_num_epoch = None, .0, 0 17 | best_model_path = os.path.join(args.model_dir, args.best_model_name, str(args.session)) 18 | if os.path.exists(best_model_path): 19 | shutil.rmtree(best_model_path) 20 | os.makedirs(best_model_path) 21 | 22 | criterion = nn.CrossEntropyLoss() 23 | if args.optim == 'SGD': 24 | optimizer = optim.SGD(model.parameters(), lr=args.lr, 25 | momentum=0.9, weight_decay=args.weight_decay) 26 | elif args.optim == 'Adam': 27 | optimizer = optim.Adam(model.parameters(), lr=args.lr, 28 | weight_decay=args.weight_decay) 29 | elif args.optim == 'RMSprop': 30 | optimizer = optim.RMSprop(model.parameters(), lr=args.lr, 31 | weight_decay=args.weight_decay, momentum=0.9) 32 | 33 | # print('-'*10+'Training'+'-'*10) 34 | print('Device id:{} Initial lr:{} Optimizer:{} network:{} depth:{} num_class:{}'.format( 35 | args.cuda_id, args.lr, args.optim, args.net_type, args.depth, args.num_class)) 36 | 37 | lr_scheduler = LRScheduler(args.lr, args.lr_decay_epoch) 38 | for epoch in range(args.num_epoch): 39 | since = time.time() 40 | print('Epoch {}/{}'.format(epoch, args.num_epoch)) 41 | for phase in ['train', 'val']: 42 | if phase == 'train': 43 | optimizer = lr_scheduler(optimizer, epoch) 44 | model.train(True) 45 | else: 46 | model.train(False) 47 | 48 | running_loss, running_corrects = .0, .0 49 | 50 | for data in dset_loaders[phase]: 51 | inputs, labels, _ = data 52 | inputs, labels = Variable(inputs.cuda()), Variable(labels.cuda()) 53 | 54 | optimizer.zero_grad() 55 | outputs = model(inputs) 56 | if isinstance(outputs, tuple): 57 | outputs = outputs[0] 58 | _, preds = torch.max(outputs.data, 1) 59 | if args.wloss == False: 60 | loss = criterion(outputs, labels) 61 | else: 62 | loss = weighted_loss(outputs, labels, args) 63 | 64 | if phase == 'train': 65 | loss.backward() 66 | optimizer.step() 67 | 68 | running_loss += loss.data 69 | running_corrects += torch.sum(preds == labels.data) 70 | 71 | epoch_loss = 1.0*running_loss.cpu().tolist()/dset_size[phase] 72 | epoch_acc = 1.0*running_corrects.cpu().tolist()/dset_size[phase] 73 | elapse_time = time.time() - since 74 | print("In {}, Number case:{} Loss:{:.4f} Acc:{:.4f} Time:{}".format( 75 | phase, dset_size[phase], epoch_loss, epoch_acc, elapse_time)) 76 | 77 | if phase == "val" and epoch_acc >= best_acc: 78 | best_acc = epoch_acc 79 | best_num_epoch = epoch 80 | val_metric_str = str(epoch).zfill(2) + '-' + str(round(best_acc, 3)) 81 | test_acc, test_mse = eval_test(args, model, dset_loaders, dset_size, "test") 82 | test_metric_str = "-" + str(round(test_acc, 3)) + "-" + str(round(test_mse, 3)) + ".pth" 83 | args.best_model_path = os.path.join(best_model_path, val_metric_str + test_metric_str) 84 | torch.save(model.cpu(), args.best_model_path) 85 | print("---On test_set: acc is {:.3f}, mse is {:.3f}".format(test_acc, test_mse)) 86 | model.cuda() 87 | 88 | print('='*80) 89 | print ('Validation best_acc: {} best_num_epoch: {}'.format(best_acc, best_num_epoch)) 90 | -------------------------------------------------------------------------------- /ClsKL/vis/conf_matrix.py: -------------------------------------------------------------------------------- 1 | 2 | import numpy as np 3 | import matplotlib.pyplot as plt 4 | 5 | from plot_util import plot_confusion_matrix 6 | 7 | # Show Confusion Matrix of ResNet-34 Cross-Entropy on Mannual Cropped Knee Joints 8 | np.set_printoptions(precision=2) 9 | plt.figure() 10 | cnf_matrix = np.array([[556, 43, 39, 1, 0], [131, 52, 108, 5, 0], [41, 32, 335, 37, 2], 11 | [1, 1, 33, 167, 21], [0, 0, 1, 7, 43]]) 12 | class_names=["0", "1", "2", "3", "4"] 13 | plot_confusion_matrix(cnf_matrix, classes=class_names, 14 | title='VGG-19-Ordinal on Manual') 15 | 16 | # # Show Confusion Matrix of Propsed Methods on Mannual Cropped Knee Joints 17 | # np.set_printoptions(precision=2) 18 | # plt.figure() 19 | # cnf_matrix = np.array([[484, 125, 29, 1, 0], [125, 104, 60, 7, 0], [66, 92, 246, 43, 0], 20 | # [0, 7, 34, 169, 13], [0, 0, 0, 3, 48]]) 21 | # class_names=["0", "1", "2", "3", "4"] 22 | # plot_confusion_matrix(cnf_matrix, classes=class_names, 23 | # title='Proposed on Manual') 24 | 25 | # # Show Confusion Matrix of ResNet-34 Cross-Entropy on Mannual Cropped Knee Joints 26 | # np.set_printoptions(precision=2) 27 | # plt.figure() 28 | # cnf_matrix = np.array([[545, 44, 49, 1, 0], [178, 36, 80, 2, 0], [123, 31, 276, 17, 0], 29 | # [9, 3, 45, 162, 4], [0, 0, 1, 9, 41]]) 30 | # class_names=["0", "1", "2", "3", "4"] 31 | # plot_confusion_matrix(cnf_matrix, classes=class_names, 32 | # title='ResNet-34-CE on Manual') 33 | 34 | # # Show Confusion Matrix of Propsed Methods on Mannual Cropped Knee Joints 35 | # np.set_printoptions(precision=2) 36 | # plt.figure() 37 | # cnf_matrix = np.array([[473, 114, 48, 1, 0], [119, 114, 58, 5, 0], [78, 98, 229, 36, 1], 38 | # [0, 5, 36, 163, 19], [0, 0, 0, 7, 42]]) 39 | # class_names=["0", "1", "2", "3", "4"] 40 | # plot_confusion_matrix(cnf_matrix, classes=class_names, 41 | # title='Proposed on Automatic') 42 | 43 | # # Show Confusion Matrix of ResNet-34 Cross-Entropy on Mannual Cropped Knee Joints 44 | # np.set_printoptions(precision=2) 45 | # plt.figure() 46 | # cnf_matrix = np.array([[537, 47, 52, 0, 0], [176, 42, 75, 3, 0], [131, 51, 249, 11, 0], 47 | # [6, 7, 48, 154, 8], [0, 0, 0, 8, 41]]) 48 | # class_names=["0", "1", "2", "3", "4"] 49 | # plot_confusion_matrix(cnf_matrix, classes=class_names, 50 | # title='ResNet-34-CE on Automatic') 51 | 52 | 53 | 54 | # plt.tight_layout() 55 | plt.show() 56 | -------------------------------------------------------------------------------- /ClsKL/vis/knee_tsne.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | import os, sys, pdb 4 | import numpy as np 5 | import deepdish as dd 6 | from time import time 7 | import matplotlib.pyplot as plt 8 | from sklearn import manifold 9 | from yellowbrick.text import TSNEVisualizer 10 | 11 | 12 | # knee = dd.io.load('./data/feas1646_auto.h5') 13 | knee = dd.io.load('./data/tsne/vgg19_feas1656_manual.h5') 14 | X = knee["data"] 15 | y = knee["target"] 16 | 17 | n_samples, n_features = X.shape 18 | 19 | tsne = manifold.TSNE(n_components=2, perplexity=20, early_exaggeration=4.0, learning_rate=1000, n_iter=1000, 20 | n_iter_without_progress=50, min_grad_norm=0, init='pca', method='exact', verbose=2) 21 | Y = tsne.fit_transform(X) 22 | 23 | 24 | plt.figure(figsize=(6, 5)) 25 | colors = ['b', 'g', 'r', 'y', 'k'] 26 | target_ids = [0, 1, 2, 3, 4] 27 | target_labels = ["Grade 0", "Grade 1", "Grade 2", "Grade 3", "Grade 4"] 28 | for i, c, label in zip(target_ids, colors, target_labels): 29 | newY = np.array([Y[ind] for ind, e in enumerate(y) if e==i]) 30 | plt.scatter(newY[:, 0], newY[:, 1], c=c, label=label) 31 | plt.legend() 32 | plt.title("Features of VGG-19-Ordinal on Manual") 33 | plt.savefig("vgg19_tsne.pdf") 34 | # plt.tight_layout() 35 | plt.show() 36 | -------------------------------------------------------------------------------- /ClsKL/vis/plot_util.py: -------------------------------------------------------------------------------- 1 | import itertools 2 | import numpy as np 3 | import matplotlib.pyplot as plt 4 | 5 | def plot_confusion_matrix(cm, classes, 6 | normalize=False, 7 | title='Confusion matrix', 8 | cmap=plt.cm.Blues): 9 | """ 10 | This function prints and plots the confusion matrix. 11 | Normalization can be applied by setting `normalize=True`. 12 | """ 13 | if normalize: 14 | cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis] 15 | print("Normalized confusion matrix") 16 | else: 17 | print('Confusion matrix, without normalization') 18 | 19 | print(cm) 20 | 21 | plt.imshow(cm, interpolation='nearest', cmap=cmap) 22 | plt.title(title, fontsize=20) 23 | # plt.colorbar() 24 | tick_marks = np.arange(len(classes)) 25 | plt.xticks(tick_marks, classes, rotation=45, fontsize=14) 26 | plt.yticks(tick_marks, classes, fontsize=14) 27 | 28 | fmt = '.2f' if normalize else 'd' 29 | thresh = cm.max() / 2. 30 | for i, j in itertools.product(range(cm.shape[0]), range(cm.shape[1])): 31 | plt.text(j, i, format(cm[i, j], fmt), 32 | horizontalalignment="center", 33 | color="white" if cm[i, j] > thresh else "black") 34 | 35 | # plt.tight_layout() 36 | plt.ylabel('True label', fontsize=18) 37 | plt.xlabel('Predicted label', fontsize=18) 38 | plt.savefig("confusion_matrix.pdf") 39 | -------------------------------------------------------------------------------- /DetJoint/README.md: -------------------------------------------------------------------------------- 1 | # Knee Detection 2 | Aims to automatically detect knee joints from X-ray image 3 | -------------------------------------------------------------------------------- /DetJoint/eval/model_eval.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | import os, sys, pdb 4 | import matplotlib.pyplot as plt 5 | import numpy as np 6 | 7 | def ablation_compare(): 8 | num_box = [1, 2, 3, 4, 5, 6] 9 | with_wd = [0.922, 0.868, 0.902, 0.872, 0.871, 0.855] 10 | no_wd = [0.911, 0.864, 0.865, 0.856, 0.851, 0.828] 11 | 12 | plt.plot(num_box, with_wd, marker='D', linestyle='--', color='r', linewidth=3.0, 13 | label='use weight decay') 14 | plt.plot(num_box, no_wd, marker='o', linestyle='--', color='b', linewidth=3.0, 15 | label='no weight decay') 16 | plt.xlim(0.6, 6.4) 17 | plt.ylim(0.8, 0.96) 18 | plt.legend() 19 | # plt.tight_layout() 20 | 21 | plt.xlabel('Number of bounding box') 22 | plt.ylabel('Recall') 23 | # plt.title('Knee joint detection') 24 | plt.grid(True) 25 | plt.savefig("yolo_para.pdf") 26 | plt.show() 27 | 28 | if __name__ == "__main__": 29 | ablation_compare() 30 | -------------------------------------------------------------------------------- /DetJoint/preprocess/extract_yolo_det.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | import os, sys, pdb 4 | import deepdish as dd 5 | import glob 6 | from skimage import io 7 | import scipy.misc as misc 8 | 9 | # def expand_bbox() 10 | 11 | def extract_detected_knees(data_dir, det_dir, results_dir, expand=0.3): 12 | img_list = glob.glob(os.path.join(data_dir, "*.png")) 13 | for cur_img in img_list: 14 | full_name = os.path.basename(cur_img) 15 | cur_name = os.path.splitext(full_name)[0] 16 | h5_path = os.path.join(det_dir, cur_name+".h5") 17 | det_dict = dd.io.load(h5_path) 18 | 19 | img = misc.imread(cur_img) 20 | classes = det_dict["classes"] 21 | coors = det_dict["coors"] 22 | ind = 0 23 | 24 | for label, coor in zip(classes, coors): 25 | ind += 1 26 | # thumb = img[coor[1]:coor[3], coor[0]:coor[2]] 27 | 28 | x_len = coor[2] - coor[0] 29 | x_len_expand = x_len * (1 + expand) 30 | x_mid = (coor[2] + coor[0]) / 2 31 | x_start = int(x_mid - x_len_expand / 2.0) 32 | x_end = int(x_mid + x_len_expand / 2.0) 33 | x_start = 0 if x_start < 0 else x_start 34 | x_end = 2560 if x_end > 2560 else x_end 35 | 36 | y_len = coor[3] - coor[1] 37 | y_len_expand = y_len * (1 + expand) 38 | y_mid = (coor[3] + coor[1]) / 2 39 | y_start = int(y_mid - y_len_expand / 2.0) 40 | y_end = int(y_mid + y_len_expand / 2.0) 41 | y_start = 0 if y_start < 0 else y_start 42 | y_end = 2048 if y_end > 2048 else y_end 43 | 44 | thumb = img[y_start:y_end, x_start:x_end] 45 | thumb = misc.imresize(thumb, (299, 299)) 46 | save_path = os.path.join(results_dir, str(label), cur_name + '_' + str(ind) + '.png') 47 | misc.imsave(save_path, thumb) 48 | 49 | if __name__ == "__main__": 50 | raw_img_dir = "../../data/DetKneeData/test" 51 | det_result_dir = "../../data/DetKneeData/det_results" 52 | auto_test_dir = "../../data/DetKneeData/automatic_test299" 53 | 54 | extract_detected_knees(raw_img_dir, det_result_dir, auto_test_dir) 55 | -------------------------------------------------------------------------------- /DetJoint/preprocess/save_det_h5.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | import os, sys, pdb 4 | import numpy as np 5 | import h5py, json 6 | from scipy.io import loadmat 7 | import deepdish as dd 8 | import scipy.misc as misc 9 | import numpy as np 10 | 11 | FILE_PATH = os.path.abspath(__file__) 12 | PRJ_PATH = os.path.dirname(os.path.dirname(FILE_PATH)) 13 | sys.path.append(PRJ_PATH) 14 | from yolo_v2.proj_utils.local_utils import getfileinfo, mkdirs 15 | from yolo_v2.proj_utils.local_utils import writeImg, imread, imresize 16 | 17 | 18 | # Load annotated bounding box mat file 19 | def load_mat(thismatfile, contourname_list=['Contours']): 20 | # First try load using h5py; then try using scipy.io.loadmat 21 | try: 22 | mat_file = h5py.File(thismatfile, 'r') 23 | for contourname in contourname_list: 24 | if contourname in list(mat_file.keys()): 25 | contour_mat = [np.transpose(mat_file[element[0]][:]) 26 | for element in mat_file[contourname]] 27 | break 28 | mat_file.close() 29 | except: 30 | loaded_mt = loadmat(thismatfile) 31 | for contourname in contourname_list: 32 | if contourname in loaded_mt.keys(): 33 | contour_mat = [] 34 | cnts = loaded_mt[contourname].tolist() 35 | if len(cnts) > 0: 36 | contour_mat = cnts[0] 37 | break 38 | 39 | return contour_mat 40 | 41 | 42 | def get_bbox(contour_mat): 43 | numCell = len(contour_mat) 44 | bbox_list = [] 45 | for icontour in range(0, numCell): 46 | thiscontour = contour_mat[icontour] 47 | xcontour = np.reshape(thiscontour[0,:], (1,-1) ) 48 | ycontour = np.reshape(thiscontour[1,:], (1,-1) ) 49 | 50 | x_min, x_max = np.min(xcontour), np.max(xcontour) 51 | y_min, y_max = np.min(ycontour), np.max(ycontour) 52 | bbox_list.append([x_min, y_min, x_max, y_max]) 53 | return bbox_list 54 | 55 | 56 | # Resize bounding box to a certain ratio 57 | def resize_mat(contour_mat, resize_ratio): 58 | numCell = len(contour_mat) 59 | res_contour = [] 60 | for icontour in range(0, numCell): 61 | thiscontour = contour_mat[icontour] 62 | res_contour.append([ele*resize_ratio for ele in thiscontour] ) 63 | return res_contour 64 | 65 | 66 | def save_h5(data_root, h5_dir, mode, ratio=0.125): 67 | all_dict_list = getfileinfo(os.path.join(data_root, mode), ['_gt'], ['.png'], '.mat') 68 | mode_dir = os.path.join(data_root, h5_dir, mode+"H5") 69 | mkdirs(mode_dir, erase=True) 70 | 71 | klg_dict = {} 72 | with open(os.path.join(data_root, mode, mode + "_kl.json"), 'r') as f: 73 | kl_dict = json.load(f) 74 | 75 | for ele in all_dict_list: 76 | info_dict = {} 77 | 78 | img_path = ele["thisfile"] 79 | cur_img = imread(img_path) 80 | mat_path = ele["thismatfile"] 81 | contour_mat = load_mat(mat_path) 82 | cur_bbox = get_bbox(contour_mat) 83 | assert len(cur_bbox) == 2, "Error, there are not 2 bbox" 84 | pat_id= os.path.splitext(os.path.basename(img_path))[0] 85 | key_r, key_l = pat_id + "R", pat_id + "L" 86 | 87 | 88 | resized_img = imresize(cur_img, ratio) 89 | assert resized_img.shape[:2] == (256, 320) 90 | if cur_bbox[0][0] > cur_bbox[1][0]: 91 | cur_bbox = cur_bbox[::-1] 92 | resized_bbox = resize_mat(cur_bbox, ratio) 93 | classes = [kl_dict[key_r], kl_dict[key_l]] 94 | 95 | info_dict["images"] = resized_img 96 | info_dict["gt_boxes"] = resized_bbox 97 | info_dict["gt_classes"] = classes 98 | info_dict["dontcare"] = [0] * len(classes) 99 | info_dict["origin_im"] = pat_id 100 | 101 | dd.io.save(os.path.join(mode_dir, pat_id+".h5"), info_dict) 102 | 103 | 104 | 105 | if __name__ == "__main__": 106 | np.random.seed(1234) 107 | 108 | knee_data_root = "../../data" 109 | data_root = os.path.join(knee_data_root, "DetKnee") 110 | h5_dir = "H5" 111 | 112 | modes = ["train", "val", "test"] 113 | for mode in modes: 114 | print("Save {} dataset".format(mode)) 115 | save_h5(data_root, h5_dir, mode) 116 | -------------------------------------------------------------------------------- /DetJoint/preprocess/split_train_val_test.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | import os, sys, pdb 4 | import numpy as np 5 | import shutil, json 6 | 7 | FILE_PATH = os.path.abspath(__file__) 8 | PRJ_PATH = os.path.dirname(os.path.dirname(FILE_PATH)) 9 | sys.path.append(PRJ_PATH) 10 | from yolo_v2.proj_utils.local_utils import getfileinfo, mkdirs 11 | 12 | 13 | def build_kl_dict(klg_json_path): 14 | kl_dict = {} 15 | with open(klg_json_path, 'r') as f: 16 | kl_dict = json.load(f) 17 | 18 | cls_dict = { 19 | "0": [], 20 | "1": [], 21 | "2": [], 22 | "3": [], 23 | "4": []} 24 | 25 | for k, v in kl_dict.items(): 26 | # organize patients according to the R leg's kl grade 27 | if k.endswith("L"): 28 | continue 29 | cls_dict[str(int(v))].append(k[:-1]) 30 | 31 | for k, v in cls_dict.items(): 32 | np.random.shuffle(cls_dict[k]) 33 | 34 | return cls_dict, kl_dict 35 | 36 | def split_train_val_test(save_dir, raw_img_dir, cls_dict, kl_dict): 37 | all_dict_list = getfileinfo(raw_img_dir, ['_gt'], ['.png'], '.mat') 38 | cls_names = ["train", "val", "test"] 39 | 40 | for mode in cls_names: 41 | cur_dir = os.path.join(save_dir, mode) 42 | mkdirs(cur_dir, erase=True) 43 | json_path = os.path.join(cur_dir, mode+"_kl.json") 44 | mode_klg_dict = {} 45 | mode_dict = {} 46 | for k, v in cls_dict.items(): 47 | cur_list = cls_dict[k] 48 | if mode == "train": 49 | end_index = int(len(cur_list) * 0.7) 50 | mode_dict[k] = cls_dict[k][:end_index] 51 | elif mode == "val": 52 | start_index = int(len(cur_list) * 0.7) 53 | end_index = int(len(cur_list) * 0.8) 54 | mode_dict[k] = cls_dict[k][start_index:end_index] 55 | elif mode == "test": 56 | start_index = int(len(cur_list) * 0.8) 57 | mode_dict[k] = cls_dict[k][start_index:] 58 | else: 59 | raise Exception("Unknow mode") 60 | 61 | for k, v in mode_dict.items(): 62 | for ele in v: 63 | key_left, key_right = ele + "L", ele + "R" 64 | mode_klg_dict[key_left] = int(kl_dict[key_left]) 65 | mode_klg_dict[key_right] = int(kl_dict[key_right]) 66 | shutil.copy(os.path.join(raw_img_dir, ele + ".png"), cur_dir) 67 | shutil.copy(os.path.join(raw_img_dir, ele + "_gt.mat"), cur_dir) 68 | 69 | with open(json_path, 'w') as fp: 70 | json.dump(mode_klg_dict, fp) 71 | 72 | if __name__ == "__main__": 73 | np.random.seed(1234) 74 | data_root = "../../data" 75 | det_data_name = "DetKnee" 76 | klg_fname = "baselineKLG.json" 77 | 78 | det_data_dir = os.path.join(data_root, det_data_name) 79 | cls_dict, kl_dict = build_kl_dict(os.path.join(det_data_dir, klg_fname)) 80 | img_annotations_dir = os.path.join(det_data_dir, "img_annotations") 81 | split_train_val_test(det_data_dir, img_annotations_dir, cls_dict, kl_dict) 82 | -------------------------------------------------------------------------------- /DetJoint/test/select_overlay.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | import os, sys, pdb 4 | import deepdish as dd 5 | import json 6 | 7 | FILE_PATH = os.path.abspath(__file__) 8 | PRJ_PATH = os.path.dirname(os.path.dirname(FILE_PATH)) 9 | sys.path.append(PRJ_PATH) 10 | 11 | from yolo_v2.knee_utils import overlay_bbox_iou 12 | from yolo_v2.proj_utils.local_utils import writeImg 13 | 14 | if __name__ == '__main__': 15 | # selection = ["9559547", "9970801", "9023935", "9488966", "9994363", "9103811", "9313330", "9607454"] 16 | select_json_path = "../../data/DetKneeData/results/selection.json" 17 | h5dir = "../../data/DetKneeData/H5/testH5" 18 | 19 | # load json file 20 | json_dict = json.load(open(select_json_path)) 21 | for ele in json_dict: 22 | h5_file = os.path.join(h5dir, ele+".h5") 23 | cur_item = dd.io.load(h5_file) 24 | cur_img = cur_item["images"] 25 | cur_det_boxes = json_dict[ele]["det"] 26 | cur_gt_boxes = json_dict[ele]["gt"] 27 | 28 | overlaid_img = overlay_bbox_iou(cur_img, cur_det_boxes, cur_gt_boxes) 29 | 30 | writeImg(overlaid_img, os.path.join(os.path.dirname(select_json_path), ele+'.png')) 31 | -------------------------------------------------------------------------------- /DetJoint/test/test_knee.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | import os, sys, pdb 4 | import torch 5 | import torchvision.transforms as standard_transforms 6 | import argparse 7 | from torch.utils import data 8 | 9 | FILE_PATH = os.path.abspath(__file__) 10 | PRJ_PATH = os.path.dirname(os.path.dirname(FILE_PATH)) 11 | sys.path.append(PRJ_PATH) 12 | 13 | from yolo_v2.proj_utils.local_utils import mkdirs 14 | from yolo_v2.darknet import Darknet19 15 | from yolo_v2.cfgs.config_knee import cfg 16 | from yolo_v2.datasets.knee import Knee 17 | from yolo_v2.test_yolo import test_eng 18 | 19 | 20 | def set_args(): 21 | parser = argparse.ArgumentParser(description = 'Testing code for Knee bone detection') 22 | parser.add_argument('--device-id', type=int, default=0) 23 | parser.add_argument('--batch-size', type=int, default=1) 24 | parser.add_argument('--model-dir', type=str, default="best_models") 25 | parser.add_argument('--model-name', type=str, default="det-epoch-500-0.000285.pth") 26 | parser.add_argument('--whichset', type=str, default="test") 27 | args = parser.parse_args() 28 | 29 | return args 30 | 31 | 32 | if __name__ == '__main__': 33 | # Config arguments 34 | args = set_args() 35 | 36 | # Setting model, testing data and result path 37 | data_root = "../../data/DetKneeData" 38 | model_root = os.path.join(data_root, args.model_dir) 39 | save_root = os.path.join(data_root, "results") 40 | mkdirs(save_root, erase=True) 41 | 42 | # Dataloader setting 43 | input_transform = standard_transforms.Compose([ 44 | standard_transforms.ToTensor(), 45 | standard_transforms.Normalize(cfg.rgb_mean, cfg.rgb_var)]) 46 | dataset = Knee(data_root, args.whichset, transform=input_transform) 47 | dataloader = data.DataLoader(dataset, batch_size=args.batch_size) 48 | # Set Darknet 49 | net = Darknet19(cfg) 50 | 51 | cuda_avail = torch.cuda.is_available() 52 | if cuda_avail: 53 | print("CUDA {} in use".format(args.device_id)) 54 | net.cuda(args.device_id) 55 | import torch.backends.cudnn as cudnn 56 | cudnn.benchmark = True 57 | else: 58 | print("CPU mode.") 59 | 60 | print('>> START testing ') 61 | if args.model_name == "": 62 | model_names = sorted([name for name in os.listdir(model_root) if name.endswith(".pth")]) 63 | for cur_name in model_names: 64 | args.model_name = cur_name 65 | test_eng(dataloader, model_root, save_root, net, args, cfg) 66 | else: 67 | model_root = os.path.join(data_root, args.model_dir) 68 | test_eng(dataloader, model_root, save_root, net, args, cfg) 69 | -------------------------------------------------------------------------------- /DetJoint/test/test_knee.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | python test_knee.py --device-id 1 \ 4 | --model-dir best_models/2wd \ 5 | --model-name det-epoch-490-0.000088.pth 6 | -------------------------------------------------------------------------------- /DetJoint/train/cal_mean_var.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | import os, sys, pdb 4 | import numpy as np 5 | 6 | FILE_PATH = os.path.abspath(__file__) 7 | PRJ_PATH = os.path.dirname(os.path.dirname(FILE_PATH)) 8 | sys.path.append(PRJ_PATH) 9 | from yolo_v2.datasets.knee import Knee 10 | 11 | if __name__ == '__main__': 12 | data_root = "../../data/DetKneeData" 13 | train_set = Knee(data_root, "train") 14 | 15 | pixel_mean = train_set.get_mean_pixel() 16 | pixel_var = train_set.get_var_pixel() 17 | print("Mean gray value is {}, variance is {}".format(pixel_mean, pixel_var)) 18 | -------------------------------------------------------------------------------- /DetJoint/train/cluster_anchor_box.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | import os, sys, pdb 4 | import numpy as np 5 | from sklearn.cluster import KMeans 6 | 7 | FILE_PATH = os.path.abspath(__file__) 8 | PRJ_PATH = os.path.dirname(os.path.dirname(FILE_PATH)) 9 | sys.path.append(PRJ_PATH) 10 | from yolo_v2.datasets.knee import Knee 11 | 12 | 13 | if __name__ == '__main__': 14 | data_root = "../../data/DetKneeData" 15 | train_set = Knee(data_root, "train") 16 | all_bbox = train_set.get_all_bbox() 17 | 18 | # [x1, y1, x2, y2] 19 | col_size_list = all_bbox[:,2] - all_bbox[:,0] 20 | row_size_list = all_bbox[:,3] - all_bbox[:,1] 21 | dim_array = np.stack([col_size_list, row_size_list], axis=1) 22 | 23 | # Clustered centered image dims 24 | cluster_num = 6 25 | # cluster_num = 5 26 | 27 | kmeans = KMeans(n_clusters=cluster_num, random_state=0).fit(dim_array) 28 | centers = kmeans.cluster_centers_ 29 | print("{} clustered centers:".format(cluster_num)) 30 | for center in centers: 31 | print("W: {:.2f}, H: {:.2f}, ratio: {:.2f}".format( 32 | center[0], center[1], center[1]/center[0])) 33 | -------------------------------------------------------------------------------- /DetJoint/train/train_knee.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | import os, sys, pdb 4 | import argparse 5 | import torch 6 | from torch.utils import data 7 | import torchvision.transforms as standard_transforms 8 | import numpy as np 9 | 10 | FILE_PATH = os.path.abspath(__file__) 11 | PRJ_PATH = os.path.dirname(os.path.dirname(FILE_PATH)) 12 | sys.path.append(PRJ_PATH) 13 | 14 | from yolo_v2.proj_utils.local_utils import mkdirs 15 | from yolo_v2.cfgs.config_knee import cfg 16 | from yolo_v2.darknet import Darknet19 17 | from yolo_v2.datasets.knee import Knee 18 | from yolo_v2.train_yolo import train_eng 19 | 20 | 21 | def set_args(): 22 | # Arguments settinge 23 | parser = argparse.ArgumentParser(description="Knee Bone Detection") 24 | parser.add_argument('--batch_size', type=int, default=8, help='batch size.') 25 | parser.add_argument('--maxepoch', type=int, default=500, help='number of epochs to train') 26 | parser.add_argument('--lr', type=float, default=2.0e-4, help='learning rate') 27 | parser.add_argument('--lr_decay', type=float, default=0.8, help='learning rate') 28 | parser.add_argument('--lr_decay_epochs', type=list, default=[60, 120, 180, 240, 300, 360, 420, 480], 29 | help='decay the learning rate at this epoch') 30 | parser.add_argument('--weight_decay', type=float, default=0.0, help='weight decay for training') 31 | parser.add_argument('--momentum', type=float, default=0.9, help='SGD momentum (default: 0.9)') 32 | parser.add_argument('--display_freq', type=int, default=10, help='plot the results per batches') 33 | parser.add_argument('--save_freq', type=int, default=10, help='how frequent to save the model') 34 | parser.add_argument('--device-id', type=int, default=0) 35 | parser.add_argument('--model-name', type=str, default='kneedet') 36 | parser.add_argument('--seed', type=int, default=1234) 37 | 38 | args = parser.parse_args() 39 | return args 40 | 41 | 42 | if __name__ == '__main__': 43 | args = set_args() 44 | np.random.seed(args.seed) 45 | 46 | # Data and Model settings 47 | data_root = "../../data/DetKneeData" 48 | model_root = os.path.join(data_root, args.model_name) 49 | mkdirs(model_root, erase=True) 50 | 51 | # Replace as mean and std 52 | input_transform = standard_transforms.Compose([ 53 | standard_transforms.ToTensor(), 54 | standard_transforms.Normalize(cfg.rgb_mean, cfg.rgb_var)]) 55 | train_dataset = Knee(data_root, "train", transform=input_transform) 56 | train_dataloader = data.DataLoader(train_dataset, batch_size=args.batch_size) 57 | val_dataset = Knee(data_root, "val", transform=input_transform) 58 | val_dataloader = data.DataLoader(val_dataset, batch_size=args.batch_size) 59 | 60 | # Set Darknet 61 | net = Darknet19(cfg) 62 | 63 | # CUDA Settings 64 | cuda_avail = torch.cuda.is_available() 65 | print("\n==== Starting training ====\n" + "===="*20) 66 | if cuda_avail: 67 | print("CUDA {} in use".format(args.device_id)) 68 | net.cuda(args.device_id) 69 | import torch.backends.cudnn as cudnn 70 | cudnn.benchmark = True 71 | else: 72 | print("CPU in use") 73 | 74 | # print ('>> START training ') 75 | train_eng(train_dataloader, val_dataloader, model_root, net, args) 76 | -------------------------------------------------------------------------------- /DetJoint/train/train_knee.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # model01 4 | python train_knee.py --device-id 7 \ 5 | --maxepoch 500 \ 6 | --weight_decay 5.0e-4 \ 7 | --batch_size 8 \ 8 | --model-name kneedet02wd 9 | -------------------------------------------------------------------------------- /DetJoint/yolo_v2/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PingjunChen/GradingKneeOA/dc19c2462111c3c5b27066e738fc11ddf57a6407/DetJoint/yolo_v2/__init__.py -------------------------------------------------------------------------------- /DetJoint/yolo_v2/cfgs/__init__.py: -------------------------------------------------------------------------------- 1 | from . import config_knee 2 | -------------------------------------------------------------------------------- /DetJoint/yolo_v2/cfgs/config_knee.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | import os, sys 4 | import numpy as np 5 | 6 | class config: 7 | def __init__(self): 8 | self.label_names = ["0", "1", "2", "3", "4"] 9 | self.num_classes = len(self.label_names) 10 | 11 | # w * h # 1 12 | self.anchors = np.asarray([[ 67.05, 65.61],]) 13 | # # w * h # 2 14 | # self.anchors = np.asarray([[ 63.10, 60.43], 15 | # [ 71.71, 71.71]]) 16 | # # w * h # 3 17 | # self.anchors = np.asarray([[ 65.3, 67.7], 18 | # [ 62.4, 57.5], 19 | # [ 75.0, 72.3]]) 20 | # # w * h # 4 21 | # self.anchors = np.asarray([[61.29, 57.40], 22 | # [74.87, 75.23], 23 | # [71.61, 63.60], 24 | # [63.74, 69.12]]) 25 | # # w * h # 5 26 | # self.anchors = np.asarray([[72.31, 63.91], 27 | # [76.06, 75.22], 28 | # [60.86, 55.39], 29 | # [63.01, 63.96], 30 | # [65.51, 73.18]]) 31 | # # w * h # 6 32 | # self.anchors = np.asarray([[74.13, 67.21], 33 | # [67.81, 60.29], 34 | # [65.62, 73.95], 35 | # [76.32, 77.20], 36 | # [62.49, 65.32], 37 | # [59.94, 55.40]]) 38 | self.num_anchors = len(self.anchors) 39 | 40 | # Image mean and std 41 | # self.rgb_mean = [0.5, 0.5, 0.5] 42 | # self.rgb_var = [0.5, 0.5, 0.5] 43 | self.rgb_mean = [0.431, 0.431, 0.431] 44 | self.rgb_var = [0.0872, 0.0872, 0.0872] 45 | 46 | # IOU scale 47 | self.iou_thresh = 0.6 48 | self.object_scale = 5. 49 | self.noobject_scale = 1. 50 | 51 | # BOX scale 52 | self.coord_scale = 1. 53 | 54 | # CLS scale 55 | # self.class_scale = 1.0 # for kl classification 56 | self.class_scale = 0.0 # only for detection 57 | 58 | # Regulable Ordinal Loss 59 | self.w_loss = False # for weighted loss 60 | 61 | # Test JI index 62 | self.JIthresh = 0.75 63 | # self.JIthresh = 0.50 64 | 65 | cfg = config() 66 | -------------------------------------------------------------------------------- /DetJoint/yolo_v2/darknet.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | import os, sys, pdb 4 | import numpy as np 5 | 6 | import torch 7 | import torch.nn as nn 8 | import torch.nn.functional as F 9 | from torch.multiprocessing import Pool 10 | from functools import partial 11 | 12 | from .utils import network as net_utils 13 | from .utils.cython_bbox import bbox_ious, anchor_intersections 14 | from .utils.cython_yolo import yolo_to_bbox 15 | from .proj_utils.model_utils import match_tensor 16 | from .proj_utils.torch_utils import to_device 17 | from .loss_utils import set_weights 18 | 19 | 20 | def _process_batch(inputs, size_spec=None, cfg=None): 21 | inp_size, out_size = size_spec 22 | x_ratio, y_ratio = float(inp_size[1])/out_size[1], float(inp_size[0])/out_size[0] 23 | bbox_pred_np, gt_boxes, gt_classes, iou_pred_np = inputs 24 | 25 | # net output 26 | hw, num_anchors, _ = bbox_pred_np.shape 27 | 28 | # gt 29 | _classes = np.zeros([hw, num_anchors, cfg.num_classes], dtype=np.float) 30 | _class_mask = np.zeros([hw, num_anchors, 1], dtype=np.float) 31 | 32 | _ious = np.zeros([hw, num_anchors, 1], dtype=np.float) 33 | _iou_mask = np.zeros([hw, num_anchors, 1], dtype=np.float) 34 | 35 | _boxes = np.zeros([hw, num_anchors, 4], dtype=np.float) 36 | _boxes[:, :, 0:2] = 0.5 # center of a box 37 | _boxes[:, :, 2:4] = 2.0 # avg(anchor / 32) 38 | _box_mask = np.zeros([hw, num_anchors, 1], dtype=np.float) + 0.001 39 | 40 | # Scale pred_bbox to real positions in image 41 | anchors = np.ascontiguousarray(cfg.anchors, dtype=np.float) 42 | bbox_pred_np = np.expand_dims(bbox_pred_np, 0) 43 | bbox_np = yolo_to_bbox( 44 | np.ascontiguousarray(bbox_pred_np, dtype=np.float), 45 | anchors, 46 | out_size[0], out_size[1], 47 | x_ratio, y_ratio) 48 | # for each prediction, calculate in 8*10 with all corresponsding anchors. 49 | bbox_np = bbox_np[0] # bbox_np = (hw, num_anchors, (x1, y1, x2, y2)) 50 | 51 | # gt_boxes_b = np.asarray(gt_boxes[b], dtype=np.float) 52 | gt_boxes_b = np.asarray(gt_boxes, dtype=np.float) 53 | # for each cell, compare predicted_bbox and gt_bbox 54 | bbox_np_b = np.reshape(bbox_np, [-1, 4]) 55 | ious = bbox_ious( 56 | np.ascontiguousarray(bbox_np_b, dtype=np.float), 57 | np.ascontiguousarray(gt_boxes_b, dtype=np.float)) 58 | 59 | # for each assumed box, find the best-matched with ground-truth 60 | best_ious = np.max(ious, axis=1).reshape(_iou_mask.shape) 61 | iou_penalty = 0 - iou_pred_np[best_ious < cfg.iou_thresh] 62 | _iou_mask[best_ious <= cfg.iou_thresh] = cfg.noobject_scale * iou_penalty 63 | 64 | # Locate the cell of each ground-truth Box 65 | cx = (gt_boxes_b[:, 0] + gt_boxes_b[:, 2]) * 0.5 / x_ratio 66 | cy = (gt_boxes_b[:, 1] + gt_boxes_b[:, 3]) * 0.5 / y_ratio 67 | cell_inds = np.floor(cy) * out_size[1] + np.floor(cx) 68 | cell_inds = cell_inds.astype(np.int) 69 | # transfer ground-truth box to 8*10 format 70 | target_boxes = np.empty(gt_boxes_b.shape, dtype=np.float) 71 | target_boxes[:, 0] = cx - np.floor(cx) # cx 72 | target_boxes[:, 1] = cy - np.floor(cy) # cy 73 | target_boxes[:, 2] = (gt_boxes_b[:, 2] - gt_boxes_b[:, 0]) # / inp_size[0] * out_size[0] # tw 74 | target_boxes[:, 3] = (gt_boxes_b[:, 3] - gt_boxes_b[:, 1]) # / inp_size[1] * out_size[1] # th 75 | 76 | # For each gt boxes, locate the best matching anchor points 77 | anchor_ious = anchor_intersections( 78 | anchors, np.ascontiguousarray(gt_boxes_b, dtype=np.float)) 79 | anchor_inds = np.argmax(anchor_ious, axis=0) 80 | ious_reshaped = np.reshape(ious, [hw, num_anchors, len(cell_inds)]) 81 | 82 | cls_weights = set_weights() 83 | for i, cell_ind in enumerate(cell_inds): 84 | a = anchor_inds[i] 85 | 86 | # IOU mask 87 | iou_pred_cell_anchor = iou_pred_np[cell_ind, a, :] # 0 ~ 1, should be close to 1 88 | _iou_mask[cell_ind, a, :] = cfg.object_scale * (1 - iou_pred_cell_anchor) 89 | _ious[cell_ind, a, :] = ious_reshaped[cell_ind, a, i] 90 | 91 | # BOX mask 92 | _box_mask[cell_ind, a, :] = cfg.coord_scale 93 | target_boxes[i, 2:4] /= anchors[a] 94 | _boxes[cell_ind, a, :] = target_boxes[i] 95 | 96 | # Classification mask 97 | _class_mask[cell_ind, a, :] = cfg.class_scale 98 | 99 | ## switch one-hot vector to weighted vector 100 | if cfg.w_loss == True: 101 | _classes[cell_ind, a, :] = cls_weights[gt_classes[i], :] 102 | else: 103 | _classes[cell_ind, a, gt_classes[i]] = 1. 104 | 105 | return _boxes, _ious, _classes, _box_mask, _iou_mask, _class_mask 106 | 107 | 108 | class Darknet19(nn.Module): 109 | def __init__(self, cfg): 110 | super(Darknet19, self).__init__() 111 | self.cfg = cfg 112 | self.register_buffer('device_id', torch.IntTensor(1)) 113 | 114 | net_cfgs = [ 115 | # conv1s 116 | [(32, 3)], 117 | ['M', (64, 3)], 118 | ['M', (128, 3), (64, 1), (128, 3)], 119 | ['M', (256, 3), (128, 1), (256, 3)], 120 | ['M', (512, 3), (256, 1), (512, 3), (256, 1), (512, 3)], 121 | # conv2 122 | ['M', (1024, 3), (512, 1), (1024, 3), (512, 1), (1024, 3)], 123 | # ------------ 124 | # conv3 125 | [(1024, 3), (1024, 3)], 126 | # conv4 127 | [(1024, 3)] 128 | ] 129 | 130 | # darknet 131 | self.conv1s, c1 = _make_layers(3, net_cfgs[0:5]) 132 | self.conv2, c2 = _make_layers(c1, net_cfgs[5]) 133 | # --- 134 | self.conv3, c3 = _make_layers(c2, net_cfgs[6]) 135 | 136 | stride = 2 137 | self.reorg = Reorg(stride=stride) 138 | # stride*stride times the channels of conv1s, then cat [conv1s, conv3] 139 | self.conv4, c4 = _make_layers((c1*(stride*stride) + c3), net_cfgs[7]) 140 | # linear 141 | out_channels = cfg.num_anchors * (cfg.num_classes + 5) 142 | self.conv5 = net_utils.Conv2d(c4, out_channels, 1, 1, relu=False) 143 | # train 144 | self.bbox_loss = None 145 | self.iou_loss = None 146 | self.cls_loss = None 147 | self.pool = Pool(processes=8) 148 | 149 | @property 150 | def loss(self): 151 | return self.bbox_loss + self.iou_loss + self.cls_loss 152 | 153 | def forward(self, im_data, gt_boxes=None, gt_classes=None): 154 | self.inp_size = im_data.size()[2:] # 256*320 155 | conv1s = self.conv1s(im_data) 156 | conv2 = self.conv2(conv1s) 157 | conv3 = self.conv3(conv2) 158 | conv1s_reorg = self.reorg(conv1s, conv3.size()[2::]) 159 | cat_1_3 = torch.cat([conv1s_reorg, conv3], 1) 160 | conv4 = self.conv4(cat_1_3) 161 | 162 | self.out_size = conv4.size()[2:] # 8*10 163 | self.x_ratio = float(self.inp_size[1])/self.out_size[1] # 32 164 | self.y_ratio = float(self.inp_size[0])/self.out_size[0] # 32 165 | 166 | conv5 = self.conv5(conv4) # batch_size, out_channels, h, w 167 | 168 | # for detection 169 | # bsize, c, h, w -> bsize, h, w, c -> bsize, h x w, num_anchors, 5+num_classes 170 | bsize, _, h, w = conv5.size() 171 | # assert bsize == 1, 'detection only support one image per batch' 172 | conv5_reshaped = conv5.permute(0, 2, 3, 1).contiguous().view(bsize, -1, self.cfg.num_anchors, self.cfg.num_classes + 5) 173 | 174 | # tx, ty, tw, th, to -> sig(tx), sig(ty), exp(tw), exp(th), sig(to) 175 | xy_pred = F.sigmoid(conv5_reshaped[:, :, :, 0:2]) 176 | wh_pred = torch.exp(conv5_reshaped[:, :, :, 2:4]) 177 | bbox_pred = torch.cat([xy_pred, wh_pred], 3) 178 | iou_pred = F.sigmoid(conv5_reshaped[:, :, :, 4:5]) 179 | score_pred = conv5_reshaped[:, :, :, 5:].contiguous() 180 | prob_pred = F.softmax(score_pred.view(-1, score_pred.size()[-1]), dim=-1).view_as(score_pred) 181 | 182 | # for training 183 | if self.training: 184 | bbox_pred_np = bbox_pred.data.cpu().numpy() 185 | gt_boxes_np = gt_boxes.numpy() 186 | gt_classes_np = gt_classes.numpy() 187 | iou_pred_np = iou_pred.data.cpu().numpy() 188 | 189 | _boxes, _ious, _classes, _box_mask, _iou_mask, _class_mask = self._build_target( 190 | bbox_pred_np, gt_boxes_np, gt_classes_np, iou_pred_np) 191 | 192 | _boxes = to_device(_boxes, self.device_id, requires_grad=False) 193 | _ious = to_device(_ious, self.device_id , requires_grad=False) 194 | _classes = to_device(_classes, self.device_id, requires_grad=False) 195 | box_mask = to_device(_box_mask, self.device_id, requires_grad=False) 196 | iou_mask = to_device(_iou_mask, self.device_id, requires_grad=False) 197 | class_mask = to_device(_class_mask, self.device_id, requires_grad=False) 198 | 199 | num_boxes = sum((len(boxes) for boxes in gt_boxes)) 200 | box_mask = box_mask.expand_as(_boxes) 201 | self.bbox_loss = nn.MSELoss(size_average=False)(bbox_pred * box_mask, _boxes * box_mask) / num_boxes 202 | self.iou_loss = nn.MSELoss(size_average=False)(iou_pred * iou_mask, _ious * iou_mask) / num_boxes 203 | 204 | class_mask = class_mask.expand_as(prob_pred) 205 | 206 | if self.cfg.w_loss == True: 207 | self.cls_loss = torch.sum((prob_pred * class_mask * _classes)**2) / num_boxes 208 | else: 209 | self.cls_loss = nn.MSELoss(size_average=False)(prob_pred * class_mask, _classes * class_mask) / num_boxes 210 | 211 | return bbox_pred, iou_pred, prob_pred 212 | 213 | 214 | def _build_target(self, bbox_pred_np, gt_boxes, gt_classes, iou_pred_np): 215 | """ 216 | :param bbox_pred_np: shape: (bsize, h x w, num_anchors, 4) : (sig(tx), sig(ty), exp(tw), exp(th)) 217 | """ 218 | bsize = bbox_pred_np.shape[0] 219 | _process_batch_func = partial(_process_batch, size_spec = (self.inp_size, self.out_size) ,cfg=self.cfg) 220 | targets = self.pool.map(_process_batch_func, ( (bbox_pred_np[b], gt_boxes[b], gt_classes[b], iou_pred_np[b])for b in range(bsize))) 221 | 222 | _boxes = np.stack(tuple((row[0] for row in targets))) 223 | _ious = np.stack(tuple((row[1] for row in targets))) 224 | _classes = np.stack(tuple((row[2] for row in targets))) 225 | _box_mask = np.stack(tuple((row[3] for row in targets))) 226 | _iou_mask = np.stack(tuple((row[4] for row in targets))) 227 | _class_mask = np.stack(tuple((row[5] for row in targets))) 228 | 229 | return _boxes, _ious, _classes, _box_mask, _iou_mask, _class_mask 230 | 231 | 232 | #### Auxiliary function and class for darknet ### 233 | def _make_layers(in_channels, net_cfg): 234 | layers = [] 235 | 236 | if len(net_cfg) > 0 and isinstance(net_cfg[0], list): 237 | for sub_cfg in net_cfg: 238 | layer, in_channels = _make_layers(in_channels, sub_cfg) 239 | layers.append(layer) 240 | else: 241 | for item in net_cfg: 242 | if item == 'M': 243 | layers.append(nn.MaxPool2d(kernel_size=2, stride=2)) 244 | else: 245 | out_channels, ksize = item 246 | layers.append(net_utils.Conv2d_BatchNorm(in_channels, out_channels, ksize, same_padding=True)) 247 | in_channels = out_channels 248 | 249 | return nn.Sequential(*layers), in_channels 250 | 251 | 252 | class Reorg(nn.Module): 253 | def __init__(self, stride=2): 254 | super(Reorg, self).__init__() 255 | self.stride = stride 256 | 257 | def forward(self, x, small_size): 258 | stride = self.stride 259 | 260 | x = match_tensor(x, (2*small_size[0], 2*small_size[1])) 261 | 262 | assert(x.data.dim() == 4) 263 | B = x.data.size(0) 264 | C = x.data.size(1) 265 | H = x.data.size(2) 266 | W = x.data.size(3) 267 | assert(H % stride == 0) 268 | assert(W % stride == 0) 269 | ws = int(stride) 270 | hs = int(stride) 271 | x = x.contiguous().view(B, C, H//hs, hs, W//ws, ws).transpose(3,4).contiguous() 272 | x = x.contiguous().view(B, C, H//hs*W//ws, hs*ws).transpose(2,3).contiguous() 273 | x = x.contiguous().view(B, C, hs*ws, H//hs, W//ws).transpose(1,2).contiguous() 274 | x = x.contiguous().view(B, hs*ws*C, H//hs, W//ws) 275 | return x 276 | -------------------------------------------------------------------------------- /DetJoint/yolo_v2/datasets/__init__.py: -------------------------------------------------------------------------------- 1 | from . import knee 2 | -------------------------------------------------------------------------------- /DetJoint/yolo_v2/datasets/knee.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | import os, sys, pdb 4 | import numpy as np 5 | import glob 6 | import deepdish as dd 7 | from torch.utils import data 8 | 9 | import torchvision.transforms as standard_transforms 10 | from ..proj_utils.local_utils import overlay_bbox 11 | 12 | class Knee(data.Dataset): 13 | def __init__(self, data_root, mode, transform=None): 14 | self.data_root = data_root 15 | assert mode in ["train", "val", "test", "most"], "Unknown mode: {}".format(mode) 16 | self.mode = mode 17 | self.data_path = os.path.join(data_root, "H5", mode + "H5") 18 | self.item_list = glob.glob(os.path.join(self.data_path, "*.h5")) 19 | self.item_num = len(self.item_list) 20 | self.transform = transform 21 | self._classes = ["0", "1", "2", "3", "4"] 22 | 23 | def __getitem__(self, index): 24 | # get bone info 25 | cur_item = dd.io.load(self.item_list[index]) 26 | # get image 27 | cur_img = cur_item["images"].astype(np.float32) 28 | if self.transform is not None: 29 | cur_img = self.transform(cur_img) 30 | 31 | 32 | cur_boxes = np.asarray(cur_item["gt_boxes"]) 33 | cur_classes = np.asarray(cur_item["gt_classes"]) 34 | cur_name = cur_item["origin_im"] 35 | 36 | return cur_img, cur_boxes, cur_classes, cur_name 37 | 38 | def __len__(self): 39 | return self.item_num 40 | 41 | @property 42 | def num_classes(self): 43 | return len(self._classes) 44 | 45 | @property 46 | def classes(self): 47 | return self._classes 48 | 49 | @property 50 | def item_names(self): 51 | return [os.path.basename(ele) for ele in self.item_list] 52 | 53 | @property 54 | def num_items(self): 55 | return self.item_num 56 | 57 | def get_all_bbox(self): 58 | bbox_all = [] 59 | for idx in range(self.item_num): 60 | cur_item = dd.io.load(self.item_list[idx]) 61 | bbox = cur_item["gt_boxes"] 62 | bbox_all.extend(bbox) 63 | 64 | return np.asarray(bbox_all) 65 | 66 | def get_mean_pixel(self): 67 | gray_list = np.zeros((self.item_num, ), np.float32) 68 | for idx in range(self.item_num): 69 | cur_item = dd.io.load(self.item_list[idx]) 70 | cur_img = cur_item["images"] 71 | gray_list[idx] = np.mean(cur_img) / 255.0 72 | 73 | pixel_mean = np.mean(gray_list) 74 | return pixel_mean 75 | 76 | def get_var_pixel(self): 77 | pixel_arr = np.zeros((self.item_num, 256, 320), np.float32) 78 | 79 | for idx in range(self.item_num): 80 | cur_item = dd.io.load(self.item_list[idx]) 81 | cur_img = cur_item["images"] 82 | 83 | var_img = cur_img[:, :, 0] / 255.0 84 | pixel_arr[idx] = var_img 85 | pixel_var = np.var(pixel_arr) 86 | return pixel_var 87 | 88 | def overlayImgs(self, save_path): 89 | for idx in range(self.item_num): 90 | cur_item = dd.io.load(self.item_list[idx]) 91 | cur_img = cur_item["images"] 92 | cur_boxes = cur_item["gt_boxes"] 93 | cur_classes = cur_item["gt_classes"] 94 | cur_name = cur_item["origin_im"] 95 | 96 | overlay_img = overlay_bbox(cur_img, cur_boxes, len=3).astype(np.uint8) 97 | 98 | for ind, box in enumerate(cur_boxes): 99 | font = cv2.FONT_HERSHEY_SIMPLEX 100 | color = (0, 255, 0) 101 | col_mid = int((box[2] + box[0]) / 2.0) 102 | row_mid = int((box[3] + box[1]) / 2.0) 103 | cv2.putText(overlay_img, str(cur_classes[ind]), (col_mid, row_mid), font, 1, color, 2) 104 | 105 | cv2.imwrite(os.path.join(save_path, cur_name+".png"), overlay_img) 106 | 107 | 108 | if __name__ == "__main__": 109 | data_root = "../../data/" 110 | input_transform = standard_transforms.Compose([ 111 | standard_transforms.ToTensor(), 112 | standard_transforms.Normalize([0.5]*3, [0.5]*3)]) 113 | val_dataset = Knee(data_root, "validation", transform=input_transform) 114 | val_dataloader = data.DataLoader(val_dataset, batch_size=4) 115 | 116 | epoch_num = 10 117 | for idx_epoch in range(epoch_num): 118 | for cur_batch, data in enumerate(val_dataloader): 119 | cur_img, cur_boxes, cur_classes, cur_name = data 120 | import pdb; pdb.set_trace() 121 | print("In Epoch: {}, Batch: {}".format(idx_epoch, cur_batch)) 122 | -------------------------------------------------------------------------------- /DetJoint/yolo_v2/knee_utils.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | import os, sys, pdb 4 | import numpy as np 5 | import cv2 6 | import deepdish as dd 7 | 8 | from .proj_utils.torch_utils import to_device 9 | from .utils.cython_yolo import yolo_to_bbox 10 | from .proj_utils.local_utils import change_val 11 | 12 | # knee detection and classification forward 13 | def knee_det_cls(cls, img, cfg=None): 14 | results = {'bbox':[],'iou':[], 'prob':[]} 15 | 16 | batch_data = to_device(img, cls.device_id, volatile=True) 17 | bbox_pred, iou_pred, prob_pred = cls.forward(batch_data) 18 | 19 | bbox_pred = bbox_pred.cpu().data.numpy() 20 | iou_pred = iou_pred.cpu().data.numpy() 21 | prob_pred = prob_pred.cpu().data.numpy() 22 | 23 | H, W = cls.out_size 24 | x_ratio, y_ratio = cls.x_ratio, cls.y_ratio 25 | 26 | bbox_pred = yolo_to_bbox( 27 | np.ascontiguousarray(bbox_pred, dtype=np.float), 28 | np.ascontiguousarray(cfg.anchors, dtype=np.float), 29 | H, W, x_ratio, y_ratio) 30 | 31 | results['bbox'] = np.array(bbox_pred) 32 | results['iou'] = np.array(iou_pred) 33 | results['prob'] = np.array(prob_pred) 34 | 35 | return results 36 | 37 | 38 | # Evaluate knee detection and classification 39 | def evaluate_det_cls(gt_boxes, gt_classes, pr_boxes, pr_classes, num=5, overlap_ratio=0.75): 40 | t_box_num = 0 41 | all_box_num = len(pr_boxes) 42 | pred_matrix = np.zeros((num, num), dtype=np.int) 43 | overlap_list = [] 44 | 45 | 46 | box_num = 0 47 | for gt_box, gt_class in zip(gt_boxes, gt_classes): 48 | for pr_box, pr_class in zip(pr_boxes, pr_classes): 49 | cur_overlap = evaluate_box_JI(gt_box, pr_box) 50 | # cur_overlap = evaluate_box_dice(gt_box, pr_box) 51 | if cur_overlap >= overlap_ratio: 52 | pred_matrix[gt_class, pr_class] += 1 53 | t_box_num += 1 54 | overlap_list.append(cur_overlap) 55 | if cur_overlap > 0.92 or (cur_overlap > 0.75 and cur_overlap < 0.78): 56 | box_num += 1 57 | 58 | if cur_overlap < 0.50 and cur_overlap > 0.4: 59 | box_num = 2 60 | 61 | return pred_matrix, t_box_num, all_box_num, overlap_list, box_num 62 | 63 | 64 | def overlay_bbox_iou(img, pred_boxes, gt_boxes, len=3): 65 | # Draw gt boxes 66 | for bb in gt_boxes: 67 | x_min_, y_min_, x_max_, y_max_ = bb 68 | x_min_, y_min_, x_max_, y_max_ = int(x_min_),int( y_min_), int(x_max_), int(y_max_) 69 | img[:,:,0] = change_val(img[:,:,0], 255, len, x_min_, y_min_, x_max_, y_max_) 70 | img[:,:,1] = change_val(img[:,:,1], 0, len, x_min_, y_min_, x_max_, y_max_) 71 | img[:,:,2] = change_val(img[:,:,2], 0, len, x_min_, y_min_, x_max_, y_max_) 72 | 73 | for bb in pred_boxes: 74 | x_min_, y_min_, x_max_, y_max_ = bb 75 | x_min_, y_min_, x_max_, y_max_ = int(x_min_),int( y_min_), int(x_max_), int(y_max_) 76 | img[:,:,0] = change_val(img[:,:,0], 0, len, x_min_, y_min_, x_max_, y_max_) 77 | img[:,:,1] = change_val(img[:,:,1], 255, len, x_min_, y_min_, x_max_, y_max_) 78 | img[:,:,2] = change_val(img[:,:,2], 0, len, x_min_, y_min_, x_max_, y_max_) 79 | 80 | max_iou = 0.0 81 | for gt_bb in gt_boxes: 82 | cur_iou = evaluate_box_JI(bb, gt_bb) 83 | # cur_iou = evaluate_box_dice(bb, gt_bb) 84 | if cur_iou > max_iou: 85 | max_iou = cur_iou 86 | 87 | text_loc = (x_min_, y_min_ - 10) 88 | img = cv2.putText(img.copy(), str(round(max_iou, 3)), text_loc, 89 | cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2) 90 | return img 91 | 92 | 93 | def evaluate_box_JI(gt_box, pr_box): 94 | if pr_box[0] >= gt_box[2] or pr_box[2] <= gt_box[0] or pr_box[1] >= gt_box[3] or pr_box[3] <= gt_box[1]: 95 | return 0.0 96 | else: 97 | gt_area = (gt_box[2] - gt_box[0]) * (gt_box[3] - gt_box[1]) 98 | pr_area = (pr_box[2] - pr_box[0]) * (pr_box[3] - pr_box[1]) 99 | overlap_start = (max(gt_box[0], pr_box[0]), max(gt_box[1], pr_box[1])) 100 | overlap_end = (min(gt_box[2], pr_box[2]), min(gt_box[3], pr_box[3])) 101 | 102 | area_numerator = (overlap_end[1] - overlap_start[1]) * (overlap_end[0] - overlap_start[0]) 103 | area_denominator = gt_area + pr_area - area_numerator 104 | overlap_ratio = area_numerator / area_denominator 105 | 106 | return overlap_ratio 107 | 108 | 109 | def evaluate_box_dice(gt_box, pr_box): 110 | if pr_box[0] >= gt_box[2] or pr_box[2] <= gt_box[0] or pr_box[1] >= gt_box[3] or pr_box[3] <= gt_box[1]: 111 | return 0.0 112 | else: 113 | gt_area = (gt_box[2] - gt_box[0]) * (gt_box[3] - gt_box[1]) 114 | pr_area = (pr_box[2] - pr_box[0]) * (pr_box[3] - pr_box[1]) 115 | overlap_start = (max(gt_box[0], pr_box[0]), max(gt_box[1], pr_box[1])) 116 | overlap_end = (min(gt_box[2], pr_box[2]), min(gt_box[3], pr_box[3])) 117 | 118 | area_numerator = 2.0 * (overlap_end[1] - overlap_start[1]) * (overlap_end[0] - overlap_start[0]) 119 | area_denominator = gt_area + pr_area 120 | overlap_ratio = area_numerator / area_denominator 121 | 122 | return overlap_ratio 123 | 124 | 125 | # Save detected knee bbox for classification and segmentation 126 | def save_pred_box_coors(save_dir, gt_boxes, gt_classes, pr_boxes, img_name, overlap_ratio=0.75): 127 | bone_dict = {"coors": [], "classes": []} 128 | for gt_box, gt_class in zip(gt_boxes, gt_classes): 129 | for ind, pr_box in enumerate(pr_boxes): 130 | if evaluate_box_JI(gt_box, pr_box) >= overlap_ratio: 131 | bone_dict["coors"].append(pr_box*8) 132 | bone_dict["classes"].append(gt_class) 133 | dd.io.save(os.path.join(save_dir, img_name+".h5"), bone_dict) 134 | -------------------------------------------------------------------------------- /DetJoint/yolo_v2/loss_utils.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | import os, sys, pdb 4 | import numpy as np 5 | 6 | def set_weights(): 7 | init_weights = np.array([[1, 3, 5, 7, 9], 8 | [3, 1, 3, 5, 7], 9 | [5, 3, 1, 3, 5], 10 | [7, 5, 3, 1, 3], 11 | [9, 7, 5, 3, 1]], dtype=np.float) 12 | 13 | adjusted_weights = init_weights + 1.0 14 | np.fill_diagonal(adjusted_weights, 0) 15 | 16 | return adjusted_weights 17 | -------------------------------------------------------------------------------- /DetJoint/yolo_v2/make.sh: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env bash 2 | 3 | CUDA_PATH=/usr/local/cuda/ 4 | 5 | cd utils 6 | python build.py build_ext --inplace 7 | cd ../ 8 | -------------------------------------------------------------------------------- /DetJoint/yolo_v2/proj_utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PingjunChen/GradingKneeOA/dc19c2462111c3c5b27066e738fc11ddf57a6407/DetJoint/yolo_v2/proj_utils/__init__.py -------------------------------------------------------------------------------- /DetJoint/yolo_v2/proj_utils/generic_utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import time 3 | import sys 4 | import six 5 | 6 | def get_from_module(identifier, module_params, module_name, 7 | instantiate=False, kwargs=None): 8 | if isinstance(identifier, six.string_types): 9 | res = module_params.get(identifier) 10 | if not res: 11 | raise Exception('Invalid ' + str(module_name) + ': ' + 12 | str(identifier)) 13 | if instantiate and not kwargs: 14 | return res() 15 | elif instantiate and kwargs: 16 | return res(**kwargs) 17 | else: 18 | return res 19 | elif type(identifier) is dict: 20 | name = identifier.pop('name') 21 | res = module_params.get(name) 22 | if res: 23 | return res(**identifier) 24 | else: 25 | raise Exception('Invalid ' + str(module_name) + ': ' + 26 | str(identifier)) 27 | return identifier 28 | -------------------------------------------------------------------------------- /DetJoint/yolo_v2/proj_utils/local_utils.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | import os, math, re 4 | import numpy as np 5 | import scipy 6 | import scipy.misc as misc 7 | import matplotlib.pyplot as plt 8 | from PIL import Image, ImageDraw 9 | from sklearn.neighbors import NearestNeighbors 10 | 11 | from sklearn.utils import shuffle 12 | from skimage import color, measure, morphology 13 | from numba import jit, autojit 14 | import random, shutil 15 | 16 | import cv2 17 | 18 | from .generic_utils import get_from_module 19 | 20 | 21 | def get(identifier): 22 | return get_from_module(identifier, globals(), 'local_utils') 23 | 24 | 25 | def mkdirs(folders, erase=False): 26 | if type(folders) is not list: 27 | folders = [folders] 28 | for fold in folders: 29 | if not os.path.exists(fold): 30 | os.makedirs(fold) 31 | else: 32 | if erase: 33 | shutil.rmtree(fold) 34 | os.makedirs(fold) 35 | 36 | # Replace pixel values with specific color 37 | def change_val(img,val, len, x_min, y_min, x_max, y_max): 38 | left_len = (len-1)//2 39 | right_len = (len-1) - left_len 40 | row_size, col_size = img.shape[0:2] 41 | for le in range(-left_len, right_len + 1): 42 | y_min_ = max(0, y_min + le ) 43 | x_min_ = max(0, x_min + le ) 44 | 45 | y_max_ = min(row_size, y_max - le ) 46 | x_max_ = min(col_size, x_max - le ) 47 | 48 | img[y_min_:y_max_, x_min_:x_min_+1] = val 49 | img[y_min_:y_min_+1, x_min_:x_max_] = val 50 | img[y_min_:y_max_, x_max_:x_max_+1] = val 51 | img[y_max_:y_max_+1, x_min_:x_max_] = val 52 | return img 53 | 54 | 55 | # Overlay bbox 56 | def overlay_bbox(img, bbox, len=1, rgb=(255, 0, 0)): 57 | for bb in bbox: 58 | x_min_, y_min_, x_max_, y_max_ = bb 59 | x_min_, y_min_, x_max_, y_max_ = int(x_min_),int( y_min_), int(x_max_), int(y_max_) 60 | img[:,:,0] = change_val(img[:,:,0], rgb[0], len, x_min_, y_min_, x_max_, y_max_) 61 | img[:,:,1] = change_val(img[:,:,1], rgb[1], len, x_min_, y_min_, x_max_, y_max_) 62 | img[:,:,2] = change_val(img[:,:,2], rgb[2], len, x_min_, y_min_, x_max_, y_max_) 63 | return img 64 | 65 | 66 | def writeImg(array, savepath): 67 | scipy.misc.imsave(savepath, array) 68 | 69 | 70 | def imread(imgfile): 71 | assert os.path.exists(imgfile), '{} does not exist!'.format(imgfile) 72 | srcBGR = cv2.imread(imgfile) 73 | destRGB = cv2.cvtColor(srcBGR, cv2.COLOR_BGR2RGB) 74 | return destRGB 75 | 76 | 77 | def imresize(img, resizeratio=1): 78 | '''Take care of cv2 reshape squeeze behevaior''' 79 | if resizeratio == 1: 80 | return img 81 | outshape = ( int(img.shape[1] * resizeratio) , int(img.shape[0] * resizeratio)) 82 | temp = cv2.resize(img, outshape).astype(float) 83 | #temp = misc.imresize(img, size=outshape).astype(float) 84 | if len(img.shape) == 3 and img.shape[2] == 1: 85 | temp = np.reshape(temp, temp.shape + (1,)) 86 | return temp 87 | 88 | 89 | def imshow(img, size=None): 90 | if size is not None: 91 | plt.figure(figsize = size) 92 | else: 93 | plt.figure() 94 | plt.imshow(img) 95 | plt.show() 96 | 97 | 98 | def normalize_img(X): 99 | min_, max_ = np.min(X), np.max(X) 100 | X = (X - min_)/ (max_ - min_ + 1e-9) 101 | X = X*255 102 | return X.astype(np.uint8) 103 | 104 | 105 | def Indexflow(Totalnum, batch_size, random=True): 106 | numberofchunk = int(Totalnum + batch_size - 1)// int(batch_size) # the floor 107 | #Chunkfile = np.zeros((batch_size, row*col*channel)) 108 | totalIndx = np.arange(Totalnum).astype(np.int) 109 | if random is True: 110 | totalIndx = np.random.permutation(totalIndx) 111 | 112 | chunkstart = 0 113 | for chunkidx in range(int(numberofchunk)): 114 | thisnum = min(batch_size, Totalnum - chunkidx*batch_size) 115 | thisInd = totalIndx[chunkstart: chunkstart + thisnum] 116 | chunkstart += thisnum 117 | yield thisInd 118 | 119 | 120 | def getfileinfo(imgdir, contourextList, ImgExtList, LabelExt, test_mode = False): 121 | '''return a list of dictionary {'thisfile':os.path.join(imgdir,f), 'thismatfile':thismatfile} 122 | ''' 123 | alllist = [f for f in os.listdir(imgdir)] 124 | alllist = sorted(alllist) 125 | 126 | returnList = [] 127 | for f in alllist: 128 | if os.path.isfile(os.path.join(imgdir,f)) and \ 129 | os.path.splitext(f)[1] in ImgExtList: 130 | if test_mode is False: 131 | flag = 0 132 | for contourext in contourextList: 133 | thismatfile = os.path.join(imgdir,os.path.splitext(f)[0] + contourext + LabelExt) 134 | if os.path.isfile(thismatfile): 135 | returnList.append({'thisfile':os.path.join(imgdir,f), 'thismatfile':thismatfile}) 136 | flag = 1 137 | break 138 | if flag == 0: 139 | print(("Image: {s} does not have matfile".format(s = os.path.splitext(f)[0] ))) 140 | else: 141 | returnList.append({'thisfile':os.path.join(imgdir,f), 'thismatfile': None}) 142 | return returnList 143 | 144 | 145 | def split_img(img, windowsize=1000, board = 0, fixed_window = False, step_size = None, tuple_slice = False): 146 | ''' 147 | img dimension: channel, row, col 148 | output: 149 | (IndexDict, PackList) 150 | IndexDict is a dictionry, the key is the actual patch size, the values is the list of identifier, 151 | PackList: list of (thisPatch,org_slice ,extract_slice, thisSize,identifier), the index of Packlist 152 | corresponds to the identifier. 153 | org_slice: slice coordinated at the orignal image. 154 | extract_slice: slice coordinate at the extract thisPatch, 155 | the length of org_slice should be equal to extract_slice. 156 | 157 | fixed_window: if true, it forces the extracted patches has to be of size window_size. 158 | we don't pad the original image to make mod(imgsize, windowsize)==0, instead, if the remaining is small, 159 | we expand the left board to lefter to compensate the smaller reminging patches. 160 | 161 | The default behavior is False: get all window_size patches, and collect the remining patches as it is. 162 | 163 | step_size: if step_size is smaller than (windowsize-2*board), we extract the patches with overlapping. 164 | which means the org_slice is overlapping. 165 | 166 | eg: 167 | lenght = 17 168 | img = np.arange(2*lenght*lenght).reshape(2,lenght,lenght) 169 | 170 | nm = np.zeros(img.shape).astype(np.int) 171 | 172 | AllDict, PackList = split_img(img, windowsize=7, board = 0, step_size= 2,fixed_window = True) 173 | 174 | print img 175 | 176 | print '---------------------------------------' 177 | 178 | print AllDict.keys() 179 | 180 | for key in AllDict.keys(): 181 | iden_list = AllDict[key] 182 | for iden in iden_list: 183 | thispatch = PackList[iden][0] 184 | org_slice = PackList[iden][1] 185 | extract_slice = PackList[iden][2] 186 | 187 | nm[:,org_slice[0],org_slice[1]] = thispatch[:,extract_slice[0],extract_slice[1]] 188 | print thispatch[:,extract_slice[0],extract_slice[1]] 189 | print nm 190 | print sum(nm-img) 191 | ''' 192 | IndexDict = {} 193 | identifier = -1 194 | PackList = [] 195 | row_size, col_size = img.shape[1], img.shape[2] 196 | if windowsize is not None and type(windowsize) is int: 197 | windowsize = (windowsize, windowsize) 198 | 199 | if windowsize is None or (row_size <= windowsize[0] and col_size<=windowsize[1] and (not fixed_window)): 200 | pad_img = img 201 | rowsize, colsize = pad_img.shape[1:] 202 | 203 | org_slice = (slice(0, rowsize), slice(0, colsize)) 204 | extract_slice = org_slice 205 | crop_patch_slice = (slice(0, rowsize), slice(0, colsize)) 206 | thisSize = (rowsize, colsize ) 207 | identifier = identifier + 1 208 | 209 | org_slice_tuple = (0, 0) 210 | if thisSize in IndexDict: 211 | IndexDict[thisSize].append(identifier) 212 | else: 213 | IndexDict[thisSize] = [] 214 | IndexDict[thisSize].append(identifier) 215 | PackList.append((crop_patch_slice, org_slice ,extract_slice, thisSize,identifier, org_slice_tuple)) 216 | 217 | else: 218 | 219 | hidden_windowsize = (windowsize[0]-2*board, windowsize[1]-2*board) 220 | for each_size in hidden_windowsize: 221 | if each_size <= 0: 222 | raise RuntimeError('windowsize can not be smaller than board*2.') 223 | 224 | if type(step_size) is int: 225 | step_size = (step_size, step_size) 226 | if step_size is None: 227 | step_size = hidden_windowsize 228 | 229 | numRowblocks = int(math.ceil(float(row_size)/step_size[0])) 230 | numColblocks = int(math.ceil(float(col_size)/step_size[1])) 231 | 232 | # sanity check, make sure the image is at least of size window_size to the left-hand side if fixed_windows is true 233 | # which means, -----*******|-----, left to the vertical board of original image is at least window_size. 234 | row_addition_board, col_addition_board = 0, 0 235 | addition_board = 0 236 | if fixed_window: 237 | if row_size + 2 * board < windowsize[0]: # means we need to add more on board. 238 | row_addition_board = windowsize[0] - (row_size + 2 * board ) 239 | if col_size + 2 * board < windowsize[1]: # means we need to add more on board. 240 | col_addition_board = windowsize[1] - (col_size + 2 * board) 241 | addition_board = row_addition_board if row_addition_board > col_addition_board else col_addition_board 242 | 243 | left_pad = addition_board + board 244 | pad4d = ((0,0),( left_pad , board), ( left_pad , board )) 245 | pad_img = np.pad(img, pad4d, 'symmetric').astype(img.dtype) 246 | 247 | thisrowstart, thiscolstart =0, 0 248 | thisrowend, thiscolend = 0,0 249 | for row_idx in range(numRowblocks): 250 | thisrowlen = min(hidden_windowsize[0], row_size - row_idx * step_size[0]) 251 | row_step_len = min(step_size[0], row_size - row_idx * step_size[0]) 252 | 253 | thisrowstart = 0 if row_idx == 0 else thisrowstart + step_size[0] 254 | 255 | thisrowend = thisrowstart + thisrowlen 256 | 257 | row_shift = 0 258 | if fixed_window: 259 | if thisrowlen < hidden_windowsize[0]: 260 | row_shift = hidden_windowsize[0] - thisrowlen 261 | 262 | for col_idx in range(numColblocks): 263 | thiscollen = min(hidden_windowsize[1], col_size - col_idx * step_size[1]) 264 | col_step_len = min(step_size[1], col_size - col_idx * step_size[1]) 265 | 266 | thiscolstart = 0 if col_idx == 0 else thiscolstart + step_size[1] 267 | 268 | thiscolend = thiscolstart + thiscollen 269 | 270 | col_shift = 0 271 | if fixed_window: 272 | # we need to shift the patch to left to make it at least windowsize. 273 | if thiscollen < hidden_windowsize[1]: 274 | col_shift = hidden_windowsize[1] - thiscollen 275 | 276 | # 277 | #----board----******************----board---- 278 | # 279 | crop_r_start = thisrowstart - board - row_shift + left_pad 280 | crop_c_start = thiscolstart - board - col_shift + left_pad 281 | crop_r_end = thisrowend + board + left_pad 282 | crop_c_end = thiscolend + board + left_pad 283 | 284 | #we need to handle the tricky board condition 285 | # thispatch will be of size (:,:, windowsize+ 2*board) 286 | #thisPatch = pad_img[:,crop_r_start:crop_r_end, crop_c_start:crop_c_end].copy() 287 | crop_patch_slice = (slice(crop_r_start, crop_r_end), slice(crop_c_start, crop_c_end)) 288 | org_slice_tuple = (crop_r_start-left_pad, crop_c_start -left_pad ) 289 | 290 | thisSize = (thisrowlen + 2*board + row_shift, thiscollen + 2*board + col_shift) 291 | 292 | 293 | org_slice = (slice(thisrowstart, thisrowend), slice(thiscolstart, thiscolend)) 294 | # slice on a cooridinate of the original image 295 | extract_slice = (slice(board + row_shift, board + thisrowlen + row_shift), 296 | slice(board + col_shift, board + col_shift + thiscollen)) 297 | # extract on local coordinate of a patch 298 | 299 | identifier = identifier +1 300 | PackList.append((crop_patch_slice,org_slice ,extract_slice, thisSize,identifier, org_slice_tuple)) 301 | 302 | if thisSize in IndexDict: 303 | IndexDict[thisSize].append(identifier) 304 | else: 305 | IndexDict[thisSize] = [] 306 | IndexDict[thisSize].append(identifier) 307 | 308 | PackDict = {} 309 | for this_size in list(IndexDict.keys()): 310 | iden_list = IndexDict[this_size] 311 | this_len = len(iden_list) 312 | org_slice_list = [] 313 | extract_slice_list = [] 314 | slice_tuple_list = [] 315 | BatchData = np.zeros( (this_len, img.shape[0]) + tuple(this_size) ) 316 | for idx, iden in enumerate(iden_list): 317 | crop_patch_slice = PackList[iden][0] 318 | BatchData[idx,...] = pad_img[:,crop_patch_slice[0],crop_patch_slice[1]].copy() 319 | org_slice_list.append(PackList[iden][1]) 320 | extract_slice_list.append(PackList[iden][2]) 321 | slice_tuple_list.append(PackList[iden][-1]) 322 | 323 | PackDict[this_size]= (BatchData,org_slice_list,extract_slice_list, slice_tuple_list) 324 | 325 | return PackDict 326 | -------------------------------------------------------------------------------- /DetJoint/yolo_v2/proj_utils/model_utils.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | from torch.autograd import Variable 6 | from .torch_utils import * 7 | from .local_utils import Indexflow, split_img, imshow 8 | from collections import deque, OrderedDict 9 | import functools 10 | 11 | 12 | def match_tensor(out, refer_shape): 13 | 14 | skiprow,skipcol = refer_shape 15 | row, col = out.size()[2], out.size()[3] 16 | if skipcol >= col: 17 | pad_col = skipcol - col 18 | left_pad_col = pad_col // 2 19 | right_pad_col = pad_col - left_pad_col 20 | out = F.pad(out, (left_pad_col, right_pad_col, 0,0), mode='reflect') 21 | else: 22 | crop_col = col - skipcol 23 | left_crop_col = crop_col // 2 24 | right_col = left_crop_col + skipcol 25 | out = out[:,:,:, left_crop_col:right_col] 26 | 27 | if skiprow >= row: 28 | pad_row = skiprow - row 29 | left_pad_row = pad_row // 2 30 | right_pad_row = pad_row - left_pad_row 31 | out = F.pad(out, (0,0, left_pad_row,right_pad_row), mode='reflect') 32 | else: 33 | crop_row = row - skiprow 34 | left_crop_row = crop_row // 2 35 | 36 | right_row = left_crop_row + skiprow 37 | 38 | out = out[:,:,left_crop_row:right_row, :] 39 | return out 40 | -------------------------------------------------------------------------------- /DetJoint/yolo_v2/proj_utils/plot_utils.py: -------------------------------------------------------------------------------- 1 | try: 2 | from visdom import Visdom 3 | except: 4 | print('Better install visdom') 5 | import numpy as np 6 | import random 7 | 8 | import scipy.misc 9 | from scipy.misc import imsave 10 | 11 | from .local_utils import imshow, writeImg, normalize_img 12 | _port = 8899 13 | 14 | 15 | 16 | class plot_scalar(object): 17 | def __init__(self, name='default', env='main', rate= 1, handler=None, port = _port): 18 | 19 | self.__dict__.update(locals()) 20 | self.values = [] 21 | self.steps = [] 22 | if self.handler is None: 23 | self.handler = Visdom(port=port) 24 | self.count = 0 25 | 26 | def plot(self,values, step = None): 27 | org_type_chk = type(values) is list 28 | if not org_type_chk: 29 | values = [values] 30 | 31 | len_val = len(values) 32 | if step is None: 33 | step = list(range(self.count, self.count+len_val)) 34 | 35 | self.count += len_val 36 | self.steps.extend(step) 37 | self.values.extend(values) 38 | 39 | if self.count % self.rate == 0 or org_type_chk: 40 | self.flush() 41 | 42 | def reset(self): 43 | self.steps = [] 44 | self.values = [] 45 | 46 | def flush(self): 47 | #print('flush the plot. :)') 48 | assert type(self.values) is list, 'values have to be list' 49 | if type(self.values[0]) is not list: 50 | self.values = [self.values] 51 | 52 | n_lines = len(self.values) 53 | repeat_steps = [self.steps]*n_lines 54 | steps = np.array(repeat_steps).transpose() 55 | values = np.array(self.values).transpose() 56 | 57 | assert not np.isnan(values).any(), 'nan error in loss!!!' 58 | res = self.handler.line( 59 | X = steps, 60 | Y= values, 61 | win= self.name, 62 | update='append', 63 | opts=dict(title = self.name), 64 | env = self.env 65 | ) 66 | 67 | if res != self.name: 68 | self.handler.line( 69 | X=steps, 70 | Y=values, 71 | win=self.name, 72 | env=self.env, 73 | opts=dict(title=self.name) 74 | ) 75 | 76 | self.reset() 77 | -------------------------------------------------------------------------------- /DetJoint/yolo_v2/proj_utils/torch_utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | from copy import copy 3 | import torch 4 | import torch.nn.functional as F 5 | from torch.autograd import Variable 6 | import torch.nn as nn 7 | 8 | 9 | def to_variable(x, requires_grad=True, var=True,volatile=False): 10 | 11 | if type(x) is Variable: 12 | return x 13 | if type(x) is np.ndarray: 14 | x = torch.from_numpy(x.astype(np.float32)) 15 | if var: 16 | x = Variable(x, requires_grad=requires_grad, volatile=volatile) 17 | x.volatile = volatile 18 | 19 | return x 20 | 21 | 22 | def to_device(src, ref, var = True, volatile = False, requires_grad=True): 23 | requires_grad = requires_grad and (not volatile) 24 | src = to_variable(src, var=var, volatile=volatile,requires_grad=requires_grad) 25 | return src.cuda(ref.get_device()) if ref.is_cuda else src 26 | 27 | 28 | def set_lr(optimizer, lr): 29 | for param_group in optimizer.param_groups: 30 | param_group['lr'] = lr 31 | 32 | return optimizer 33 | 34 | 35 | def tensor_to_img(t_img, mean, std): 36 | t_img.squeeze_() 37 | for t, m, s in zip(t_img, mean, std): 38 | t.mul_(s).add_(m) 39 | 40 | np_img = t_img.numpy().transpose(1, 2, 0) 41 | np_img = (np_img * 255.0).astype(np.uint8) 42 | 43 | return np_img 44 | -------------------------------------------------------------------------------- /DetJoint/yolo_v2/test_yolo.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | import os, sys, pdb 4 | import numpy as np 5 | from scipy.misc import toimage 6 | import matplotlib.pyplot as plt 7 | import torch 8 | import deepdish as dd 9 | import json 10 | 11 | from .utils.timer import Timer 12 | from .utils import yolo as yolo_utils 13 | from .proj_utils.plot_utils import plot_scalar 14 | from .proj_utils.local_utils import writeImg, mkdirs 15 | from .proj_utils.torch_utils import tensor_to_img 16 | from .knee_utils import knee_det_cls 17 | from .knee_utils import evaluate_det_cls, save_pred_box_coors 18 | from .knee_utils import overlay_bbox_iou 19 | 20 | 21 | def test_eng(dataloader, model_root, save_root, net, args, cfg): 22 | net.eval() 23 | weightspath = os.path.join(model_root, args.model_name) 24 | weights_dict = torch.load(weightspath, map_location=lambda storage, loc: storage) 25 | print("===="*20) 26 | print('Model name is {}'.format(weightspath)) 27 | net.load_state_dict(weights_dict) 28 | print("Num images: {}".format(len(dataloader))) 29 | _t = {'im_detect': Timer(), 'misc': Timer()} 30 | 31 | evalate_matrix = np.zeros((5, 5), dtype=np.int) 32 | true_box_num, total_box_num = 0, 0 33 | all_overlap = [] 34 | 35 | # import pdb; pdb.set_trace() 36 | for ind, data in enumerate(dataloader): 37 | cur_img, cur_boxes, cur_classes, cur_name = data 38 | # detection 39 | _t['im_detect'].tic() 40 | result_dict = knee_det_cls(net, cur_img, cfg=cfg) 41 | bbox_pred, iou_pred, prob_pred = result_dict['bbox'], result_dict['iou'], result_dict['prob'] 42 | detect_time = _t['im_detect'].toc() 43 | # postprocessing 44 | _t['misc'].tic() 45 | bboxes, scores, cls_inds = yolo_utils.postprocess_bbox(bbox_pred, iou_pred, prob_pred, cur_img.shape[2:], cfg, thresh=0.12) 46 | utils_time = _t['misc'].toc() 47 | 48 | gt_boxes = cur_boxes.squeeze().numpy() 49 | gt_classes = cur_classes.squeeze().numpy() 50 | _, true_num, total_num, overlap_list, _ = evaluate_det_cls(gt_boxes, gt_classes, bboxes, cls_inds, 51 | cfg.num_classes, cfg.JIthresh) 52 | # save_pred_box_coors(save_root, gt_boxes, gt_classes, bboxes, cur_name[0]) 53 | if total_num != 2: # Check wrong detection file 54 | print("Name: {}, num: {}".format(cur_name[0], total_num)) 55 | true_box_num += true_num 56 | total_box_num += total_num 57 | all_overlap.extend(overlap_list) 58 | 59 | # if cur_name[0] in selection: 60 | # box_dict = {} 61 | # box_dict["det"] = bboxes.tolist() 62 | # box_dict["gt"] = gt_boxes.tolist() 63 | # select_dict[cur_name[0]] = box_dict 64 | 65 | 66 | if (ind+1) % 100 == 0: 67 | print('{}/{} detection time {:.4f}, post_processing time {:.4f}'.format( 68 | ind+1, len(dataloader), detect_time, utils_time)) 69 | # Overlay gt box and prediction box and overlap value 70 | img_np = tensor_to_img(cur_img, cfg.rgb_mean, cfg.rgb_var) 71 | overlaid_img = overlay_bbox_iou(img_np, bboxes, gt_boxes) 72 | writeImg(overlaid_img, os.path.join(save_root, cur_name[0]+'.png')) 73 | 74 | # # save to json 75 | # with open(os.path.join(save_root, 'selection.json'), 'w') as outfile: 76 | # json.dump(select_dict, outfile) 77 | 78 | print("---Detection accuracy---") 79 | print("Number of object: {}, Mean IoU is: {}".format(len(all_overlap), np.mean(all_overlap))) 80 | print("True predicted knee: {}\n All predicted knee: {}\n Total accuracy: {:.4f}\n".format( 81 | true_box_num, total_box_num, true_box_num*1.0/total_box_num)) 82 | -------------------------------------------------------------------------------- /DetJoint/yolo_v2/train_yolo.py: -------------------------------------------------------------------------------- 1 | # -*- coding: utf-8 -*- 2 | 3 | import os, sys, pdb 4 | import numpy as np 5 | import torch 6 | import datetime 7 | import cv2 8 | 9 | from .proj_utils.plot_utils import plot_scalar 10 | from .proj_utils.torch_utils import to_device, set_lr 11 | from .proj_utils.local_utils import mkdirs 12 | 13 | 14 | def train_eng(train_dataloader, val_dataloader, model_root, net, args): 15 | net.train() 16 | lr = args.lr 17 | optimizer = torch.optim.SGD(net.parameters(), lr=lr, momentum=args.momentum, weight_decay=args.weight_decay) 18 | 19 | loss_train_plot = plot_scalar(name = "loss_train", env= args.model_name, rate = args.display_freq) 20 | loss_bbox_plot = plot_scalar(name = "loss_bbox", env= args.model_name, rate = args.display_freq) 21 | loss_iou_plot = plot_scalar(name = "loss_iou", env= args.model_name, rate = args.display_freq) 22 | loss_cls_plot = plot_scalar(name = "loss_cls", env= args.model_name, rate = args.display_freq) 23 | 24 | print("Start training...") 25 | best_loss = 1.0e6 # given a high loss value 26 | for cur_epoch in range(1, args.maxepoch+1): 27 | train_loss, bbox_loss, iou_loss, cls_loss = 0., 0., 0., 0. 28 | mini_batch_num = 0 29 | for cur_batch, data in enumerate(train_dataloader): 30 | cur_imgs, cur_boxes, cur_classes, cur_names = data 31 | # forward 32 | im_data = to_device(cur_imgs, net.device_id) 33 | _ = net(im_data, cur_boxes, cur_classes) 34 | 35 | # backward 36 | loss = net.loss 37 | bbox_loss_val = net.bbox_loss.data.cpu().numpy().mean() 38 | iou_loss_val = net.iou_loss.data.cpu().numpy().mean() 39 | cls_loss_val = net.cls_loss.data.cpu().numpy().mean() 40 | train_loss_val = loss.data.cpu().numpy().mean() 41 | 42 | bbox_loss += bbox_loss_val 43 | iou_loss += iou_loss_val 44 | cls_loss += cls_loss_val 45 | train_loss += train_loss_val 46 | loss_train_plot.plot(train_loss_val) 47 | 48 | optimizer.zero_grad() 49 | loss.backward() 50 | optimizer.step() 51 | 52 | mini_batch_num += 1 53 | # Print error information 54 | if (cur_batch + 1) % args.display_freq == 0: 55 | num_samples = args.display_freq * args.batch_size 56 | ttl_loss = train_loss/num_samples 57 | bb_loss, iou_loss = bbox_loss/num_samples, iou_loss/num_samples 58 | cls_loss = cls_loss/num_samples 59 | print_str = "Train:Epoch:{:>3}/{:>3}, {:>3}/{:>3} total loss:{:.6f}, bbox_loss:{:.6f}, iou_loss:{:.6f}, cls_loss:{:.2f}" 60 | print(print_str.format(cur_epoch, args.maxepoch, mini_batch_num, len(train_dataloader), ttl_loss, bb_loss, iou_loss, cls_loss)) 61 | train_loss, bbox_loss, iou_loss, cls_loss = 0, 0., 0., 0. 62 | 63 | if cur_epoch % args.save_freq == 0: 64 | # Validate current model's performance 65 | best_loss = validate(val_dataloader, net, cur_epoch, model_root, best_loss, args) 66 | net.train() 67 | 68 | # Adjust learing rate 69 | if cur_epoch in args.lr_decay_epochs: 70 | lr *= args.lr_decay 71 | optimizer = set_lr(optimizer, lr) 72 | print("Current lr is {}".format(lr)) 73 | 74 | 75 | def validate(val_dataloader, net, cur_epoch, model_root, best_loss, args): 76 | net.eval() 77 | bbox_loss_val, iou_loss_val, cls_loss_val, total_loss_val = 0., 0., 0., 0. 78 | 79 | for ind, data in enumerate(val_dataloader): 80 | cur_imgs, cur_boxes, cur_classes, cur_names = data 81 | # forward 82 | im_data = to_device(cur_imgs, net.device_id) 83 | _ = net(im_data, cur_boxes, cur_classes) 84 | bbox_loss_val += net.bbox_loss.data.cpu().numpy().mean() 85 | iou_loss_val += net.iou_loss.data.cpu().numpy().mean() 86 | cls_loss_val += net.cls_loss.data.cpu().numpy().mean() 87 | total_loss_val += net.loss.data.cpu().numpy().mean() 88 | 89 | total_loss_val /= (ind + 1) 90 | 91 | # Saving new model 92 | if total_loss_val < best_loss: 93 | best_loss = total_loss_val 94 | 95 | bbox_loss_val /= (ind + 1) 96 | iou_loss_val /= (ind + 1) 97 | cls_loss_val /= (ind + 1) 98 | 99 | print("Validation--Epoch {:>2}, total loss: {:.5f}, bbox_loss: {:.5f}, iou_loss: {:.5f}, cls_loss:{:.2f}".format( 100 | cur_epoch, total_loss_val, bbox_loss_val, iou_loss_val, cls_loss_val)) 101 | 102 | weights_name = "det-epoch-" + str(cur_epoch).zfill(3) + "-" + "{:.6f}".format(best_loss) + '.pth' 103 | model_path = os.path.join(model_root, weights_name) 104 | torch.save(net.state_dict(), model_path) 105 | print('save weights at {}'.format(model_path)) 106 | 107 | return best_loss 108 | -------------------------------------------------------------------------------- /DetJoint/yolo_v2/utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PingjunChen/GradingKneeOA/dc19c2462111c3c5b27066e738fc11ddf57a6407/DetJoint/yolo_v2/utils/__init__.py -------------------------------------------------------------------------------- /DetJoint/yolo_v2/utils/bbox.pyx: -------------------------------------------------------------------------------- 1 | # -------------------------------------------------------- 2 | # Fast R-CNN 3 | # Copyright (c) 2015 Microsoft 4 | # Licensed under The MIT License [see LICENSE for details] 5 | # Written by Sergey Karayev 6 | # -------------------------------------------------------- 7 | 8 | cimport cython 9 | import numpy as np 10 | cimport numpy as np 11 | 12 | DTYPE = np.float 13 | ctypedef np.float_t DTYPE_t 14 | 15 | cdef extern from "math.h": 16 | double abs(double m) 17 | double log(double x) 18 | 19 | 20 | def bbox_overlaps(np.ndarray[DTYPE_t, ndim=2] boxes, 21 | np.ndarray[DTYPE_t, ndim=2] query_boxes): 22 | return bbox_overlaps_c(boxes, query_boxes) 23 | 24 | cdef np.ndarray[DTYPE_t, ndim=2] bbox_overlaps_c( 25 | np.ndarray[DTYPE_t, ndim=2] boxes, 26 | np.ndarray[DTYPE_t, ndim=2] query_boxes): 27 | """ 28 | Parameters 29 | ---------- 30 | boxes: (N, 4) ndarray of float 31 | query_boxes: (K, 4) ndarray of float 32 | Returns 33 | ------- 34 | overlaps: (N, K) ndarray of overlap between boxes and query_boxes 35 | """ 36 | cdef unsigned int N = boxes.shape[0] 37 | cdef unsigned int K = query_boxes.shape[0] 38 | cdef np.ndarray[DTYPE_t, ndim=2] overlaps = np.zeros((N, K), dtype=DTYPE) 39 | cdef DTYPE_t iw, ih, box_area 40 | cdef DTYPE_t ua 41 | cdef unsigned int k, n 42 | for k in range(K): 43 | box_area = ( 44 | (query_boxes[k, 2] - query_boxes[k, 0] + 1) * 45 | (query_boxes[k, 3] - query_boxes[k, 1] + 1) 46 | ) 47 | for n in range(N): 48 | iw = ( 49 | min(boxes[n, 2], query_boxes[k, 2]) - 50 | max(boxes[n, 0], query_boxes[k, 0]) + 1 51 | ) 52 | if iw > 0: 53 | ih = ( 54 | min(boxes[n, 3], query_boxes[k, 3]) - 55 | max(boxes[n, 1], query_boxes[k, 1]) + 1 56 | ) 57 | if ih > 0: 58 | ua = float( 59 | (boxes[n, 2] - boxes[n, 0] + 1) * 60 | (boxes[n, 3] - boxes[n, 1] + 1) + 61 | box_area - iw * ih 62 | ) 63 | overlaps[n, k] = iw * ih / ua 64 | return overlaps 65 | 66 | 67 | def bbox_intersections( 68 | np.ndarray[DTYPE_t, ndim=2] boxes, 69 | np.ndarray[DTYPE_t, ndim=2] query_boxes): 70 | return bbox_intersections_c(boxes, query_boxes) 71 | 72 | 73 | cdef np.ndarray[DTYPE_t, ndim=2] bbox_intersections_c( 74 | np.ndarray[DTYPE_t, ndim=2] boxes, 75 | np.ndarray[DTYPE_t, ndim=2] query_boxes): 76 | """ 77 | For each query box compute the intersection ratio covered by boxes 78 | ---------- 79 | Parameters 80 | ---------- 81 | boxes: (N, 4) ndarray of float 82 | query_boxes: (K, 4) ndarray of float 83 | Returns 84 | ------- 85 | overlaps: (N, K) ndarray of intersec between boxes and query_boxes 86 | """ 87 | cdef unsigned int N = boxes.shape[0] 88 | cdef unsigned int K = query_boxes.shape[0] 89 | cdef np.ndarray[DTYPE_t, ndim=2] intersec = np.zeros((N, K), dtype=DTYPE) 90 | cdef DTYPE_t iw, ih, box_area 91 | cdef DTYPE_t ua 92 | cdef unsigned int k, n 93 | for k in range(K): 94 | box_area = ( 95 | (query_boxes[k, 2] - query_boxes[k, 0] + 1) * 96 | (query_boxes[k, 3] - query_boxes[k, 1] + 1) 97 | ) 98 | for n in range(N): 99 | iw = ( 100 | min(boxes[n, 2], query_boxes[k, 2]) - 101 | max(boxes[n, 0], query_boxes[k, 0]) + 1 102 | ) 103 | if iw > 0: 104 | ih = ( 105 | min(boxes[n, 3], query_boxes[k, 3]) - 106 | max(boxes[n, 1], query_boxes[k, 1]) + 1 107 | ) 108 | if ih > 0: 109 | intersec[n, k] = iw * ih / box_area 110 | return intersec 111 | 112 | 113 | def bbox_ious( 114 | np.ndarray[DTYPE_t, ndim=2] boxes, 115 | np.ndarray[DTYPE_t, ndim=2] query_boxes): 116 | return bbox_ious_c(boxes, query_boxes) 117 | 118 | 119 | cdef np.ndarray[DTYPE_t, ndim=2] bbox_ious_c( 120 | np.ndarray[DTYPE_t, ndim=2] boxes, 121 | np.ndarray[DTYPE_t, ndim=2] query_boxes): 122 | """ 123 | For each query box compute the IOU covered by boxes 124 | ---------- 125 | Parameters 126 | ---------- 127 | boxes: (N, 4) ndarray of float 128 | query_boxes: (K, 4) ndarray of float 129 | Returns 130 | ------- 131 | overlaps: (N, K) ndarray of intersec between boxes and query_boxes 132 | """ 133 | cdef unsigned int N = boxes.shape[0] 134 | cdef unsigned int K = query_boxes.shape[0] 135 | cdef np.ndarray[DTYPE_t, ndim=2] intersec = np.zeros((N, K), dtype=DTYPE) 136 | cdef DTYPE_t iw, ih, qbox_area, box_area, inter_area 137 | cdef unsigned int k, n 138 | for k in range(K): 139 | qbox_area = ( 140 | (query_boxes[k, 2] - query_boxes[k, 0] + 1) * 141 | (query_boxes[k, 3] - query_boxes[k, 1] + 1) 142 | ) 143 | for n in range(N): 144 | iw = ( 145 | min(boxes[n, 2], query_boxes[k, 2]) - 146 | max(boxes[n, 0], query_boxes[k, 0]) + 1 147 | ) 148 | if iw > 0: 149 | ih = ( 150 | min(boxes[n, 3], query_boxes[k, 3]) - 151 | max(boxes[n, 1], query_boxes[k, 1]) + 1 152 | ) 153 | if ih > 0: 154 | box_area = ( 155 | (boxes[n, 2] - boxes[n, 0] + 1) * 156 | (boxes[n, 3] - boxes[n, 1] + 1) 157 | ) 158 | inter_area = iw * ih 159 | intersec[n, k] = inter_area / (qbox_area + box_area - inter_area) 160 | return intersec 161 | 162 | 163 | def anchor_intersections( 164 | np.ndarray[DTYPE_t, ndim=2] anchors, 165 | np.ndarray[DTYPE_t, ndim=2] query_boxes): 166 | return anchor_intersections_c(anchors, query_boxes) 167 | 168 | 169 | cdef np.ndarray[DTYPE_t, ndim=2] anchor_intersections_c( 170 | np.ndarray[DTYPE_t, ndim=2] anchors, 171 | np.ndarray[DTYPE_t, ndim=2] query_boxes): 172 | """ 173 | For each query box compute the intersection ratio covered by anchors 174 | ---------- 175 | Parameters 176 | ---------- 177 | boxes: (N, 2) ndarray of float 178 | query_boxes: (K, 4) ndarray of float 179 | Returns 180 | ------- 181 | overlaps: (N, K) ndarray of intersec between boxes and query_boxes 182 | """ 183 | cdef unsigned int N = anchors.shape[0] 184 | cdef unsigned int K = query_boxes.shape[0] 185 | cdef np.ndarray[DTYPE_t, ndim=2] intersec = np.zeros((N, K), dtype=DTYPE) 186 | cdef DTYPE_t iw, ih, anchor_area, inter_area 187 | cdef DTYPE_t boxw, boxh 188 | cdef unsigned int k, n 189 | for n in range(N): 190 | anchor_area = anchors[n, 0] * anchors[n, 1] 191 | for k in range(K): 192 | boxw = (query_boxes[k, 2] - query_boxes[k, 0] + 1) 193 | boxh = (query_boxes[k, 3] - query_boxes[k, 1] + 1) 194 | iw = min(anchors[n, 0], boxw) 195 | ih = min(anchors[n, 1], boxh) 196 | inter_area = iw * ih 197 | intersec[n, k] = inter_area / (anchor_area + boxw * boxh - inter_area) 198 | 199 | return intersec 200 | 201 | 202 | def bbox_intersections_self( 203 | np.ndarray[DTYPE_t, ndim=2] boxes): 204 | return bbox_intersections_self_c(boxes) 205 | 206 | 207 | cdef np.ndarray[DTYPE_t, ndim=2] bbox_intersections_self_c( 208 | np.ndarray[DTYPE_t, ndim=2] boxes): 209 | """ 210 | For each query box compute the intersection ratio covered by boxes 211 | ---------- 212 | Parameters 213 | ---------- 214 | boxes: (N, 4) ndarray of float 215 | Returns 216 | ------- 217 | overlaps: (N, N) ndarray of intersec between boxes and query_boxes 218 | """ 219 | cdef unsigned int N = boxes.shape[0] 220 | cdef np.ndarray[DTYPE_t, ndim=2] intersec = np.zeros((N, N), dtype=DTYPE) 221 | cdef DTYPE_t iw, ih, box_area 222 | cdef unsigned int k, n 223 | 224 | for k in range(N): 225 | box_area = ( 226 | (boxes[k, 2] - boxes[k, 0] + 1) * 227 | (boxes[k, 3] - boxes[k, 1] + 1) 228 | ) 229 | for n in range(k+1, N): 230 | iw = ( 231 | min(boxes[n, 2], boxes[k, 2]) - 232 | max(boxes[n, 0], boxes[k, 0]) + 1 233 | ) 234 | if iw > 0: 235 | ih = ( 236 | min(boxes[n, 3], boxes[k, 3]) - 237 | max(boxes[n, 1], boxes[k, 1]) + 1 238 | ) 239 | if ih > 0: 240 | intersec[k, n] = iw * ih / box_area 241 | return intersec 242 | 243 | 244 | def bbox_similarities( 245 | np.ndarray[DTYPE_t, ndim=2] boxes, 246 | np.ndarray[DTYPE_t, ndim=2] query_boxes): 247 | return bbox_similarities_c(boxes, query_boxes) 248 | 249 | cdef np.ndarray[DTYPE_t, ndim=2] bbox_similarities_c( 250 | np.ndarray[DTYPE_t, ndim=2] boxes, 251 | np.ndarray[DTYPE_t, ndim=2] query_boxes): 252 | """ 253 | For each query box compute the intersection ratio covered by boxes 254 | ---------- 255 | Parameters 256 | ---------- 257 | boxes: (N, 4) ndarray of float 258 | query_boxes: (K, 4) ndarray of float (dets) 259 | Returns 260 | ------- 261 | overlaps: (N, K) ndarray of similarity scores between boxes and query_boxes 262 | """ 263 | cdef unsigned int N = boxes.shape[0] 264 | cdef unsigned int K = query_boxes.shape[0] 265 | cdef np.ndarray[DTYPE_t, ndim=2] sims = np.zeros((N, K), dtype=DTYPE) 266 | cdef DTYPE_t cx1, cy1, w1, h1 267 | cdef DTYPE_t cx2, cy2, w2, h2 268 | 269 | cdef DTYPE_t loc_dist, shape_dist 270 | 271 | cdef unsigned int k, n 272 | for n in range(N): 273 | cx1 = (boxes[n, 0] + boxes[n, 2]) * 0.5 274 | cy1 = (boxes[n, 1] + boxes[n, 3]) * 0.5 275 | w1 = boxes[n, 2] - boxes[n, 0] + 1 276 | h1 = boxes[n, 3] - boxes[n, 1] + 1 277 | 278 | for k in range(K): 279 | cx2 = (query_boxes[k, 0] + query_boxes[k, 2]) * 0.5 280 | cy2 = (query_boxes[k, 1] + query_boxes[k, 3]) * 0.5 281 | w2 = query_boxes[k, 2] - query_boxes[k, 0] + 1 282 | h2 = query_boxes[k, 3] - query_boxes[k, 1] + 1 283 | 284 | loc_dist = abs(cx1 - cx2) / (w1 + w2) + abs(cy1 - cy2) / (h1 + h2) 285 | shape_dist = abs(w2 * h2 / (w1 * h1) - 1.0) 286 | 287 | sims[n, k] = -log(loc_dist + 0.001) - shape_dist * shape_dist + 1 288 | 289 | return sims -------------------------------------------------------------------------------- /DetJoint/yolo_v2/utils/build.py: -------------------------------------------------------------------------------- 1 | # -------------------------------------------------------- 2 | # Fast R-CNN 3 | # Copyright (c) 2015 Microsoft 4 | # Licensed under The MIT License [see LICENSE for details] 5 | # Written by Ross Girshick 6 | # -------------------------------------------------------- 7 | 8 | import os 9 | from os.path import join as pjoin 10 | import numpy as np 11 | from distutils.core import setup 12 | from distutils.extension import Extension 13 | from Cython.Distutils import build_ext 14 | 15 | 16 | def find_in_path(name, path): 17 | "Find a file in a search path" 18 | for dir in path.split(os.pathsep): 19 | binpath = pjoin(dir, name) 20 | if os.path.exists(binpath): 21 | return os.path.abspath(binpath) 22 | return None 23 | 24 | 25 | def locate_cuda(): 26 | """Locate the CUDA environment on the system 27 | 28 | Returns a dict with keys 'home', 'nvcc', 'include', and 'lib64' 29 | and values giving the absolute path to each directory. 30 | 31 | Starts by looking for the CUDAHOME env variable. If not found, everything 32 | is based on finding 'nvcc' in the PATH. 33 | """ 34 | 35 | # first check if the CUDAHOME env variable is in use 36 | if 'CUDAHOME' in os.environ: 37 | home = os.environ['CUDAHOME'] 38 | nvcc = pjoin(home, 'bin', 'nvcc') 39 | else: 40 | # otherwise, search the PATH for NVCC 41 | default_path = pjoin(os.sep, 'usr', 'local', 'cuda', 'bin') 42 | nvcc = find_in_path('nvcc', os.environ['PATH'] + os.pathsep + default_path) 43 | if nvcc is None: 44 | raise EnvironmentError('The nvcc binary could not be ' 45 | 'located in your $PATH. Either add it to your path, or set $CUDAHOME') 46 | home = os.path.dirname(os.path.dirname(nvcc)) 47 | 48 | cudaconfig = {'home': home, 'nvcc': nvcc, 49 | 'include': pjoin(home, 'include'), 50 | 'lib64': pjoin(home, 'lib64')} 51 | for k, v in cudaconfig.items(): 52 | if not os.path.exists(v): 53 | raise EnvironmentError('The CUDA %s path could not be located in %s' % (k, v)) 54 | 55 | return cudaconfig 56 | 57 | 58 | CUDA = locate_cuda() 59 | 60 | # Obtain the numpy include directory. This logic works across numpy versions. 61 | try: 62 | numpy_include = np.get_include() 63 | except AttributeError: 64 | numpy_include = np.get_numpy_include() 65 | 66 | 67 | def customize_compiler_for_nvcc(self): 68 | """inject deep into distutils to customize how the dispatch 69 | to gcc/nvcc works. 70 | 71 | If you subclass UnixCCompiler, it's not trivial to get your subclass 72 | injected in, and still have the right customizations (i.e. 73 | distutils.sysconfig.customize_compiler) run on it. So instead of going 74 | the OO route, I have this. Note, it's kindof like a wierd functional 75 | subclassing going on.""" 76 | 77 | # tell the compiler it can processes .cu 78 | self.src_extensions.append('.cu') 79 | 80 | # save references to the default compiler_so and _comple methods 81 | default_compiler_so = self.compiler_so 82 | super = self._compile 83 | 84 | # now redefine the _compile method. This gets executed for each 85 | # object but distutils doesn't have the ability to change compilers 86 | # based on source extension: we add it. 87 | def _compile(obj, src, ext, cc_args, extra_postargs, pp_opts): 88 | print(extra_postargs) 89 | if os.path.splitext(src)[1] == '.cu': 90 | # use the cuda for .cu files 91 | self.set_executable('compiler_so', CUDA['nvcc']) 92 | # use only a subset of the extra_postargs, which are 1-1 translated 93 | # from the extra_compile_args in the Extension class 94 | postargs = extra_postargs['nvcc'] 95 | else: 96 | postargs = extra_postargs['gcc'] 97 | 98 | super(obj, src, ext, cc_args, postargs, pp_opts) 99 | # reset the default compiler_so, which we might have changed for cuda 100 | self.compiler_so = default_compiler_so 101 | 102 | # inject our redefined _compile method into the class 103 | self._compile = _compile 104 | 105 | 106 | # run the customize_compiler 107 | class custom_build_ext(build_ext): 108 | def build_extensions(self): 109 | customize_compiler_for_nvcc(self.compiler) 110 | build_ext.build_extensions(self) 111 | 112 | 113 | ext_modules = [ 114 | Extension( 115 | "cython_bbox", 116 | ["bbox.pyx"], 117 | extra_compile_args={'gcc': ["-Wno-cpp", "-Wno-unused-function"]}, 118 | include_dirs=[numpy_include] 119 | ), 120 | Extension( 121 | "cython_yolo", 122 | ["yolo.pyx"], 123 | extra_compile_args={'gcc': ["-Wno-cpp", "-Wno-unused-function"]}, 124 | include_dirs=[numpy_include] 125 | ), 126 | Extension( 127 | "nms.cpu_nms", 128 | ["nms/cpu_nms.pyx"], 129 | extra_compile_args={'gcc': ["-Wno-cpp", "-Wno-unused-function"]}, 130 | include_dirs=[numpy_include] 131 | ), 132 | Extension('nms.gpu_nms', 133 | ['nms/nms_kernel.cu', 'nms/gpu_nms.pyx'], 134 | library_dirs=[CUDA['lib64']], 135 | libraries=['cudart'], 136 | language='c++', 137 | runtime_library_dirs=[CUDA['lib64']], 138 | # this syntax is specific to this build system 139 | # we're only going to use certain compiler args with nvcc and not with gcc 140 | # the implementation of this trick is in customize_compiler() below 141 | extra_compile_args={'gcc': ["-Wno-unused-function"], 142 | 'nvcc': ['-arch=sm_35', 143 | '--ptxas-options=-v', 144 | '-c', 145 | '--compiler-options', 146 | "'-fPIC'"]}, 147 | include_dirs=[numpy_include, CUDA['include']] 148 | ), 149 | # Extension( 150 | # 'pycocotools._mask', 151 | # sources=['pycocotools/maskApi.c', 'pycocotools/_mask.pyx'], 152 | # include_dirs=[numpy_include, 'pycocotools'], 153 | # extra_compile_args={ 154 | # 'gcc': ['-Wno-cpp', '-Wno-unused-function', '-std=c99']}, 155 | # ), 156 | ] 157 | 158 | setup( 159 | name='mot_utils', 160 | ext_modules=ext_modules, 161 | # inject our custom trigger 162 | cmdclass={'build_ext': custom_build_ext}, 163 | ) 164 | -------------------------------------------------------------------------------- /DetJoint/yolo_v2/utils/im_transform.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import cv2 3 | 4 | 5 | def imcv2_recolor(im, a=.1): 6 | # t = [np.random.uniform()] 7 | # t += [np.random.uniform()] 8 | # t += [np.random.uniform()] 9 | # t = np.array(t) * 2. - 1. 10 | t = np.random.uniform(-1, 1, 3) 11 | 12 | # random amplify each channel 13 | im = im.astype(np.float) 14 | im *= (1 + t * a) 15 | mx = 255. * (1 + a) 16 | up = np.random.uniform(-1, 1) 17 | im = np.power(im / mx, 1. + up * .5) 18 | # return np.array(im * 255., np.uint8) 19 | return im 20 | 21 | 22 | def imcv2_affine_trans(im): 23 | # Scale and translate 24 | h, w, c = im.shape 25 | scale = np.random.uniform() / 10. + 1. 26 | max_offx = (scale - 1.) * w 27 | max_offy = (scale - 1.) * h 28 | offx = int(np.random.uniform() * max_offx) 29 | offy = int(np.random.uniform() * max_offy) 30 | 31 | im = cv2.resize(im, (0, 0), fx=scale, fy=scale) 32 | im = im[offy: (offy + h), offx: (offx + w)] 33 | flip = np.random.uniform() > 0.5 34 | if flip: 35 | im = cv2.flip(im, 1) 36 | 37 | return im, [scale, [offx, offy], flip] 38 | -------------------------------------------------------------------------------- /DetJoint/yolo_v2/utils/network.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from torch.autograd import Variable 4 | import numpy as np 5 | 6 | 7 | class Conv2d(nn.Module): 8 | def __init__(self, in_channels, out_channels, kernel_size, stride=1, relu=True, same_padding=False): 9 | super(Conv2d, self).__init__() 10 | padding = int((kernel_size - 1) / 2) if same_padding else 0 11 | self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding=padding) 12 | self.relu = nn.LeakyReLU(0.1, inplace=True) if relu else None 13 | 14 | def forward(self, x): 15 | x = self.conv(x) 16 | if self.relu is not None: 17 | x = self.relu(x) 18 | return x 19 | 20 | 21 | class Conv2d_BatchNorm(nn.Module): 22 | def __init__(self, in_channels, out_channels, kernel_size, stride=1, relu=True, same_padding=False): 23 | super(Conv2d_BatchNorm, self).__init__() 24 | padding = int((kernel_size - 1) / 2) if same_padding else 0 25 | 26 | self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding=padding, bias=False) 27 | self.bn = nn.BatchNorm2d(out_channels, momentum=0.01) 28 | self.relu = nn.LeakyReLU(0.1, inplace=True) if relu else None 29 | 30 | def forward(self, x): 31 | x = self.conv(x) 32 | x = self.bn(x) 33 | if self.relu is not None: 34 | x = self.relu(x) 35 | return x 36 | 37 | 38 | class FC(nn.Module): 39 | def __init__(self, in_features, out_features, relu=True): 40 | super(FC, self).__init__() 41 | self.fc = nn.Linear(in_features, out_features) 42 | self.relu = nn.ReLU(inplace=True) if relu else None 43 | 44 | def forward(self, x): 45 | x = self.fc(x) 46 | if self.relu is not None: 47 | x = self.relu(x) 48 | return x 49 | 50 | 51 | def save_net(fname, net): 52 | import h5py 53 | h5f = h5py.File(fname, mode='w') 54 | for k, v in list(net.state_dict().items()): 55 | h5f.create_dataset(k, data=v.cpu().numpy()) 56 | 57 | 58 | def load_net(fname, net): 59 | import h5py 60 | h5f = h5py.File(fname, mode='r') 61 | for k, v in list(net.state_dict().items()): 62 | param = torch.from_numpy(np.asarray(h5f[k])) 63 | v.copy_(param) 64 | 65 | 66 | def load_pretrained_npy(faster_rcnn_model, fname): 67 | params = np.load(fname).item() 68 | # vgg16 69 | vgg16_dict = faster_rcnn_model.rpn.features.state_dict() 70 | for name, val in list(vgg16_dict.items()): 71 | # # print name 72 | # # print val.size() 73 | # # print param.size() 74 | if name.find('bn.') >= 0: 75 | continue 76 | i, j = int(name[4]), int(name[6]) + 1 77 | ptype = 'weights' if name[-1] == 't' else 'biases' 78 | key = 'conv{}_{}'.format(i, j) 79 | param = torch.from_numpy(params[key][ptype]) 80 | 81 | if ptype == 'weights': 82 | param = param.permute(3, 2, 0, 1) 83 | 84 | val.copy_(param) 85 | 86 | # fc6 fc7 87 | frcnn_dict = faster_rcnn_model.state_dict() 88 | pairs = {'fc6.fc': 'fc6', 'fc7.fc': 'fc7'} 89 | for k, v in list(pairs.items()): 90 | key = '{}.weight'.format(k) 91 | param = torch.from_numpy(params[v]['weights']).permute(1, 0) 92 | frcnn_dict[key].copy_(param) 93 | 94 | key = '{}.bias'.format(k) 95 | param = torch.from_numpy(params[v]['biases']) 96 | frcnn_dict[key].copy_(param) 97 | 98 | 99 | def np_to_variable(x, is_cuda=True, dtype=torch.FloatTensor, volatile=False): 100 | v = Variable(torch.from_numpy(x).type(dtype), volatile=volatile) 101 | if is_cuda: 102 | v = v.cuda() 103 | return v 104 | 105 | 106 | def variable_to_np_tf(x): 107 | return x.data.cpu().numpy().transpose([0, 2, 3, 1]) 108 | 109 | 110 | def set_trainable(model, requires_grad): 111 | for param in model.parameters(): 112 | param.requires_grad = requires_grad 113 | 114 | 115 | def weights_normal_init(model, dev=0.01): 116 | if isinstance(model, list): 117 | for m in model: 118 | weights_normal_init(m, dev) 119 | else: 120 | for m in model.modules(): 121 | if isinstance(m, nn.Conv2d): 122 | m.weight.data.normal_(0.0, dev) 123 | elif isinstance(m, nn.Linear): 124 | m.weight.data.normal_(0.0, dev) 125 | 126 | 127 | def clip_gradient(model, clip_norm): 128 | """Computes a gradient clipping coefficient based on gradient norm.""" 129 | totalnorm = 0 130 | for p in model.parameters(): 131 | if p.requires_grad: 132 | modulenorm = p.grad.data.norm() 133 | totalnorm += modulenorm ** 2 134 | totalnorm = np.sqrt(totalnorm) 135 | 136 | norm = clip_norm / max(totalnorm, clip_norm) 137 | for p in model.parameters(): 138 | if p.requires_grad: 139 | p.grad.mul_(norm) 140 | -------------------------------------------------------------------------------- /DetJoint/yolo_v2/utils/nms/.gitignore: -------------------------------------------------------------------------------- 1 | *.c 2 | *.cpp 3 | *.so 4 | -------------------------------------------------------------------------------- /DetJoint/yolo_v2/utils/nms/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PingjunChen/GradingKneeOA/dc19c2462111c3c5b27066e738fc11ddf57a6407/DetJoint/yolo_v2/utils/nms/__init__.py -------------------------------------------------------------------------------- /DetJoint/yolo_v2/utils/nms/cpu_nms.pyx: -------------------------------------------------------------------------------- 1 | # -------------------------------------------------------- 2 | # Fast R-CNN 3 | # Copyright (c) 2015 Microsoft 4 | # Licensed under The MIT License [see LICENSE for details] 5 | # Written by Ross Girshick 6 | # -------------------------------------------------------- 7 | 8 | import numpy as np 9 | cimport numpy as np 10 | 11 | cdef inline np.float32_t max(np.float32_t a, np.float32_t b): 12 | return a if a >= b else b 13 | 14 | cdef inline np.float32_t min(np.float32_t a, np.float32_t b): 15 | return a if a <= b else b 16 | 17 | def cpu_nms(np.ndarray[np.float32_t, ndim=2] dets, np.float thresh): 18 | cdef np.ndarray[np.float32_t, ndim=1] x1 = dets[:, 0] 19 | cdef np.ndarray[np.float32_t, ndim=1] y1 = dets[:, 1] 20 | cdef np.ndarray[np.float32_t, ndim=1] x2 = dets[:, 2] 21 | cdef np.ndarray[np.float32_t, ndim=1] y2 = dets[:, 3] 22 | cdef np.ndarray[np.float32_t, ndim=1] scores = dets[:, 4] 23 | 24 | cdef np.ndarray[np.float32_t, ndim=1] areas = (x2 - x1 + 1) * (y2 - y1 + 1) 25 | cdef np.ndarray[np.int_t, ndim=1] order = scores.argsort()[::-1] 26 | 27 | cdef int ndets = dets.shape[0] 28 | cdef np.ndarray[np.int_t, ndim=1] suppressed = \ 29 | np.zeros((ndets), dtype=np.int) 30 | 31 | # nominal indices 32 | cdef int _i, _j 33 | # sorted indices 34 | cdef int i, j 35 | # temp variables for box i's (the box currently under consideration) 36 | cdef np.float32_t ix1, iy1, ix2, iy2, iarea 37 | # variables for computing overlap with box j (lower scoring box) 38 | cdef np.float32_t xx1, yy1, xx2, yy2 39 | cdef np.float32_t w, h 40 | cdef np.float32_t inter, ovr 41 | 42 | keep = [] 43 | for _i in range(ndets): 44 | i = order[_i] 45 | if suppressed[i] == 1: 46 | continue 47 | keep.append(i) 48 | ix1 = x1[i] 49 | iy1 = y1[i] 50 | ix2 = x2[i] 51 | iy2 = y2[i] 52 | iarea = areas[i] 53 | for _j in range(_i + 1, ndets): 54 | j = order[_j] 55 | if suppressed[j] == 1: 56 | continue 57 | xx1 = max(ix1, x1[j]) 58 | yy1 = max(iy1, y1[j]) 59 | xx2 = min(ix2, x2[j]) 60 | yy2 = min(iy2, y2[j]) 61 | w = max(0.0, xx2 - xx1 + 1) 62 | h = max(0.0, yy2 - yy1 + 1) 63 | inter = w * h 64 | ovr = inter / (iarea + areas[j] - inter) 65 | if ovr >= thresh: 66 | suppressed[j] = 1 67 | 68 | return keep 69 | -------------------------------------------------------------------------------- /DetJoint/yolo_v2/utils/nms/gpu_nms.hpp: -------------------------------------------------------------------------------- 1 | void _nms(int* keep_out, int* num_out, const float* boxes_host, int boxes_num, 2 | int boxes_dim, float nms_overlap_thresh, int device_id); 3 | -------------------------------------------------------------------------------- /DetJoint/yolo_v2/utils/nms/gpu_nms.pyx: -------------------------------------------------------------------------------- 1 | # -------------------------------------------------------- 2 | # Faster R-CNN 3 | # Copyright (c) 2015 Microsoft 4 | # Licensed under The MIT License [see LICENSE for details] 5 | # Written by Ross Girshick 6 | # -------------------------------------------------------- 7 | 8 | import numpy as np 9 | cimport numpy as np 10 | 11 | assert sizeof(int) == sizeof(np.int32_t) 12 | 13 | cdef extern from "gpu_nms.hpp": 14 | void _nms(np.int32_t*, int*, np.float32_t*, int, int, float, int) 15 | 16 | def gpu_nms(np.ndarray[np.float32_t, ndim=2] dets, np.float thresh, 17 | np.int32_t device_id=0): 18 | cdef int boxes_num = dets.shape[0] 19 | cdef int boxes_dim = dets.shape[1] 20 | cdef int num_out 21 | cdef np.ndarray[np.int32_t, ndim=1] \ 22 | keep = np.zeros(boxes_num, dtype=np.int32) 23 | cdef np.ndarray[np.float32_t, ndim=1] \ 24 | scores = dets[:, 4] 25 | cdef np.ndarray[np.int_t, ndim=1] \ 26 | order = scores.argsort()[::-1] 27 | cdef np.ndarray[np.float32_t, ndim=2] \ 28 | sorted_dets = dets[order, :] 29 | _nms(&keep[0], &num_out, &sorted_dets[0, 0], boxes_num, boxes_dim, thresh, device_id) 30 | keep = keep[:num_out] 31 | return list(order[keep]) 32 | -------------------------------------------------------------------------------- /DetJoint/yolo_v2/utils/nms/nms_kernel.cu: -------------------------------------------------------------------------------- 1 | // ------------------------------------------------------------------ 2 | // Faster R-CNN 3 | // Copyright (c) 2015 Microsoft 4 | // Licensed under The MIT License [see fast-rcnn/LICENSE for details] 5 | // Written by Shaoqing Ren 6 | // ------------------------------------------------------------------ 7 | 8 | #include "gpu_nms.hpp" 9 | #include 10 | #include 11 | 12 | #define CUDA_CHECK(condition) \ 13 | /* Code block avoids redefinition of cudaError_t error */ \ 14 | do { \ 15 | cudaError_t error = condition; \ 16 | if (error != cudaSuccess) { \ 17 | std::cout << cudaGetErrorString(error) << std::endl; \ 18 | } \ 19 | } while (0) 20 | 21 | #define DIVUP(m,n) ((m) / (n) + ((m) % (n) > 0)) 22 | int const threadsPerBlock = sizeof(unsigned long long) * 8; 23 | 24 | __device__ inline float devIoU(float const * const a, float const * const b) { 25 | float left = max(a[0], b[0]), right = min(a[2], b[2]); 26 | float top = max(a[1], b[1]), bottom = min(a[3], b[3]); 27 | float width = max(right - left + 1, 0.f), height = max(bottom - top + 1, 0.f); 28 | float interS = width * height; 29 | float Sa = (a[2] - a[0] + 1) * (a[3] - a[1] + 1); 30 | float Sb = (b[2] - b[0] + 1) * (b[3] - b[1] + 1); 31 | return interS / (Sa + Sb - interS); 32 | } 33 | 34 | __global__ void nms_kernel(const int n_boxes, const float nms_overlap_thresh, 35 | const float *dev_boxes, unsigned long long *dev_mask) { 36 | const int row_start = blockIdx.y; 37 | const int col_start = blockIdx.x; 38 | 39 | // if (row_start > col_start) return; 40 | 41 | const int row_size = 42 | min(n_boxes - row_start * threadsPerBlock, threadsPerBlock); 43 | const int col_size = 44 | min(n_boxes - col_start * threadsPerBlock, threadsPerBlock); 45 | 46 | __shared__ float block_boxes[threadsPerBlock * 5]; 47 | if (threadIdx.x < col_size) { 48 | block_boxes[threadIdx.x * 5 + 0] = 49 | dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 0]; 50 | block_boxes[threadIdx.x * 5 + 1] = 51 | dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 1]; 52 | block_boxes[threadIdx.x * 5 + 2] = 53 | dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 2]; 54 | block_boxes[threadIdx.x * 5 + 3] = 55 | dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 3]; 56 | block_boxes[threadIdx.x * 5 + 4] = 57 | dev_boxes[(threadsPerBlock * col_start + threadIdx.x) * 5 + 4]; 58 | } 59 | __syncthreads(); 60 | 61 | if (threadIdx.x < row_size) { 62 | const int cur_box_idx = threadsPerBlock * row_start + threadIdx.x; 63 | const float *cur_box = dev_boxes + cur_box_idx * 5; 64 | int i = 0; 65 | unsigned long long t = 0; 66 | int start = 0; 67 | if (row_start == col_start) { 68 | start = threadIdx.x + 1; 69 | } 70 | for (i = start; i < col_size; i++) { 71 | if (devIoU(cur_box, block_boxes + i * 5) > nms_overlap_thresh) { 72 | t |= 1ULL << i; 73 | } 74 | } 75 | const int col_blocks = DIVUP(n_boxes, threadsPerBlock); 76 | dev_mask[cur_box_idx * col_blocks + col_start] = t; 77 | } 78 | } 79 | 80 | void _set_device(int device_id) { 81 | int current_device; 82 | CUDA_CHECK(cudaGetDevice(¤t_device)); 83 | if (current_device == device_id) { 84 | return; 85 | } 86 | // The call to cudaSetDevice must come before any calls to Get, which 87 | // may perform initialization using the GPU. 88 | CUDA_CHECK(cudaSetDevice(device_id)); 89 | } 90 | 91 | void _nms(int* keep_out, int* num_out, const float* boxes_host, int boxes_num, 92 | int boxes_dim, float nms_overlap_thresh, int device_id) { 93 | _set_device(device_id); 94 | 95 | float* boxes_dev = NULL; 96 | unsigned long long* mask_dev = NULL; 97 | 98 | const int col_blocks = DIVUP(boxes_num, threadsPerBlock); 99 | 100 | CUDA_CHECK(cudaMalloc(&boxes_dev, 101 | boxes_num * boxes_dim * sizeof(float))); 102 | CUDA_CHECK(cudaMemcpy(boxes_dev, 103 | boxes_host, 104 | boxes_num * boxes_dim * sizeof(float), 105 | cudaMemcpyHostToDevice)); 106 | 107 | CUDA_CHECK(cudaMalloc(&mask_dev, 108 | boxes_num * col_blocks * sizeof(unsigned long long))); 109 | 110 | dim3 blocks(DIVUP(boxes_num, threadsPerBlock), 111 | DIVUP(boxes_num, threadsPerBlock)); 112 | dim3 threads(threadsPerBlock); 113 | nms_kernel<<>>(boxes_num, 114 | nms_overlap_thresh, 115 | boxes_dev, 116 | mask_dev); 117 | 118 | std::vector mask_host(boxes_num * col_blocks); 119 | CUDA_CHECK(cudaMemcpy(&mask_host[0], 120 | mask_dev, 121 | sizeof(unsigned long long) * boxes_num * col_blocks, 122 | cudaMemcpyDeviceToHost)); 123 | 124 | std::vector remv(col_blocks); 125 | memset(&remv[0], 0, sizeof(unsigned long long) * col_blocks); 126 | 127 | int num_to_keep = 0; 128 | for (int i = 0; i < boxes_num; i++) { 129 | int nblock = i / threadsPerBlock; 130 | int inblock = i % threadsPerBlock; 131 | 132 | if (!(remv[nblock] & (1ULL << inblock))) { 133 | keep_out[num_to_keep++] = i; 134 | unsigned long long *p = &mask_host[0] + i * col_blocks; 135 | for (int j = nblock; j < col_blocks; j++) { 136 | remv[j] |= p[j]; 137 | } 138 | } 139 | } 140 | *num_out = num_to_keep; 141 | 142 | CUDA_CHECK(cudaFree(boxes_dev)); 143 | CUDA_CHECK(cudaFree(mask_dev)); 144 | } 145 | -------------------------------------------------------------------------------- /DetJoint/yolo_v2/utils/nms/py_cpu_nms.py: -------------------------------------------------------------------------------- 1 | # -------------------------------------------------------- 2 | # Fast R-CNN 3 | # Copyright (c) 2015 Microsoft 4 | # Licensed under The MIT License [see LICENSE for details] 5 | # Written by Ross Girshick 6 | # -------------------------------------------------------- 7 | 8 | import numpy as np 9 | 10 | def py_cpu_nms(dets, thresh): 11 | """Pure Python NMS baseline.""" 12 | x1 = dets[:, 0] 13 | y1 = dets[:, 1] 14 | x2 = dets[:, 2] 15 | y2 = dets[:, 3] 16 | scores = dets[:, 4] 17 | 18 | areas = (x2 - x1 + 1) * (y2 - y1 + 1) 19 | order = scores.argsort()[::-1] 20 | 21 | keep = [] 22 | while order.size > 0: 23 | i = order[0] 24 | keep.append(i) 25 | xx1 = np.maximum(x1[i], x1[order[1:]]) 26 | yy1 = np.maximum(y1[i], y1[order[1:]]) 27 | xx2 = np.minimum(x2[i], x2[order[1:]]) 28 | yy2 = np.minimum(y2[i], y2[order[1:]]) 29 | 30 | w = np.maximum(0.0, xx2 - xx1 + 1) 31 | h = np.maximum(0.0, yy2 - yy1 + 1) 32 | inter = w * h 33 | ovr = inter / (areas[i] + areas[order[1:]] - inter) 34 | 35 | inds = np.where(ovr <= thresh)[0] 36 | order = order[inds + 1] 37 | 38 | return keep 39 | -------------------------------------------------------------------------------- /DetJoint/yolo_v2/utils/nms_wrapper.py: -------------------------------------------------------------------------------- 1 | # -------------------------------------------------------- 2 | # Fast R-CNN 3 | # Copyright (c) 2015 Microsoft 4 | # Licensed under The MIT License [see LICENSE for details] 5 | # Written by Ross Girshick 6 | # -------------------------------------------------------- 7 | 8 | 9 | 10 | from .nms.cpu_nms import cpu_nms 11 | from .nms.gpu_nms import gpu_nms 12 | 13 | 14 | def nms(dets, thresh, force_cpu=True): 15 | """Dispatch to either CPU or GPU NMS implementations.""" 16 | 17 | if dets.shape[0] == 0: 18 | return [] 19 | if force_cpu: 20 | return cpu_nms(dets, thresh) 21 | return gpu_nms(dets, thresh) 22 | -------------------------------------------------------------------------------- /DetJoint/yolo_v2/utils/timer.py: -------------------------------------------------------------------------------- 1 | # -------------------------------------------------------- 2 | # Fast R-CNN 3 | # Copyright (c) 2015 Microsoft 4 | # Licensed under The MIT License [see LICENSE for details] 5 | # Written by Ross Girshick 6 | # -------------------------------------------------------- 7 | 8 | import time 9 | 10 | 11 | class Timer(object): 12 | """A simple timer.""" 13 | def __init__(self): 14 | self.total_time = 0. 15 | self.calls = 0 16 | self.start_time = 0. 17 | self.diff = 0. 18 | self.average_time = 0. 19 | 20 | def tic(self): 21 | # using time.time instead of time.clock because time time.clock 22 | # does not normalize for multithreading 23 | self.start_time = time.time() 24 | 25 | def toc(self, average=True): 26 | self.diff = time.time() - self.start_time 27 | self.total_time += self.diff 28 | self.calls += 1 29 | self.average_time = self.total_time / self.calls 30 | if average: 31 | return self.average_time 32 | else: 33 | return self.diff 34 | 35 | def clear(self): 36 | self.total_time = 0. 37 | self.calls = 0 38 | self.start_time = 0. 39 | self.diff = 0. 40 | self.average_time = 0. 41 | -------------------------------------------------------------------------------- /DetJoint/yolo_v2/utils/yolo.py: -------------------------------------------------------------------------------- 1 | import cv2 2 | import os 3 | import numpy as np 4 | from .im_transform import imcv2_affine_trans, imcv2_recolor 5 | # from box import BoundBox, box_iou, prob_compare 6 | from .nms_wrapper import nms 7 | from .cython_yolo import yolo_to_bbox 8 | 9 | 10 | def clip_boxes(boxes, im_shape): 11 | """ 12 | Clip boxes to image boundaries. 13 | """ 14 | if boxes.shape[0] == 0: 15 | return boxes 16 | 17 | # x1 >= 0 18 | boxes[:, 0::4] = np.maximum(np.minimum(boxes[:, 0::4], im_shape[1] - 1), 0) 19 | # y1 >= 0 20 | boxes[:, 1::4] = np.maximum(np.minimum(boxes[:, 1::4], im_shape[0] - 1), 0) 21 | # x2 < im_shape[1] 22 | boxes[:, 2::4] = np.maximum(np.minimum(boxes[:, 2::4], im_shape[1] - 1), 0) 23 | # y2 < im_shape[0] 24 | boxes[:, 3::4] = np.maximum(np.minimum(boxes[:, 3::4], im_shape[0] - 1), 0) 25 | return boxes 26 | 27 | 28 | def nms_detections(pred_boxes, scores, nms_thresh): 29 | dets = np.hstack((pred_boxes, 30 | scores[:, np.newaxis])).astype(np.float32) 31 | keep = nms(dets, nms_thresh) 32 | return keep 33 | 34 | 35 | def _offset_boxes(boxes, im_shape, scale, offs, flip): 36 | if len(boxes) == 0: 37 | return boxes 38 | boxes = np.asarray(boxes, dtype=np.float) 39 | boxes *= scale 40 | boxes[:, 0::2] -= offs[0] 41 | boxes[:, 1::2] -= offs[1] 42 | boxes = clip_boxes(boxes, im_shape) 43 | 44 | if flip: 45 | boxes_x = np.copy(boxes[:, 0]) 46 | boxes[:, 0] = im_shape[1] - boxes[:, 2] 47 | boxes[:, 2] = im_shape[1] - boxes_x 48 | 49 | return boxes 50 | 51 | 52 | def preprocess_train(data): 53 | im_path, blob, inp_size = data 54 | boxes, gt_classes = blob['boxes'], blob['gt_classes'] 55 | 56 | im = cv2.imread(im_path) 57 | ori_im = np.copy(im) 58 | 59 | im, trans_param = imcv2_affine_trans(im) 60 | scale, offs, flip = trans_param 61 | boxes = _offset_boxes(boxes, im.shape, scale, offs, flip) 62 | 63 | if inp_size is not None: 64 | w, h = inp_size 65 | boxes[:, 0::2] *= float(w) / im.shape[1] 66 | boxes[:, 1::2] *= float(h) / im.shape[0] 67 | im = cv2.resize(im, (w, h)) 68 | im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB) 69 | im = imcv2_recolor(im) 70 | 71 | boxes = np.asarray(boxes, dtype=np.int) 72 | return im, boxes, gt_classes, [], ori_im 73 | 74 | 75 | def preprocess_test(data): 76 | 77 | im, _, inp_size = data 78 | if isinstance(im, str): 79 | im = cv2.imread(im) 80 | ori_im = np.copy(im) 81 | 82 | if inp_size is not None: 83 | w, h = inp_size 84 | im = cv2.resize(im, (w, h)) 85 | im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB) 86 | im = im / 255. 87 | 88 | return im, [], [], [], ori_im 89 | 90 | 91 | def postprocess_bbox(bbox_pred, iou_pred, prob_pred, im_shape, cfg, thresh=0.05): 92 | """ 93 | bbox_pred: prcossed (bsize, HxW, num_anchors, 4) ndarray of float (xs, ys, xe, ye) 94 | iou_pred: (bsize, HxW, num_anchors, 1) 95 | prob_pred: (bsize, HxW, num_anchors, num_classes) 96 | """ 97 | num_classes, num_anchors = cfg.num_classes, cfg.num_anchors 98 | anchors = cfg.anchors 99 | 100 | assert bbox_pred.shape[0] == 1, 'postprocess only support one image per batch' 101 | 102 | bbox_pred = np.reshape(bbox_pred, [-1, 4]) 103 | bbox_pred = bbox_pred.astype(np.int) 104 | iou_pred = np.reshape(iou_pred, [-1]) 105 | prob_pred = np.reshape(prob_pred, [-1, num_classes]) 106 | 107 | cls_inds = np.argmax(prob_pred, axis=1) 108 | prob_pred = prob_pred[(np.arange(prob_pred.shape[0]), cls_inds)] 109 | scores = iou_pred * prob_pred 110 | 111 | # threshold 112 | 113 | keep = np.where(scores >= thresh) 114 | bbox_pred = bbox_pred[keep] 115 | scores = scores[keep] 116 | cls_inds = cls_inds[keep] 117 | 118 | keep = nms_detections(bbox_pred, scores, 0.3) 119 | bbox_pred = bbox_pred[keep] 120 | scores = scores[keep] 121 | cls_inds = cls_inds[keep] 122 | 123 | # clip 124 | bbox_pred = clip_boxes(bbox_pred, im_shape) 125 | 126 | return bbox_pred, scores, cls_inds 127 | 128 | def _bbox_targets_perimage(im_shape, gt_boxes, cls_inds, dontcare_areas, cfg): 129 | # num_classes, num_anchors = cfg.num_classes, cfg.num_anchors 130 | # anchors = cfg.anchors 131 | H, W = cfg.out_size 132 | gt_boxes = np.asarray(gt_boxes, dtype=np.float) 133 | # TODO: dontcare areas 134 | dontcare_areas = np.asarray(dontcare_areas, dtype=np.float) 135 | 136 | # locate the cell of each gt_boxe 137 | cell_w = float(im_shape[1]) / W 138 | cell_h = float(im_shape[0]) / H 139 | cx = (gt_boxes[:, 0] + gt_boxes[:, 2]) * 0.5 / cell_w 140 | cy = (gt_boxes[:, 1] + gt_boxes[:, 3]) * 0.5 / cell_h 141 | cell_inds = np.floor(cy) * W + np.floor(cx) 142 | cell_inds = cell_inds.astype(np.int) 143 | 144 | # [x1, y1, x2, y2], [class] 145 | # gt_boxes[:, 0::2] /= im_shape[1] 146 | # gt_boxes[:, 1::2] /= im_shape[0] 147 | # gt_boxes[:, 0] = cx - np.floor(cx) 148 | # gt_boxes[:, 1] = cy - np.floor(cy) 149 | # gt_boxes[:, 2] = (gt_boxes[:, 2] - gt_boxes[:, 0]) / im_shape[1] 150 | # gt_boxes[:, 3] = (gt_boxes[:, 3] - gt_boxes[:, 1]) / im_shape[0] 151 | 152 | bbox_target = [[] for _ in range(H*W)] 153 | cls_target = [[] for _ in range(H*W)] 154 | for i, ind in enumerate(cell_inds): 155 | bbox_target[ind].append(gt_boxes[i]) 156 | cls_target[ind].append(cls_inds[i]) 157 | return bbox_target, cls_target 158 | 159 | 160 | def get_bbox_targets(images, gt_boxes, cls_inds, dontcares, cfg): 161 | bbox_targets = [] 162 | cls_targets = [] 163 | for i, im in enumerate(images): 164 | bbox_target, cls_target = _bbox_targets_perimage(im.shape, gt_boxes[i], cls_inds[i], dontcares[i], cfg) 165 | bbox_targets.append(bbox_target) 166 | cls_targets.append(cls_target) 167 | return bbox_targets, cls_targets 168 | 169 | 170 | def draw_detection(im, bboxes, scores, cls_inds, cfg, thr=0.3): 171 | # draw image 172 | colors = cfg.colors 173 | labels = cfg.label_names 174 | 175 | imgcv = np.copy(im) 176 | h, w, _ = imgcv.shape 177 | for i, box in enumerate(bboxes): 178 | if scores[i] < thr: 179 | continue 180 | cls_indx = cls_inds[i] 181 | 182 | thick = int((h + w) / 300) 183 | cv2.rectangle(imgcv, 184 | (box[0], box[1]), (box[2], box[3]), 185 | colors[cls_indx], thick) 186 | mess = '%s: %.3f' % (labels[cls_indx], scores[i]) 187 | cv2.putText(imgcv, mess, (box[0], box[1] - 12), 188 | 0, 1e-3 * h, colors[cls_indx], thick // 3) 189 | 190 | return imgcv 191 | -------------------------------------------------------------------------------- /DetJoint/yolo_v2/utils/yolo.pyx: -------------------------------------------------------------------------------- 1 | cimport cython 2 | import numpy as np 3 | cimport numpy as np 4 | 5 | DTYPE = np.float 6 | ctypedef np.float_t DTYPE_t 7 | 8 | cdef extern from "math.h": 9 | double abs(double m) 10 | double log(double x) 11 | 12 | 13 | def yolo_to_bbox( 14 | np.ndarray[DTYPE_t, ndim=4] bbox_pred, 15 | np.ndarray[DTYPE_t, ndim=2] anchors,int H, int W, DTYPE_t x_ratio, DTYPE_t y_ratio): 16 | return yolo_to_bbox_c(bbox_pred, anchors,H, W, x_ratio, y_ratio) 17 | 18 | cdef yolo_to_bbox_c( 19 | np.ndarray[DTYPE_t, ndim=4] bbox_pred, 20 | np.ndarray[DTYPE_t, ndim=2] anchors,int H, int W, DTYPE_t x_ratio, DTYPE_t y_ratio): 21 | """ 22 | Parameters 23 | ---------- 24 | bbox_pred: (bsize, HxW, num_anchors, 4) ndarray of float (sig(tx), sig(ty), exp(tw), exp(th)) 25 | anchors: (num_anchors, 2) (pw, ph) 26 | Returns 27 | ------- 28 | bbox_out: (HxWxnum_anchors, 4) ndarray of bbox (x1, y1, x2, y2) rescaled to (0, 1) 29 | """ 30 | cdef unsigned int bsize = bbox_pred.shape[0] 31 | cdef unsigned int num_anchors = anchors.shape[0] 32 | cdef np.ndarray[DTYPE_t, ndim=4] bbox_out = np.zeros((bsize, H*W, num_anchors, 4), dtype=DTYPE) 33 | 34 | cdef DTYPE_t cx, cy, bw, bh 35 | cdef unsigned int row, col, a, ind 36 | 37 | for b in range(bsize): 38 | for row in range(H): 39 | for col in range(W): 40 | ind = row * W + col 41 | for a in range(num_anchors): 42 | cx = (bbox_pred[b, ind, a, 0] + col) * x_ratio 43 | cy = (bbox_pred[b, ind, a, 1] + row) * y_ratio 44 | 45 | bw = bbox_pred[b, ind, a, 2] * anchors[a][0] * 0.5 46 | bh = bbox_pred[b, ind, a, 3] * anchors[a][1] * 0.5 47 | 48 | bbox_out[b, ind, a, 0] = cx - bw 49 | bbox_out[b, ind, a, 1] = cy - bh 50 | bbox_out[b, ind, a, 2] = cx + bw 51 | bbox_out[b, ind, a, 3] = cy + bh 52 | 53 | return bbox_out 54 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2018 Pingjun Chen 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Knee Osteoarthritis Analysis 2 | 3 | #### Pipeline of knee osteoarthritis grading pipeline, which includes knee joints detection and knee OA grading. 4 | 5 | ![klg_pipeline](klg_pipeline.png) 6 | 7 | ### DetJoint 8 | - Detecting two knee joints in X-ray images using a customized YOLOv2 model. 9 | 10 | ### ClsKL 11 | - Classifying the KL grade of detected knee joints with a novel ordinal loss. 12 | 13 | 14 | ## Data & Models 15 | Knee joint detection (DetJoint) and KL grading (ClsKL) training/testing datasets, as well as best models, can be downloaded from [KneeXrayData](https://data.mendeley.com/datasets/56rmx5bjcr/1), around 7G. 16 | 17 | 18 | ## Citation 19 | Please consider `cite` the paper if you use the code or data for your research. 20 | ``` 21 | @article{chen2019fully, 22 | title={Fully Automatic Knee Osteoarthritis Severity Grading Using Deep Neural Networks with a Novel Ordinal Loss}, 23 | author={Chen, Pingjun and Gao, Linlin and Shi, Xiaoshuang and Allen, Kyle and Yang Lin}, 24 | journal={Computerized Medical Imaging and Graphics},, 25 | volume={75}, 26 | pages={84--92}, 27 | year={2019}, 28 | doi={https://doi.org/10.1016/j.compmedimag.2019.06.002}, 29 | publisher={Elsevier} 30 | } 31 | ``` 32 | -------------------------------------------------------------------------------- /klg_pipeline.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/PingjunChen/GradingKneeOA/dc19c2462111c3c5b27066e738fc11ddf57a6407/klg_pipeline.png --------------------------------------------------------------------------------