Spaces:
Sleeping
Sleeping
import numpy as np | |
import gradio as gr | |
import glob | |
import imageio | |
from PIL import Image | |
import numpy as np | |
import torch | |
import torch.nn as nn | |
import torch.optim as optim | |
from torchvision.models import resnet18, ResNet18_Weights | |
import albumentations as A | |
import tqdm | |
import random | |
from sklearn.metrics import f1_score | |
import torchvision.models as models | |
import numpy as np | |
import torch | |
import torch.nn as nn | |
import albumentations as A | |
import imageio | |
from PIL import Image | |
import gradio as gr | |
import os | |
import glob | |
import cv2 | |
# 創建模型 | |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
# 設置模型路徑 | |
model_path = r"MRIy14.pt" | |
# 設定設備 | |
device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
class SimpleCNN(nn.Module): | |
def __init__(self, num_classes): | |
super(SimpleCNN, self).__init__() | |
self.conv1 = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, padding=1) | |
self.relu1 = nn.ReLU() | |
self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2) | |
self.conv2 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, padding=1) | |
self.relu2 = nn.ReLU() | |
self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2) | |
self.conv3 = nn.Conv2d(in_channels=32, out_channels=32, kernel_size=3, padding=1) | |
self.relu3 = nn.ReLU() | |
self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2) | |
self.fc = nn.Linear(32 * 28 * 28, num_classes) | |
def forward(self, x): | |
# 保存各層輸出以供熱力圖生成 | |
x1 = self.conv1(x) | |
x1_relu = self.relu1(x1) | |
x1_pool = self.pool1(x1_relu) | |
x2 = self.conv2(x1_pool) | |
x2_relu = self.relu2(x2) | |
x2_pool = self.pool2(x2_relu) | |
x3 = self.conv3(x2_pool) | |
x3_relu = self.relu3(x3) | |
x3_pool = self.pool3(x3_relu) | |
x_flatten = x3_pool.view(x3_pool.size(0), -1) | |
out = self.fc(x_flatten) | |
# 返回最終輸出和用於熱力圖的中間特徵圖 (例如 x3_relu 或 x3_pool) | |
return out, x3_relu # 或者 return out, x3_pool | |
# 創建模型 | |
model = torch.load(model_path, map_location='cpu') | |
model = model.to(device) | |
model.eval() | |
transform = A.Compose([ | |
A.Resize(224, 224), | |
]) | |
''' | |
def predict(image, model=model, transform=transform): | |
#print(f"Input image path: {image}") | |
im = np.array(image) | |
im = Image.fromarray(image).convert('L') | |
im = np.array(im) | |
im = A.Resize(224, 224)(image=im)['image'] | |
im = np.stack([im]*3, axis=0)[None, ...] #1, 3, 224, 224 | |
im = im/np.max(im) | |
im_d = torch.from_numpy(im).to(device).float() | |
with torch.no_grad(): outputs = model(im_d) | |
_, predicted = torch.max(outputs, 1) | |
prediction = predicted.cpu().numpy() | |
transformed_image = transform(image=im)['image'] | |
#transformed_image = np.transpose(transformed_image, (2, 1, 0))[None, ...] | |
transformed_image = np.stack([transformed_image] * 3, axis=0)[None, ...] | |
transformed_image = torch.from_numpy(transformed_image.astype(np.float32) / np.max(transformed_image)).to(device) | |
with torch.no_grad(): | |
output = model(transformed_image.float()) | |
prediction = torch.argmax(output, dim=1).item() | |
print(f"Prediction: {prediction}") | |
print(f"Prediction: {prediction}") | |
if prediction == 1: | |
return "Cancer" | |
else: | |
return "Healthy" | |
demo = gr.Interface(predict, gr.Image(), outputs=["text"]) | |
demo.launch() | |
''' | |
import torch.nn.functional as F | |
def generate_heatmap(model, last_conv_layer, input_tensor, pred_class): | |
""" | |
使用 Grad-CAM 生成更精确的热力图。 | |
""" | |
# 获取最后卷积层的特征图和模型输出 | |
model.eval() | |
features = [] | |
def hook_function(module, input, output): | |
features.append(output) | |
hook = last_conv_layer.register_forward_hook(hook_function) | |
model_output = model(input_tensor) | |
hook.remove() | |
# 获取预测类别的得分 | |
class_score = model_output[:, pred_class] | |
# 计算梯度 | |
model.zero_grad() | |
class_score.backward(retain_graph=True) | |
# 获取梯度和特征图 | |
gradients = features[0].grad # 获取最后卷积层的梯度 | |
pooled_gradients = torch.mean(gradients, dim=[0, 2, 3]) # 全局平均池化 | |
# 获取特征图,并应用梯度加权 | |
features = features[0] | |
for i in range(features.shape[1]): | |
features[:, i, :, :] *= pooled_gradients[i] | |
# 创建热力图 | |
heatmap = torch.mean(features, dim=1).squeeze() | |
heatmap = np.maximum(heatmap.detach().cpu().numpy(), 0) | |
heatmap /= np.max(heatmap) | |
return heatmap | |
def overlay_heatmap(image, heatmap, intensity=0.5, threshold=0.5): | |
""" | |
优化热力图叠加逻辑。 | |
""" | |
heatmap = cv2.resize(heatmap, (image.shape[1], image.shape[0])) | |
heatmap = np.uint8(255 * heatmap) | |
heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET) | |
# 只在高信心区域应用热力图 | |
mask = heatmap > np.max(heatmap) * threshold | |
superimposed_img = image.copy() | |
superimposed_img[mask] = superimposed_img[mask] * (1 - intensity) + heatmap[mask] * intensity | |
superimposed_img = np.clip(superimposed_img, 0, 255).astype(np.uint8) | |
return superimposed_img | |
def predict(image): | |
""" | |
預測並生成熱力圖。 | |
""" | |
# 處理圖片 | |
processed_image = transform(image=image)['image'] | |
processed_image = np.transpose(processed_image, (2, 0, 1)).astype(np.float32) | |
input_tensor = torch.from_numpy(processed_image[None, ...]).to(device) # 更有效的轉換方法 | |
# 預測 | |
with torch.no_grad(): | |
outputs, last_conv_output = model(input_tensor) | |
_, predicted = torch.max(outputs, 1) | |
prediction = predicted.item() | |
# 生成熱力圖 | |
heatmap = generate_heatmap(last_conv_output, prediction) | |
result_image = overlay_heatmap(np.array(image), heatmap) | |
return prediction, result_image | |
# 使用 Gradio | |
demo = gr.Interface(fn=predict, inputs=gr.Image(), outputs=["text", "image"]) | |
demo.launch() |