Spaces:
Sleeping
Sleeping
| from __future__ import annotations | |
| import argparse | |
| from pathlib import Path | |
| import numpy as np | |
| import torch | |
| from torch import nn | |
| from torch.utils.data import DataLoader | |
| from torchvision import models, transforms | |
| import ocular_preliminary_benchmark as base | |
| def parse_args() -> argparse.Namespace: | |
| parser = argparse.ArgumentParser( | |
| description="Extract ImageNet EfficientNetB3 embeddings in benchmark sample order." | |
| ) | |
| parser.add_argument("--data-root", type=Path, required=True) | |
| parser.add_argument("--output-dir", type=Path, default=Path("runs/phase1/features")) | |
| parser.add_argument("--input-size", type=int, default=384) | |
| parser.add_argument("--batch-size", type=int, default=32) | |
| parser.add_argument("--num-workers", type=int, default=0) | |
| parser.add_argument("--device", default="auto") | |
| parser.add_argument("--force", action="store_true") | |
| return parser.parse_args() | |
| def build_model(device: torch.device) -> nn.Module: | |
| weights = models.EfficientNet_B3_Weights.DEFAULT | |
| model = models.efficientnet_b3(weights=weights) | |
| model.classifier = nn.Identity() | |
| model.eval() | |
| return model.to(device) | |
| def main() -> None: | |
| args = parse_args() | |
| args.output_dir.mkdir(parents=True, exist_ok=True) | |
| feature_path = args.output_dir / f"efficientnet_b3_features_{args.input_size}.npy" | |
| label_path = args.output_dir / f"efficientnet_b3_labels_{args.input_size}.npy" | |
| if feature_path.exists() and label_path.exists() and not args.force: | |
| print(f"Feature cache already exists: {feature_path}") | |
| return | |
| device = base.get_device(args.device) | |
| samples = base.discover_samples(args.data_root.expanduser().resolve(), None) | |
| labels = np.asarray([sample.label for sample in samples], dtype=np.int64) | |
| transform = transforms.Compose( | |
| [ | |
| transforms.Resize(int(args.input_size * 1.14)), | |
| transforms.CenterCrop(args.input_size), | |
| transforms.ToTensor(), | |
| transforms.Normalize(base.IMAGENET_MEAN, base.IMAGENET_STD), | |
| ] | |
| ) | |
| dataset = base.FundusDataset(samples, transform) | |
| loader = DataLoader( | |
| dataset, | |
| batch_size=args.batch_size, | |
| shuffle=False, | |
| num_workers=args.num_workers, | |
| pin_memory=device.type == "cuda", | |
| persistent_workers=args.num_workers > 0, | |
| ) | |
| model = build_model(device) | |
| features: list[np.ndarray] = [] | |
| with torch.no_grad(): | |
| for batch_index, (images, _) in enumerate(loader, start=1): | |
| images = images.to(device, non_blocking=device.type == "cuda") | |
| embedding = model(images) | |
| if device.type == "mps" and hasattr(torch, "mps"): | |
| torch.mps.synchronize() | |
| features.append(embedding.detach().cpu().numpy().astype(np.float32)) | |
| if batch_index <= 3 or batch_index % 25 == 0: | |
| print( | |
| f"extracted {min(batch_index * args.batch_size, len(samples))}/{len(samples)}", | |
| flush=True, | |
| ) | |
| matrix = np.concatenate(features, axis=0) | |
| np.save(feature_path, matrix) | |
| np.save(label_path, labels) | |
| print(f"saved {feature_path} shape={matrix.shape}") | |
| print(f"saved {label_path} shape={labels.shape}") | |
| if __name__ == "__main__": | |
| main() | |