nick-localhost/Sign-language-detection

1from data import DETRData
2from model import DETR
3import torch
4from torch import load
5from torch.utils.data import DataLoader 
6from matplotlib import pyplot as plt 
7from utils.boxes import rescale_bboxes
8from utils.setup import get_classes
9from utils.logger import get_logger
10from utils.rich_handlers import TestHandler, DetectionHandler
11
12# Initialize loggers and handlers
13logger = get_logger("test")
14test_handler = TestHandler()
15detection_handler = DetectionHandler()
16
17logger.print_banner()
18
19num_classes = 3
20test_dataset = DETRData('data/test', train=False) 
21test_dataloader = DataLoader(test_dataset, shuffle=True, batch_size=4, drop_last=True) 
22model = DETR(num_classes=num_classes)
23model.eval()
24model.load_pretrained('pretrained/4426_model.pt')
25
26X, y = next(iter(test_dataloader))
27
28logger.test("Running inference on test batch...")
29
30import time
31start_time = time.time()
32result = model(X) 
33inference_time = (time.time() - start_time) * 1000  # Convert to ms
34
35probabilities = result['pred_logits'].softmax(-1)[:,:,:-1] 
36max_probs, max_classes = probabilities.max(-1)
37keep_mask = max_probs > 0.95
38batch_indices, query_indices = torch.where(keep_mask) 
39
40bboxes = rescale_bboxes(result['pred_boxes'][batch_indices, query_indices,:], (224,224))
41classes = max_classes[batch_indices, query_indices]
42probas = max_probs[batch_indices, query_indices]
43
44# Log inference timing
45detection_handler.log_inference_time(inference_time)
46
47# Prepare detection results for logging
48detections = []
49for i in range(len(classes)):
50    detections.append({
51        'class': get_classes()[classes[i].item()],
52        'confidence': probas[i].item(),
53        'bbox': bboxes[i].detach().numpy().tolist()
54    })
55
56# Log detection results
57detection_handler.log_detections(detections) 
58
59CLASSES = get_classes()
60
61fig, ax = plt.subplots(2,2) 
62axs = ax.flatten()
63for idx, (img, ax) in enumerate(zip(X, axs)): 
64    ax.imshow(img.permute(1,2,0))
65    for batch_idx, box_class, box_prob, bbox in zip(batch_indices, classes, probas, bboxes): 
66        if batch_idx == idx: 
67            xmin, ymin, xmax, ymax = bbox.detach().numpy()
68            print(xmin, ymin, xmax, ymax) 
69            ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin, fill=False, color=(0.000, 0.447, 0.741), linewidth=3))
70            text = f'{CLASSES[box_class]}: {box_prob:0.2f}'
71            ax.text(xmin, ymin, text, fontsize=15, bbox=dict(facecolor='yellow', alpha=0.5))
72
73fig.tight_layout() 
74plt.show()