1import torch
2import numpy as np
3from huggingface_hub import hf_hub_download
4from safetensors.torch import load_file
5from ultralytics import YOLO
6from erdes.models.components.cls_model import SwinUnetrClassifier
7from erdes.data.components.utils import resize
8
9# --- 1. Load YOLO for ocular globe detection ---
10yolo = YOLO(hf_hub_download("pcvlab/yolov8_ocular_ultrasound_globe_detection", "yolov8_ocular_ultrasound_globe_detection.pt"))
11
12# --- 2. Crop your POCUS ultrasound video using YOLO (finds largest globe bbox across all frames) ---
13def crop_video(video_path, model, conf=0.8):
14 # First pass: find the largest bounding box across all frames
15 area_max, cropping_bbox = 0, None
16 for frame in model.predict(video_path, stream=True, verbose=False, conf=conf):
17 if len(frame.boxes.xywhn):
18 bbox = frame.boxes.xywhn[0].cpu().numpy()
19 area = bbox[2] * bbox[3]
20 if area > area_max:
21 area_max, cropping_bbox = area, bbox
22
23 if cropping_bbox is None:
24 raise ValueError("YOLO could not detect ocular globe in video.")
25
26 # Second pass: crop every frame with the largest bbox
27 frames = []
28 for frame in model.predict(video_path, stream=True, verbose=False, conf=conf):
29 img = frame.orig_img # [H, W, C] BGR
30 h, w, _ = img.shape
31 x_c, y_c, bw, bh = cropping_bbox
32 x1, y1 = int((x_c - bw/2) * w), int((y_c - bh/2) * h)
33 x2, y2 = int((x_c + bw/2) * w), int((y_c + bh/2) * h)
34 frames.append(img[y1:y2, x1:x2])
35
36 return np.stack(frames) # [D, H, W, C]
37
38frames = crop_video("your_video.mp4", yolo) # [D, H, W, C]
39
40# --- 3. Preprocess ---
41video = torch.from_numpy(frames).float() # [D, H, W, C]
42video = video.permute(3, 0, 1, 2) # [C, D, H, W]
43if video.shape[0] == 3:
44 video = video.mean(dim=0, keepdim=True) # grayscale [1, D, H, W]
45video = resize((96, 128, 128))(video) / 255.0 # pad + resize + normalize
46video = video.unsqueeze(0) # [1, 1, 96, 128, 128]
47
48# --- 4. Load model and run inference ---
49model = SwinUnetrClassifier(in_channels=1, num_classes=1, img_size=[96, 128, 128], feature_size=24, pooling="avg")
50weights = load_file(hf_hub_download("pcvlab/swinunetr_normal_vs_pvd", "model.safetensors"))
51model.load_state_dict(weights)
52model.eval()
53
54with torch.no_grad():
55 logit = model(video)
56 prob = torch.sigmoid(logit).item()
57 pred = int(prob > 0.5)
58
59labels = {'0': 'Normal', '1': 'Posterior Vitreous Detachment'}
60print(f"Prediction: {labels[str(pred)]} (confidence: {prob:.3f})")
