| 18 | from src.python.neuralforge.models.resnet import ResNet18 |
| 19 | |
| 20 | class PredictionThread(QThread): |
| 21 | finished = pyqtSignal(list, list, str) |
| 22 | error = pyqtSignal(str) |
| 23 | |
| 24 | def __init__(self, model, image_path, classes, device): |
| 25 | super().__init__() |
| 26 | self.model = model |
| 27 | self.image_path = image_path |
| 28 | self.classes = classes |
| 29 | self.device = device |
| 30 | |
| 31 | def run(self): |
| 32 | try: |
| 33 | image = Image.open(self.image_path).convert('RGB') |
| 34 | |
| 35 | transform = transforms.Compose([ |
| 36 | transforms.Resize(256), |
| 37 | transforms.CenterCrop(224), |
| 38 | transforms.ToTensor(), |
| 39 | transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) |
| 40 | ]) |
| 41 | |
| 42 | image_tensor = transform(image).unsqueeze(0).to(self.device) |
| 43 | |
| 44 | with torch.no_grad(): |
| 45 | outputs = self.model(image_tensor) |
| 46 | probabilities = F.softmax(outputs, dim=1) |
| 47 | |
| 48 | top5_prob, top5_idx = torch.topk(probabilities, min(5, len(self.classes)), dim=1) |
| 49 | |
| 50 | predictions = [] |
| 51 | confidences = [] |
| 52 | |
| 53 | for idx, prob in zip(top5_idx[0].cpu().numpy(), top5_prob[0].cpu().numpy()): |
| 54 | predictions.append(self.classes[idx]) |
| 55 | confidences.append(float(prob) * 100) |
| 56 | |
| 57 | main_prediction = predictions[0] |
| 58 | |
| 59 | self.finished.emit(predictions, confidences, main_prediction) |
| 60 | |
| 61 | except Exception as e: |
| 62 | self.error.emit(str(e)) |
| 63 | |
| 64 | class NeuralForgeGUI(QMainWindow): |
| 65 | def __init__(self): |