Spaces:
Sleeping
Sleeping
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| import os | |
| import pandas as pd | |
| import numpy as np | |
| from tqdm import tqdm | |
| from .file_io import FileLoader, FileSaver, get_folder_loader | |
| from .models import get_model | |
| from .extraction import Extractor | |
| from .helper import set_gpus, get_available_devices, find_max_overlap | |
| class CXAS(nn.Module): | |
| def __init__(self, | |
| model_name:str='UNet_ResNet50_default', | |
| gpus:str=''): | |
| """ | |
| Create Chest X-Ray anatomy segmentation model | |
| Parameters | |
| ---------- | |
| model_name: which model/weight to load, if weights are not stored, they will be downloaded at '~/weights/' | |
| gpus: on which gpu to perform inference on | |
| """ | |
| super(CXAS,self).__init__() | |
| self.gpus = set_gpus(gpus) | |
| self.model = get_model(model_name, gpus) | |
| self.fileloader = FileLoader(gpus) | |
| self.filesaver = FileSaver() | |
| self.extractor = Extractor() | |
| self.eval() | |
| def seg(self, filename) -> np.array: | |
| """ | |
| Create segmentation of image file, can store predictions in desired directory in desired format | |
| Parameters | |
| ---------- | |
| filename: path of file to process, currently supported types [.dcm, .jpg, .png] | |
| do_store: bool indicating whether to store prediction | |
| output_directory: desired path of output directory | |
| storage_type: desired type to store segmentation prediction as, currently supported types [dicom-seg, jpg, png, npy, npz, json] | |
| Returns | |
| ------- | |
| prediction: model output dictionary containing [feats: network features , logits: unnormalized network logit scores, data: input data, segmentation_preds: thresholded multi-label segmentations] | |
| """ | |
| assert os.path.isfile(filename) | |
| file_dict = self.fileloader.load_file(filename) | |
| file_dict['filename'] = [file_dict['filename']] | |
| file_dict['file_size'] = [file_dict['file_size']] | |
| with torch.no_grad(): | |
| predictions = self.model(file_dict) | |
| for i in range(len(predictions['filename'])): | |
| mask = self.resize_to_numpy( | |
| segmentation = predictions['segmentation_preds'][i], | |
| file_size = predictions['file_size'][i] | |
| ) # output resized np.ndarray with img_size | |
| return mask | |
| def process_file(self, | |
| filename: str, | |
| do_store:bool=False, | |
| output_directory:str='./', | |
| create:bool = False, | |
| storage_type:str='npy' | |
| ) -> np.array: | |
| """ | |
| Create segmentation of image file, can store predictions in desired directory in desired format | |
| Parameters | |
| ---------- | |
| filename: path of file to process, currently supported types [.dcm, .jpg, .png] | |
| do_store: bool indicating whether to store prediction | |
| output_directory: desired path of output directory | |
| storage_type: desired type to store segmentation prediction as, currently supported types [dicom-seg, jpg, png, npy, npz, json] | |
| Returns | |
| ------- | |
| prediction: model output dictionary containing [feats: network features , logits: unnormalized network logit scores, data: input data, segmentation_preds: thresholded multi-label segmentations] | |
| """ | |
| assert os.path.isfile(filename) | |
| if not create: | |
| assert os.path.isdir(output_directory) | |
| else: | |
| os.makedirs(output_directory, exist_ok=True) | |
| file_dict = self.fileloader.load_file(filename) | |
| file_dict['filename'] = [file_dict['filename']] | |
| file_dict['file_size'] = [file_dict['file_size']] | |
| with torch.no_grad(): | |
| predictions = self.model(file_dict) | |
| if do_store: | |
| self.store_prediction(predictions, output_directory, storage_type) | |
| return predictions | |
| def process_folder(self, | |
| input_directory_name: str, | |
| output_directory:str, | |
| storage_type:str = 'npy', | |
| create:bool = False, | |
| batch_size:int=1 | |
| ) -> None: | |
| """ | |
| Create segmentations for all image files in directory, stores predictions in desired output directory in desired format | |
| Parameters | |
| ---------- | |
| input_directory_name: path of file to process, currently supported types [.dcm, .jpg, .png] | |
| output_directory: desired path of output directory | |
| storage_type: desired type to store segmentation prediction as, currently supported types [dicom-seg, jpg, png, npy, npz, json] | |
| create: whether to create the output directory | |
| batch_size: batch size used for the forward passes of the model | |
| """ | |
| assert os.path.isdir(input_directory_name) | |
| if not create: | |
| assert os.path.isdir(output_directory) | |
| else: | |
| os.makedirs(output_directory, exist_ok=True) | |
| dataloader = get_folder_loader(input_directory_name, self.gpus, batch_size, ) | |
| if storage_type == 'json': | |
| from .io_utils.create_annotations import get_coco_json_format, \ | |
| create_category_annotation | |
| from .io_utils.mask_to_coco import binary_mask_to_rle, toBox, mask_to_annotation | |
| from .label_mapper import id2label_dict, category_ids | |
| import json | |
| coco_format = get_coco_json_format() | |
| coco_format["categories"] = create_category_annotation(category_ids) | |
| coco_format["images"] = [] | |
| coco_format["annotations"] = [] | |
| img_id = 1 | |
| base_ann_id = 1 | |
| for file_dict in tqdm(dataloader): | |
| if (type(self.gpus) is list) and len(self.gpus) > 0: | |
| file_dict['data'] = file_dict['data'].to('{}'.format(self.gpus[0])) | |
| else: | |
| file_dict['data'] = file_dict['data'].to(self.gpus) | |
| with torch.no_grad(): | |
| predictions = self.model(file_dict) | |
| if storage_type == 'json': | |
| for i in range(len(predictions['filename'])): | |
| mask = self.resize_to_numpy( | |
| segmentation = predictions['segmentation_preds'][i], | |
| file_size = predictions['file_size'][i] | |
| ) | |
| annotations = mask_to_annotation( | |
| mask = mask, | |
| base_ann_id = base_ann_id, | |
| img_id = img_id | |
| ) | |
| base_ann_id += len(annotations) | |
| coco_format["images"] += [{'id':img_id, 'file_name': predictions['filename'][i]}] | |
| coco_format["annotations"] += annotations | |
| img_id += 1 | |
| else: | |
| self.store_prediction(predictions, output_directory, storage_type) | |
| if storage_type == 'json': | |
| os.makedirs(output_directory,exist_ok=True) | |
| out_path = os.path.join(output_directory, input_directory_name.split('/')[-1]+'.json') | |
| with open(out_path,"w") as outfile: | |
| json.dump(coco_format, outfile) | |
| def store_prediction(self, | |
| predictions: dict, | |
| output_directory:str, | |
| storage_type:str) -> None: | |
| """ | |
| Store all elements in batch | |
| Parameters | |
| ---------- | |
| predictions: model output dictionary containing [feats: network features , logits: unnormalized network logit scores, data: input data, segmentation_preds: thresholded multi-label segmentations] | |
| output_directory: desired path of output directory | |
| storage_type: desired type to store segmentation prediction as, currently supported types [dicom-seg, jpg, png, npy, npz, json] | |
| """ | |
| for i in range(len(predictions['filename'])): | |
| pred = self.resize_to_numpy( | |
| segmentation = predictions['segmentation_preds'][i], | |
| file_size = predictions['file_size'][i] | |
| ) | |
| self.filesaver.save_prediction(pred, output_directory, predictions['filename'][i], storage_type) | |
| def resize_to_numpy(self, | |
| segmentation: torch.Tensor, | |
| file_size, | |
| ) -> torch.Tensor: | |
| """ | |
| Resize binary torch prediction mask to desired size | |
| Parameters | |
| ---------- | |
| segmentation: model output dictionary containing [feats: network features , logits: unnormalized network logit scores, data: input data, segmentation_preds: thresholded multi-label segmentations] | |
| file_size: desired path of output directory | |
| """ | |
| pred = segmentation.float() | |
| pred = F.interpolate(pred.unsqueeze(0), file_size, mode='nearest') | |
| pred = pred[0].bool().to('cpu').numpy() | |
| return pred | |
| def extract_features_for_file(self, | |
| filename: str, | |
| feat_to_extract: str, | |
| draw:bool = False, | |
| create:bool = False, | |
| do_store:bool=False, | |
| output_directory:str='./', | |
| storage_type:str='npy' | |
| ) -> dict: | |
| """ | |
| Create segmentation of image file and extract features in relation to the segmentation. Can store predictions in desired output directory in desired format. | |
| Parameters | |
| ---------- | |
| filename: path of file to process, currently supported types [.dcm, .jpg, .png] | |
| feat_to_extract: which features to extract in relation to the segmentation | |
| draw: draw the origin of the features | |
| create: whether to create the output directory | |
| do_store: bool indicating whether to store prediction | |
| output_directory: desired path of output directory | |
| storage_type: desired type to store segmentation prediction as, currently supported types [dicom-seg, jpg, png, npy, npz, json] | |
| Returns | |
| ------- | |
| features: extracted feature score and if so designated its visualization | |
| """ | |
| assert os.path.isfile(filename) | |
| if not create: | |
| assert os.path.isdir(output_directory) | |
| else: | |
| os.makedirs(output_directory, exist_ok=True) | |
| predictions = self.process_file( | |
| filename, | |
| do_store = do_store, | |
| output_directory = output_directory, | |
| storage_type = storage_type, | |
| ) | |
| feat_dict = self.extractor.extract( | |
| file = predictions['segmentation_preds'][0].cpu().bool().numpy(), | |
| method = feat_to_extract, | |
| draw=draw, | |
| ) | |
| if 'score' in feat_dict.keys(): | |
| print('The {} for the file {} is '.format(feat_to_extract, filename.split('/')[-1]),feat_dict['score']) | |
| if do_store: | |
| scores = [{**{key:feat_dict[key] for key in feat_dict.keys() if key != 'drawing'}, | |
| 'filename': predictions['filename'][0], | |
| }] | |
| pd.DataFrame(scores).to_csv(os.path.join(output_directory, filename.split('/')[-1].split('.')[0]+'.csv')) | |
| return feat_dict | |
| def extract_features_for_folder(self, | |
| input_directory_name: str, | |
| output_directory:str, | |
| feat_to_extract: str, | |
| create:bool = False, | |
| store_pred: bool = False, | |
| storage_type:str = 'npy', | |
| batch_size:int=1 | |
| ) -> None: | |
| """ | |
| Create segmentation of image file and extract features in relation to the segmentation. Can store predictions in desired output directory in desired format. | |
| Parameters | |
| ---------- | |
| input_directory_name: path of file to process, currently supported types [.dcm, .jpg, .png] | |
| output_directory: desired path of output directory | |
| storage_type: desired type to store segmentation prediction as, currently supported types [dicom-seg, jpg, png, npy, npz, json] | |
| create: whether to create the output directory | |
| batch_size: batch size used for the forward passes of the model | |
| feat_to_extract: which features to extract in relation to the segmentation | |
| draw: draw the origin of the features | |
| create: whether to create the output directory | |
| store_pred: bool indicating whether to store prediction | |
| """ | |
| assert os.path.isdir(input_directory_name) | |
| if not create: | |
| assert os.path.isdir(output_directory) | |
| else: | |
| os.makedirs(output_directory, exist_ok=True) | |
| scores = [] | |
| dataloader = get_folder_loader(input_directory_name, self.gpus, batch_size, ) | |
| if (storage_type == 'json') and store_pred: | |
| from cxas.io_utils.create_annotations import get_coco_json_format, \ | |
| create_category_annotation | |
| from cxas.io_utils.mask_to_coco import binary_mask_to_rle, toBox, mask_to_annotation | |
| from cxas.label_mapper import id2label_dict, category_ids | |
| import json | |
| coco_format = get_coco_json_format() | |
| coco_format["categories"] = create_category_annotation(category_ids) | |
| coco_format["images"] = [] | |
| coco_format["annotations"] = [] | |
| img_id = 1 | |
| base_ann_id = 1 | |
| for file_dict in tqdm(dataloader): | |
| if len(self.gpus)>0: | |
| file_dict['data'] = file_dict['data'].to('cuda:{}'.format(self.gpus[0])) | |
| with torch.no_grad(): | |
| predictions = self.model(file_dict) | |
| for i in range(len(predictions['segmentation_preds'])): | |
| extractions = self.extractor.extract( | |
| file = predictions['segmentation_preds'][i].cpu().bool().numpy(), | |
| method = feat_to_extract, | |
| draw=False, | |
| ) | |
| scores += [{**{key:extractions[key] for key in extractions.keys() if key != 'drawing'}, | |
| 'filename': predictions['filename'][i], | |
| }] | |
| if store_pred: | |
| if storage_type == 'json': | |
| for i in range(len(predictions['filename'])): | |
| mask = self.resize_to_numpy( | |
| segmentation = predictions['segmentation_preds'][i], | |
| file_size = predictions['file_size'][i] | |
| ) | |
| annotations = mask_to_annotation( | |
| mask = mask, | |
| base_ann_id = base_ann_id, | |
| img_id = img_id | |
| ) | |
| base_ann_id += len(annotations) | |
| coco_format["images"] += [{'id':img_id, 'file_name': predictions['filename'][i]}] | |
| coco_format["annotations"] += annotations | |
| img_id += 1 | |
| else: | |
| self.store_prediction(predictions, output_directory, storage_type) | |
| pd.DataFrame(scores).to_csv(os.path.join(output_directory, input_directory_name.split('/')[-1]+'.csv' if input_directory_name[-1]!='/' else | |
| input_directory_name[:-1].split('/')[-1]+'.csv')) | |
| if (storage_type == 'json') and store_pred: | |
| os.makedirs(output_directory,exist_ok=True) | |
| out_path = os.path.join(output_directory, | |
| input_directory_name.split('/')[-1] if input_directory_name[-1]!='/' else | |
| input_directory_name[:-1].split('/')[-1]+'.json') | |
| with open(out_path,"w") as outfile: | |
| json.dump(coco_format, outfile) | |
| def forward(self, image_batch) -> dict: | |
| """ | |
| Forward pass function for processing image batches. | |
| Args: | |
| image_batch (dict or numpy.ndarray): Input image batch. If it's not a dictionary, | |
| assumes the input is a numpy array and wraps it in a dictionary under the key 'data'. | |
| Returns: | |
| dict: Result of forward pass through the model. | |
| """ | |
| # If the input is not a dictionary, wrap it in a dictionary under the key 'data' | |
| if not isinstance(image_batch, dict): | |
| image_batch = {'data': image_batch} | |
| # Perform forward pass through the model and return the result | |
| return self.model(image_batch) | |