File size: 4,756 Bytes
bdc0687 8a229fe bdc0687 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 |
import json
import torch
import numpy as np
def load_annotations(annotation_file):
with open(annotation_file, 'r') as f:
data = json.load(f)
# Create mapping from image_id to category_id
image_to_category = {}
for annotation in data['annotations']:
image_to_category[annotation['image_id']] = annotation['category_id']
# Get image order and their corresponding labels
images = []
labels = []
for image in data['images']:
image_id = image['id']
file_name = image['file_name']
if image_id in image_to_category:
images.append(file_name)
labels.append(image_to_category[image_id])
return images, labels
def load_model_predictions(model_files):
models = []
all_predictions = {}
class_names = None
for model_file in model_files:
with open(model_file, 'r') as f:
data = json.load(f)
model_name = data['model']
models.append(model_name)
if class_names is None:
class_names = data['class_names']
# Store predictions for each image
for image_name, predictions in data['results'].items():
if image_name not in all_predictions:
all_predictions[image_name] = {}
all_predictions[image_name][model_name] = predictions
return models, all_predictions, class_names
def create_tensors():
# Load annotations
images, labels = load_annotations('iwildcam_demo_annotations.json')
# Load model predictions
model_files = [
'zeroshot_results_facebook_PE_Core_L14_336.json',
'zeroshot_results_google_siglip2_so400m_patch16_naflex.json',
'zeroshot_results_openai_clip_vit_large_patch14.json',
'zeroshot_results_imageomics_bioclip_v2.json',
'zeroshot_results_laion_CLIP_ViT_L_14_laion2B_s32B_b82K.json',
]
models, all_predictions, class_names = load_model_predictions(model_files)
# Create class to index mapping (0-indexed)
class_to_idx = {class_name: idx for idx, class_name in enumerate(class_names)}
# Filter images that have predictions from all models
valid_images = []
valid_labels = []
for i, image_name in enumerate(images):
if image_name in all_predictions and len(all_predictions[image_name]) == len(models):
valid_images.append(image_name)
valid_labels.append(labels[i])
print(f"Found {len(valid_images)} images with predictions from all {len(models)} models")
# Create prediction tensor: H x N x C
H = len(models) # number of models
N = len(valid_images) # number of images
C = len(class_names) # number of classes
prediction_tensor = torch.zeros(H, N, C)
for h, model_name in enumerate(models):
for n, image_name in enumerate(valid_images):
predictions = all_predictions[image_name][model_name]
for class_name, score in predictions.items():
c = class_to_idx[class_name]
prediction_tensor[h, n, c] = score
# Convert category_ids to 0-indexed labels based on class_names order
# Need to map from species names back to category IDs
from collections import OrderedDict
SPECIES_MAP = OrderedDict([
(24, "Jaguar"),
(10, "Ocelot"),
(6, "Mountain Lion"),
(101, "Common Eland"),
(102, "Waterbuck")
])
reverse_species_map = {v: k for k, v in SPECIES_MAP.items()}
# Create category_to_label mapping based on class_names order (not sorted category IDs)
category_order = [reverse_species_map[name] for name in class_names]
category_to_label = {cat: idx for idx, cat in enumerate(category_order)}
label_tensor = torch.tensor([category_to_label[cat] for cat in valid_labels])
# Save tensors
torch.save(prediction_tensor, 'iwildcam_demo.pt')
torch.save(label_tensor, 'iwildcam_demo_labels.pt')
# Save text files
with open('models.txt', 'w') as f:
for model in models:
f.write(f"{model}\n")
with open('images.txt', 'w') as f:
for image in valid_images:
f.write(f"{image}\n")
with open('classes.txt', 'w') as f:
for class_name in class_names:
f.write(f"{class_name}\n")
print(f"Saved prediction tensor of shape {prediction_tensor.shape} to iwildcam_demo.pt")
print(f"Saved label tensor of shape {label_tensor.shape} to iwildcam_demo_labels.pt")
print(f"Saved {len(models)} models to models.txt")
print(f"Saved {len(valid_images)} images to images.txt")
print(f"Saved {len(class_names)} classes to classes.txt")
return prediction_tensor, label_tensor, models, valid_images, class_names
if __name__ == "__main__":
create_tensors()
|