| 41 | self.device = device |
| 42 | |
| 43 | def run(self): |
| 44 | try: |
| 45 | image = Image.open(self.image_path).convert('RGB') |
| 46 | |
| 47 | transform = transforms.Compose([ |
| 48 | transforms.Resize(256), |
| 49 | transforms.CenterCrop(224), |
| 50 | transforms.ToTensor(), |
| 51 | transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) |
| 52 | ]) |
| 53 | |
| 54 | image_tensor = transform(image).unsqueeze(0).to(self.device) |
| 55 | |
| 56 | with torch.no_grad(): |
| 57 | outputs = self.model(image_tensor) |
| 58 | probabilities = F.softmax(outputs, dim=1) |
| 59 | |
| 60 | top5_prob, top5_idx = torch.topk(probabilities, min(5, len(self.classes)), dim=1) |
| 61 | |
| 62 | predictions = [] |
| 63 | confidences = [] |
| 64 | |
| 65 | for idx, prob in zip(top5_idx[0].cpu().numpy(), top5_prob[0].cpu().numpy()): |
| 66 | predictions.append(self.classes[idx]) |
| 67 | confidences.append(float(prob) * 100) |
| 68 | |
| 69 | main_prediction = predictions[0] |
| 70 | |
| 71 | self.finished.emit(predictions, confidences, main_prediction) |
| 72 | |
| 73 | except Exception as e: |
| 74 | self.error.emit(str(e)) |
| 75 | |
| 76 | class NeuralForgeGUI(QMainWindow): |
| 77 | def __init__(self): |