import torch
|
import torch.nn as nn
|
import torchvision
|
from torch.nn import functional as F
|
from torch.utils.data import Dataset
|
from src.utils.dataset import DatasetBase
|
import os
|
import cv2
|
import numpy as np
|
|
|
class MGCN(nn.Module):
|
def __init__(self, depth: int = 50,
|
num_classes: int = 1000,
|
mode: str = 'train',):
|
super(MGCN, self).__init__()
|
self.depth = depth
|
self.mode = mode
|
|
if depth == 18:
|
self.model = torchvision.models.resnet18(pretrained=True)
|
elif depth == 34:
|
self.model = torchvision.models.resnet34(pretrained=True)
|
elif depth == 50:
|
self.model = torchvision.models.resnet50(pretrained=True)
|
elif depth == 101:
|
self.model = torchvision.models.resnet101(pretrained=True)
|
elif depth == 152:
|
self.model = torchvision.models.resnet152(pretrained=True)
|
else:
|
raise ValueError("Unsupported MGCN depth: {}, choose from [18, 34, 50, 101, 152]".format(depth))
|
|
self.model.fc = nn.Linear(self.model.fc.in_features, num_classes) # 修改最后一层为1000类输出
|
|
def forward(self, x, device):
|
image, annot = x['batch_inputs'], x['ground_truth']
|
labels = []
|
for ann in annot:
|
assert len(ann['pLabel']) == 1, "Each image should have exactly one label."
|
labels.append(ann['pLabel'][0])
|
labels = torch.tensor(labels, dtype=torch.long, device=device)
|
output = self.model(image.to(device))
|
if self.training and self.mode == 'training':
|
# 训练模式下返回损失
|
loss = F.cross_entropy(output, labels)
|
return loss
|
else:
|
# 推理模式下返回预测结果
|
probs = F.softmax(output, dim=1)
|
preds = torch.argmax(probs, dim=1)
|
|
return {
|
"pred_classes": preds.cpu().detach().numpy(),
|
"pred_scores": probs.max(dim=1).values.cpu().detach().numpy(),
|
"ground_truth": annot
|
}
|
|
class MGCNMetricor:
|
def __init__(self, mode='train'):
|
assert mode in ["validation", "evaluation"]
|
self.mode = mode
|
|
def __call__(self, outputs):
|
pred_classes = []
|
pred_scores = []
|
ground_truth = []
|
gt_labels = []
|
imageIds = []
|
|
for output in outputs:
|
if 'ground_truth' in output:
|
ground_truth.extend(output['ground_truth'])
|
gt_labels.extend([ann['pLabel'][0] for ann in output['ground_truth']])
|
imageIds.extend([ann['id'] for ann in output['ground_truth']])
|
if 'pred_classes' in output:
|
pred_classes.append(output['pred_classes'])
|
if 'pred_scores' in output:
|
pred_scores.append(output['pred_scores'])
|
|
pred_classes = np.concatenate(pred_classes) if pred_classes else np.array([])
|
pred_scores = np.concatenate(pred_scores) if pred_scores else np.array([])
|
gt_labels = np.array(gt_labels)
|
|
# 计算准确率
|
correct = (pred_classes == gt_labels).sum().item()
|
total = len(gt_labels)
|
accuracy = correct / total if total > 0 else 0.0
|
|
metrics = {"accuracy": accuracy}
|
|
if self.mode == 'validation':
|
return metrics, {}
|
|
predictions = {}
|
for img_id, gt, pred, score in zip(imageIds, gt_labels, pred_classes, pred_scores):
|
pred = pred.item()
|
score = score.item()
|
predictions[img_id] = {
|
"id": img_id,
|
"state": "correct" if gt == pred else "false", # ["correct", "false", "leaky", "missing"]
|
"predictions": [
|
{
|
"type": "classify",
|
"categoryId": pred,
|
"score": score
|
}
|
]
|
}
|
|
return metrics, predictions
|
|
def compare(self, metrics, best_metrics):
|
"""比较两个输出的准确率"""
|
acc1 = metrics.get("accuracy", 0.0)
|
acc2 = best_metrics.get("accuracy", 0.0) if best_metrics else 0.0
|
if acc1 >= acc2:
|
return True
|
return False
|
|
|
class MGCNDataset(DatasetBase):
|
def __init__(self, data, split_id=None, phase='training',
|
input_size=(224, 224), crop_size=(224, 224),
|
transforms=None):
|
# Convert single phase to list if needed
|
phases = [phase] if isinstance(phase, str) else phase
|
super(MGCNDataset, self).__init__(data, split_id, phases, transforms)
|
self.input_size = input_size # (height, width)
|
self.crop_size = crop_size # (height, width)
|
self.std = [0.229, 0.224, 0.225] # ImageNet标准化
|
self.mean = [0.485, 0.456, 0.406]
|
|
def __len__(self):
|
return len(self.images)
|
|
def __getitem__(self, idx):
|
image_data = self.images[idx]
|
image_path = image_data.get('path')
|
if not image_path or not os.path.exists(image_path):
|
raise FileNotFoundError(f"Image path {image_path} does not exist.")
|
|
# 读取图像并进行预处理
|
image = cv2.imread(image_path, cv2.IMREAD_UNCHANGED)
|
image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
|
image = cv2.resize(image, (self.input_size[1], self.input_size[0]))
|
image = np.transpose(image, (2, 0, 1)) # 转换为C x H x W格式
|
image = torch.from_numpy(image)
|
|
# 标准化图像
|
image = image.float() / 255.0
|
image = (image - torch.tensor(self.mean).view(3, 1, 1)) / torch.tensor(self.std).view(3, 1, 1)
|
|
# 处理标注格式，注意test中包含无标注图像
|
annot = {
|
'id': image_data['id'],
|
'name': image_data['name'],
|
'path': image_path,
|
'height': image_data['height'],
|
'width': image_data['width'],
|
}
|
|
annotations = image_data.get('annotations', [])
|
annot.update(self.format_annotations(annotations))
|
return {
|
"batch_inputs": image, # shape: (C,H,W)
|
"ground_truth": annot
|
}
|
