Spaces:
Sleeping
Sleeping
File size: 5,896 Bytes
bcc0393 1e646da bcc0393 d020327 bcc0393 c81bbd5 a3f3a1f bcc0393 21f6d56 08323cf bcc0393 a3f3a1f f5fe239 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 |
from torch.utils.data import DataLoader
from torchvision import transforms
import numpy as np
import pandas as pd
import os
import cv2
from sklearn.utils import shuffle
from sklearn.model_selection import train_test_split
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms
import torch
import torch.nn as nn
class HybridCNNViT(nn.Module):
def __init__(self, in_channels: int, num_classes: int):
super(HybridCNNViT, self).__init__()
self.conv1 = nn.Conv2d(
in_channels, 64, kernel_size=7, stride=2, padding=3, bias=False)
self.bn1 = nn.BatchNorm2d(64)
self.relu = nn.ReLU(inplace=True)
self.conv2 = nn.Conv2d(64, 128, kernel_size=3,
stride=1, padding=1, bias=False)
self.bn2 = nn.BatchNorm2d(128)
self.conv3 = nn.Conv2d(128, 128, kernel_size=3,
stride=1, padding=1, bias=False)
self.bn3 = nn.BatchNorm2d(128)
self.conv4 = nn.Conv2d(128, 256, kernel_size=3,
stride=2, padding=1, bias=False)
self.bn4 = nn.BatchNorm2d(256)
self.conv5 = nn.Conv2d(256, 256, kernel_size=3,
stride=1, padding=1, bias=False)
self.bn5 = nn.BatchNorm2d(256)
self.conv6 = nn.Conv2d(256, 512, kernel_size=3,
stride=1, padding=1, bias=False)
self.bn6 = nn.BatchNorm2d(512)
self.conv7 = nn.Conv2d(512, 512, kernel_size=3,
stride=2, padding=1, bias=False)
self.bn7 = nn.BatchNorm2d(512)
# Optional MaxPooling (can be removed if strictly no max pooling)
self.maxpool = nn.MaxPool2d(kernel_size=2, stride=2)
self.classifier_conv = nn.Conv2d(
512, num_classes, kernel_size=1, stride=1, padding=0, bias=False)
self.classifier = nn.Sequential(
nn.AdaptiveAvgPool2d((1, 1)),
nn.Flatten(),
nn.Dropout(0.5)
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = self.relu(self.bn1(self.conv1(x)))
x = self.relu(self.bn2(self.conv2(x)))
x = self.relu(self.bn3(self.conv3(x)))
x = self.relu(self.bn4(self.conv4(x)))
x = self.relu(self.bn5(self.conv5(x)))
x = self.relu(self.bn6(self.conv6(x)))
x = self.relu(self.bn7(self.conv7(x)))
x = self.maxpool(x) # Comment this line if no max pooling is needed
x = self.classifier_conv(x)
x = self.classifier(x)
return x
def load_and_pad_single_image(image_path, img_size=(224, 224)):
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
img = cv2.imread(image_path)
if img is None:
raise ValueError(f"Could not read image: {image_path}")
img = cv2.resize(img, img_size)
return np.array(img)
def check_file(image_path):
# image_path = "d/Control-Axial/C-A (2).png"
# Load and preprocess the single image
image = load_and_pad_single_image(image_path)
image = np.expand_dims(image, axis=0) # Convert to batch format
# Duplicate the image 10 times
data = np.repeat(image, 10, axis=0)
# Normalize and transform the image
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[
0.229, 0.224, 0.225])
])
data = torch.tensor(data, dtype=torch.float32).permute(
0, 3, 1, 2).to(device)
# Placeholder labels for 10 images
labels = torch.tensor([0] * 10, dtype=torch.long).to(device)
data, labels = shuffle(data, labels, random_state=42)
train_data, test_data, train_labels, test_labels = train_test_split(
data, labels, test_size=0.2, random_state=42
)
train_labels = torch.tensor(train_labels, dtype=torch.long)
test_labels = torch.tensor(test_labels, dtype=torch.long)
batch_size = 1 # Since we are working with a single image
train_dataset = list(zip(train_data, train_labels))
test_dataset = list(zip(test_data, test_labels))
test_loader = DataLoader(
test_dataset, batch_size=batch_size, shuffle=False)
# Simple test with a model
output = ""
def test_model(model, test_loader, device):
global output
model.to(device)
model.eval()
with torch.no_grad():
for images, labels in test_loader:
images, labels = images.to(device), labels.to(device)
outputs = model(images)
_, predicted = torch.max(outputs.data, 1)
output = predicted
# Convert logits to probabilities
probabilities = F.softmax(outputs, dim=1)
# Get confidence score and prediction
confidence, d = torch.max(probabilities, 1)
print(confidence)
return predicted, confidence
def remove_module_from_checkpoint(checkpoint):
new_state_dict = {}
for key, value in checkpoint["model_state_dict"].items():
new_key = key.replace("module.", "")
new_state_dict[new_key] = value
checkpoint["model_state_dict"] = new_state_dict
return checkpoint
model = HybridCNNViT(3, 2)
checkpoint = torch.load(
"/home/user/app/checkpoint32.pth", weights_only=False, map_location=torch.device('cpu'))
checkpoint = remove_module_from_checkpoint(checkpoint)
model.load_state_dict(checkpoint['model_state_dict'])
model.eval()
model.to(device)
model = nn.DataParallel(model)
output, confidence = test_model(model, test_loader, device)
return "No ms detected" if output.item() == 0 else "MS Detected", confidence.item()
|