nick-localhost/Sign-language-detection

1import cv2
2import torch
3from torch import load
4from model import DETR
5import albumentations as A
6from utils.boxes import rescale_bboxes
7from utils.setup import get_classes, get_colors
8from utils.logger import get_logger
9from utils.rich_handlers import DetectionHandler, create_detection_live_display
10import sys
11import time 
12import os, requests
13
14# Initialize logger and handlers
15logger = get_logger("realtime")
16detection_handler = DetectionHandler()
17
18logger.print_banner()
19logger.realtime("Initializing real-time sign language detection...")
20
21transforms = A.Compose(
22        [   
23            A.Resize(224,224),
24            A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
25            A.ToTensorV2()
26        ]
27    )
28
29model = DETR(num_classes=3)
30model.eval()
31# model.load_pretrained('pretrained/4426_model.pt')
32
33
34
35MODEL_PATH = "pretrained/4426_model.pt"
36MODEL_URL = "https://drive.google.com/uc?export=download&id=YOUR_FILE_ID"
37
38os.makedirs(os.path.dirname(MODEL_PATH), exist_ok=True)
39
40if not os.path.exists(MODEL_PATH):
41    print(f"Downloading model from {MODEL_URL} ...")
42    response = requests.get(MODEL_URL, stream=True)
43    with open(MODEL_PATH, "wb") as f:
44        for chunk in response.iter_content(chunk_size=8192):
45            if chunk:
46                f.write(chunk)
47    print("Model downloaded successfully.")
48
49
50model.load_pretrained(MODEL_PATH)
51
52
53CLASSES = get_classes() 
54COLORS = get_colors() 
55
56logger.realtime("Starting camera capture...")
57cap = cv2.VideoCapture(0)
58
59# Initialize performance tracking
60frame_count = 0
61fps_start_time = time.time()
62
63while cap.isOpened(): 
64    ret, frame = cap.read()
65    if not ret:
66        logger.error("Failed to read frame from camera")
67        break
68        
69    # Time the inference
70    inference_start = time.time()
71    transformed = transforms(image=frame)
72    result = model(torch.unsqueeze(transformed['image'], dim=0))
73    inference_time = (time.time() - inference_start) * 1000  # Convert to ms
74
75    probabilities = result['pred_logits'].softmax(-1)[:,:,:-1] 
76    max_probs, max_classes = probabilities.max(-1)
77    keep_mask = max_probs > 0.8
78
79    batch_indices, query_indices = torch.where(keep_mask) 
80
81    bboxes = rescale_bboxes(result['pred_boxes'][batch_indices, query_indices,:], (1920,1080))
82    classes = max_classes[batch_indices, query_indices]
83    probas = max_probs[batch_indices, query_indices]
84
85    # Prepare detection results for logging
86    detections = []
87    for bclass, bprob, bbox in zip(classes, probas, bboxes): 
88        bclass_idx = bclass.detach().numpy()
89        bprob_val = bprob.detach().numpy() 
90        x1,y1,x2,y2 = bbox.detach().numpy()
91        
92        detections.append({
93            'class': CLASSES[bclass_idx],
94            'confidence': float(bprob_val),
95            'bbox': [float(x1), float(y1), float(x2), float(y2)]
96        })
97        
98        # Draw bounding boxes on frame
99        frame = cv2.rectangle(frame, (int(x1),int(y1)), (int(x2),int(y2)), COLORS[bclass_idx], 2)
100        frame_text = f"{CLASSES[bclass_idx]} - {round(float(bprob_val),4)}"
101        # frame = cv2.rectangle(frame, (int(x1),int(y1)-100), (int(x1)+700,int(y1)), COLORS[bclass_idx], -1)
102        text_size, _ = cv2.getTextSize(frame_text, cv2.FONT_HERSHEY_SIMPLEX, 0.7, 2)
103        text_w, text_h = text_size
104        frame = cv2.rectangle(frame, (int(x1), int(y1)-text_h-5), (int(x1)+text_w, int(y1)), COLORS[bclass_idx], -1)
105
106        # frame = cv2.putText(frame, frame_text, (int(x1),int(y1)), cv2.FONT_HERSHEY_DUPLEX, 2, (255,255,255), 4, cv2.LINE_AA)
107        frame = cv2.putText(frame, frame_text, (int(x1), int(y1)-5), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255,255,255), 2, cv2.LINE_AA)
108
109
110    # Calculate FPS
111    frame_count += 1
112    if frame_count % 30 == 0:  # Log every 30 frames
113        elapsed_time = time.time() - fps_start_time
114        fps = 30 / elapsed_time
115        
116        # Log detection results and performance
117        if detections:
118            detection_handler.log_detections(detections, frame_id=frame_count)
119        detection_handler.log_inference_time(inference_time, fps)
120        
121        # Reset FPS counter
122        fps_start_time = time.time()
123
124    frame_resized = cv2.resize(frame, (1280, 720))
125    cv2.imshow('Frame', frame_resized)
126
127    if cv2.waitKey(1) & 0xFF == ord('q'): 
128        logger.realtime("Stopping real-time detection...")
129        break
130
131cap.release() 
132cv2.destroyAllWindows() 
133