| import os |
| import cv2 |
| import json |
| import numpy as np |
| import base64 |
| import torch |
| import torchvision |
| import asyncio |
| import time |
| from fastapi import FastAPI, WebSocket, WebSocketDisconnect |
| from fastapi.responses import HTMLResponse |
| from ultralytics import YOLO |
| from torchvision.models.detection.faster_rcnn import FastRCNNPredictor |
| from torchvision.transforms import functional as F |
|
|
| app = FastAPI() |
|
|
| |
| DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
| models = {"yolo": None, "rcnn": None} |
| CLASSES = ["Background"] + [chr(i) for i in range(65, 91)] |
|
|
| |
| yolo_path = "models/best.pt" |
| if os.path.exists(yolo_path): |
| try: |
| models["yolo"] = YOLO(yolo_path) |
| print("✅ YOLO model loaded.") |
| except Exception as e: |
| print(f"❌ Failed to load YOLO: {e}") |
|
|
| |
| rcnn_path = "models/best_faster_rcnn.pth" |
| if os.path.exists(rcnn_path): |
| try: |
| NUM_CLASSES = 27 |
| rcnn = torchvision.models.detection.fasterrcnn_resnet50_fpn(weights=None) |
| in_features = rcnn.roi_heads.box_predictor.cls_score.in_features |
| rcnn.roi_heads.box_predictor = FastRCNNPredictor(in_features, NUM_CLASSES) |
| rcnn.load_state_dict(torch.load(rcnn_path, map_location=DEVICE)) |
| rcnn.to(DEVICE) |
| rcnn.eval() |
| models["rcnn"] = rcnn |
| print("✅ Faster R-CNN model loaded.") |
| except Exception as e: |
| print(f"❌ Failed to load R-CNN: {e}") |
|
|
| @app.get("/") |
| async def get(): |
| with open("index.html", "r") as f: |
| return HTMLResponse(f.read()) |
|
|
| @app.websocket("/ws") |
| async def websocket_endpoint(websocket: WebSocket): |
| await websocket.accept() |
| |
| |
| state = { |
| "last_pred": None, |
| "pred_count": 0, |
| "frames_to_confirm": 3 |
| } |
| |
| image_queue = asyncio.Queue(maxsize=1) |
|
|
| async def receiver(): |
| """Aggressively receives frames and keeps only the newest one.""" |
| try: |
| while True: |
| data = await websocket.receive_text() |
| |
| if image_queue.full(): |
| try: |
| image_queue.get_nowait() |
| except asyncio.QueueEmpty: |
| pass |
| await image_queue.put(data) |
| except Exception: |
| pass |
|
|
| async def processor(): |
| """Pulls from queue as fast as the model can handle.""" |
| try: |
| while True: |
| data = await image_queue.get() |
| payload = json.loads(data) |
| header, encoded = payload["image"].split(",", 1) |
| active_model = payload["model"] |
| |
| img_bytes = base64.b64decode(encoded) |
| np_arr = np.frombuffer(img_bytes, np.uint8) |
| img = cv2.imdecode(np_arr, cv2.IMREAD_COLOR) |
| |
| detected_letter = "" |
| best_label_idx = None |
| best_label_name = "" |
| |
| current_model = models.get(active_model) |
| if current_model is None: |
| |
| fallback = "rcnn" if active_model == "yolo" else "yolo" |
| current_model = models.get(fallback) |
| active_model = fallback if current_model else None |
|
|
| if active_model == "yolo" and current_model: |
| results = current_model(img, conf=0.6, verbose=False)[0] |
| highest_conf = 0.0 |
| for box in results.boxes: |
| conf = box.conf.item() |
| cls_idx = int(box.cls.item()) |
| name = current_model.names[cls_idx] |
| if conf > highest_conf: |
| highest_conf, best_label_idx, best_label_name = conf, cls_idx, name |
| x1, y1, x2, y2 = map(int, box.xyxy[0]) |
| cv2.rectangle(img, (x1, y1), (x2, y2), (255, 0, 0), 2) |
| cv2.putText(img, f"YOLO: {name} {conf:.2f}", (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 0), 2) |
|
|
| elif active_model == "rcnn" and current_model: |
| rgb_img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) |
| img_tensor = F.to_tensor(rgb_img).unsqueeze(0).to(DEVICE) |
| with torch.no_grad(): |
| predictions = current_model(img_tensor)[0] |
| |
| highest_conf = 0.0 |
| for i in range(len(predictions['boxes'])): |
| score = predictions['scores'][i].item() |
| if score > 0.75: |
| box = predictions['boxes'][i].cpu().numpy().astype(int) |
| label_idx = predictions['labels'][i].item() |
| name = CLASSES[label_idx] |
| if score > highest_conf: |
| highest_conf, best_label_idx, best_label_name = score, label_idx, name |
| cv2.rectangle(img, (box[0], box[1]), (box[2], box[3]), (0, 255, 0), 2) |
| cv2.putText(img, f"RCNN: {name} {score:.2f}", (box[0], box[1]-10), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2) |
|
|
| |
| if best_label_idx is not None: |
| if best_label_idx == state["last_pred"]: |
| state["pred_count"] += 1 |
| else: |
| state["last_pred"] = best_label_idx |
| state["pred_count"] = 1 |
| if state["pred_count"] == state["frames_to_confirm"]: |
| detected_letter = best_label_name |
| else: |
| state["last_pred"], state["pred_count"] = None, 0 |
|
|
| |
| _, buffer = cv2.imencode('.jpg', img, [cv2.IMWRITE_JPEG_QUALITY, 70]) |
| out_encoded = base64.b64encode(buffer).decode('utf-8') |
| await websocket.send_json({ |
| "image": f"data:image/jpeg;base64,{out_encoded}", |
| "detected": detected_letter |
| }) |
| except Exception: |
| pass |
|
|
| |
| await asyncio.gather(receiver(), processor()) |
