| import torch | |
| import torch.nn as nn | |
| import torchvision.models as models | |
| def load_model(model_path): | |
| # Architektur aufbauen | |
| model = models.resnet50(pretrained=False) | |
| model.fc = nn.Linear(2048, 228) | |
| # State Dict laden | |
| state_dict = torch.load(model_path, map_location=torch.device("cpu")) | |
| # Keys ggf. anpassen | |
| new_state_dict = {} | |
| for k, v in state_dict.items(): | |
| if k.startswith("predictor."): | |
| new_k = k.replace("predictor.", "") | |
| else: | |
| new_k = k | |
| new_state_dict[new_k] = v | |
| model.load_state_dict(new_state_dict) | |
| model.eval() | |
| return model | |
| def predict_attributes(model, input_tensor): | |
| with torch.no_grad(): | |
| output = model(input_tensor) | |
| prediction = torch.sigmoid(output).squeeze().numpy() | |
| threshold = 0.5 | |
| predicted_indices = [i for i, p in enumerate(prediction) if p > threshold] | |
| return predicted_indices | |