Function main

saliency_maps/scripts/eval.py:38–72

(args)

    return results

def main(args):
    print("Loading models ...")
    # TODO: Modify these three lines to adapt for your model 
    # For an example on a model finetuned for MSCXR, check demo https://github.com/YingWANGG/M2IB/blob/main/demo.ipynb
    model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device)
    processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
    tokenizer = CLIPTokenizerFast.from_pretrained("openai/clip-vit-base-patch32")
    # TODO: Modify these two lines to adapt for your input data
    # The annotation of Conceptual Captions is in TSV format without a header
    # The first column is the caption and the second is the image url
    df = pd.read_csv(args.data_path,sep='\t')
    data = list(df.itertuples(index=False))
    all_results = []
    print("Evaluating ...")
    for text, image_path in tqdm(sample(data, args.samples)):
        # Load (from a url or a local path) and preprocess image
        try:
            image = Image.open(requests.get(image_path, stream=True, timeout=5).raw) if 'http' in image_path else Image.open(image_path).convert('RGB')
        except:
            print(f"Unable to load image at {image_path}", flush=True)
            continue
        image_feat = processor(images=image, return_tensors="pt")['pixel_values'].to(device) # 3*224*224
        # Tokenize text
        text_ids = torch.tensor([tokenizer.encode(text, add_special_tokens=True)]).to(device)
        # Train information bottleneck on image and text
        vmap = vision_heatmap_iba(text_ids, image_feat, model, args.vlayer, args.vbeta, args.vvar, progbar=False)
        tmap = text_heatmap_iba(text_ids, image_feat, model, args.tlayer, args.tbeta, args.tvar, progbar=False)
        # Evaluation
        results = get_metrics(image_feat, vmap, text_ids, tmap, model)
        results['image'] = image_path
        results['text'] = text
        all_results.append(results)
    all_results = pd.DataFrame(all_results)
    print(all_results.mean(axis=0), flush=True)
    all_results.to_csv(args.output_path)

if __name__ == '__main__':
    parser = argparse.ArgumentParser('M2IB argument parser')