Lewislou commited on
Commit
367aed9
·
1 Parent(s): d4776d1

Delete cell_sribd_model.py

Browse files
Files changed (1) hide show
  1. cell_sribd_model.py +0 -99
cell_sribd_model.py DELETED
@@ -1,99 +0,0 @@
1
-
2
- import os
3
- join = os.path.join
4
- import argparse
5
- import numpy as np
6
- import torch
7
- import torch.nn as nn
8
- from collections import OrderedDict
9
- from torchvision import datasets, models, transforms
10
- from classifiers import resnet10, resnet18
11
-
12
- from utils_modify import sliding_window_inference,sliding_window_inference_large,__proc_np_hv
13
- from PIL import Image
14
- import torch.nn.functional as F
15
- from skimage import io, segmentation, morphology, measure, exposure
16
- import tifffile as tif
17
- from models.flexible_unet_convnext import FlexibleUNet_star,FlexibleUNet_hv
18
- from transformers import PretrainedConfig
19
- from typing import List
20
- from transformers import PreTrainedModel
21
- from huggingface_hub import PyTorchModelHubMixin
22
- from torch import nn
23
- class ModelConfig(PretrainedConfig):
24
- model_type = "cell_sribd"
25
- def __init__(
26
- self,
27
- version = 1,
28
- input_channels: int = 3,
29
- roi_size: int = 512,
30
- overlap: float = 0.5,
31
- device: str = 'cpu',
32
- **kwargs,
33
- ):
34
-
35
- self.device = device
36
- self.roi_size = (roi_size, roi_size)
37
- self.input_channels = input_channels
38
- self.overlap = overlap
39
- self.np_thres, self.ksize, self.overall_thres, self.obj_size_thres = 0.6, 15, 0.4, 100
40
- self.n_rays = 32
41
- self.sw_batch_size = 4
42
- self.num_classes= 4
43
- self.block_size = 2048
44
- self.min_overlap = 128
45
- self.context = 128
46
- super().__init__(**kwargs)
47
-
48
-
49
- class MyModel(PreTrainedModel):
50
- config_class = ModelConfig
51
- #print(config.input_channels)
52
- def __init__(self, config):
53
- super().__init__(config)
54
- #print(config.input_channels)
55
- self.cls_model = resnet18()
56
- self.model0 = FlexibleUNet_star(in_channels=config.input_channels,out_channels=config.n_rays+1,backbone='convnext_small',pretrained=False,n_rays=config.n_rays,prob_out_channels=1,)
57
- self.model1 = FlexibleUNet_star(in_channels=config.input_channels,out_channels=config.n_rays+1,backbone='convnext_small',pretrained=False,n_rays=config.n_rays,prob_out_channels=1,)
58
- self.model2 = FlexibleUNet_star(in_channels=config.input_channels,out_channels=config.n_rays+1,backbone='convnext_small',pretrained=False,n_rays=config.n_rays,prob_out_channels=1,)
59
- self.model3 = FlexibleUNet_hv(in_channels=config.input_channels,out_channels=2+2,backbone='convnext_small',pretrained=False,n_rays=2,prob_out_channels=2,)
60
- self.preprocess=transforms.Compose([
61
- transforms.Resize(size=256),
62
- transforms.CenterCrop(size=224),
63
- transforms.ToTensor(),
64
- transforms.Normalize([0.485, 0.456, 0.406],[0.229, 0.224, 0.225])])
65
- def load_checkpoints(self,checkpoints):
66
- self.cls_model.load_state_dict(checkpoints['cls_model'])
67
- self.model0.load_state_dict(checkpoints['class1_model']['model_state_dict'])
68
- self.model1.load_state_dict(checkpoints['class2_model']['model_state_dict'])
69
- self.model2.load_state_dict(checkpoints['class3_model']['model_state_dict'])
70
- self.model3.load_state_dict(checkpoints['class4_model'])
71
-
72
- def forward(self, pre_img_data):
73
- inputs=self.preprocess(Image.fromarray(pre_img_data)).unsqueeze(0)
74
- outputs = self.cls_model(inputs)
75
- _, preds = torch.max(outputs, 1)
76
- label=preds[0].cpu().numpy()
77
- test_npy01 = pre_img_data
78
- if label in [0,1,2] or img_data.shape[0] > 4000:
79
- if label == 0:
80
- output_label = sliding_window_inference_large(test_npy01,config.block_size,config.min_overlap,config.context, config.roi_size,config.sw_batch_size,predictor=self.model0,device=config.device)
81
- elif label == 1:
82
- output_label = sliding_window_inference_large(test_npy01,config.block_size,config.min_overlap,config.context, config.roi_size,config.sw_batch_size,predictor=self.model1,device=config.device)
83
- elif label == 2:
84
- output_label = sliding_window_inference_large(test_npy01,config.block_size,config.min_overlap,config.context, config.roi_size,config.sw_batch_size,predictor=self.model2,device=config.device)
85
- else:
86
- test_tensor = torch.from_numpy(np.expand_dims(test_npy01, 0)).permute(0, 3, 1, 2).type(torch.FloatTensor)
87
-
88
- output_hv, output_np = sliding_window_inference(test_tensor, config.roi, config.sw_batch_size, self.model3, overlap=config.overlap,device=config.device)
89
- pred_dict = {'np': output_np, 'hv': output_hv}
90
- pred_dict = OrderedDict(
91
- [[k, v.permute(0, 2, 3, 1).contiguous()] for k, v in pred_dict.items()] # NHWC
92
- )
93
- pred_dict["np"] = F.softmax(pred_dict["np"], dim=-1)[..., 1:]
94
- pred_output = torch.cat(list(pred_dict.values()), -1).cpu().numpy() # NHW3
95
- pred_map = np.squeeze(pred_output) # HW3
96
- pred_inst = __proc_np_hv(pred_map, config.np_thres, config.ksize, config.overall_thres, config.obj_size_thres)
97
- raw_pred_shape = pred_inst.shape[:2]
98
- output_label = pred_inst
99
- return output_label