Bittensorminingfactory/streetvision-roadwork-v2

1import torch2import torch.nn as nn3import torch.nn.functional as F4from PIL import Image5import torchvision.transforms as transforms6import sys7import os8 9# Add fastervit to path if needed10try:11    from fastervit import create_model12except ImportError:13    sys.path.insert(0, '/workspace/sn72-fastervit-training')14    from fastervit import create_model15 16from pathlib import Path17 18class BinaryFasterViT(nn.Module):19    def __init__(self, base_model, method='max'):20        super().__init__()21        self.base_model = base_model22        self.method = method23        24    def forward(self, x):25        logits = self.base_model(x)26        probs = F.softmax(logits, dim=1)27        if self.method == 'max':28            binary_score = probs.max(dim=1)[0]29        elif self.method == 'sum':30            binary_score = probs.sum(dim=1)31        else:32            raise ValueError(f"Unknown method: {self.method}")33        return binary_score.unsqueeze(1)34 35def load_model(checkpoint_path, method='max'):36    base_model = create_model('faster_vit_0_224', pretrained=False, num_classes=4)37    checkpoint = torch.load(checkpoint_path, map_location='cpu')38    if 'model_state_dict' in checkpoint:39        base_model.load_state_dict(checkpoint['model_state_dict'])40    else:41        base_model.load_state_dict(checkpoint)42    model = BinaryFasterViT(base_model, method=method)43    model.eval()44    return model45 46def preprocess_image(image_path):47    transform = transforms.Compose([48        transforms.Resize((224, 224)),49        transforms.ToTensor(),50        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])51    ])52    image = Image.open(image_path).convert('RGB')53    return transform(image).unsqueeze(0)54 55def predict(model, image_path):56    input_tensor = preprocess_image(image_path)57    with torch.no_grad():58        score = model(input_tensor)59    return score.item()60 61def test_on_directory(model, image_dir, threshold=0.5):62    image_paths = list(Path(image_dir).glob('*.jpg')) + list(Path(image_dir).glob('*.png'))63    results = []64    for img_path in image_paths[:200]:65        try:66            score = predict(model, str(img_path))67            results.append({'image': str(img_path), 'score': score, 'prediction': 'damage' if score > threshold else 'no_damage'})68        except Exception as e:69            print(f"Error processing {img_path}: {e}")70    return results71 72if __name__ == '__main__':73    import argparse74    parser = argparse.ArgumentParser()75    parser.add_argument('--checkpoint', default='checkpoints/best_model.pth')76    parser.add_argument('--method', default='max', choices=['max', 'sum'])77    parser.add_argument('--test-dir', default=None)78    parser.add_argument('--image', default=None)79    args = parser.parse_args()80    81    print(f"Loading model from {args.checkpoint} with method={args.method}...")82    model = load_model(args.checkpoint, method=args.method)83    print("Model loaded. Binary wrapper active.")84    85    if args.image:86        score = predict(model, args.image)87        print(f"Image: {args.image}")88        print(f"Binary score: {score:.4f}")89        print(f"Prediction: {'DAMAGE' if score > 0.5 else 'NO DAMAGE'}")90    91    if args.test_dir:92        print(f"Testing on directory: {args.test_dir}")93        results = test_on_directory(model, args.test_dir)94        scores = [r['score'] for r in results]95        if scores:96            print(f"Processed {len(results)} images")97            print(f"Score stats: min={min(scores):.4f}, max={max(scores):.4f}, mean={sum(scores)/len(scores):.4f}")98