Views
No views yet
1import torch
2from datasets import load_dataset
3from transformers import SegGptImageProcessor, SegGptForImageSegmentation
4
5model_id = "BAAI/seggpt-vit-large"
6image_processor = SegGptImageProcessor.from_pretrained(checkpoint)
7model = SegGptForImageSegmentation.from_pretrained(checkpoint)
8
9dataset_id = "EduardoPacheco/FoodSeg103"
10ds = load_dataset(dataset_id, split="train")
11# Number of labels in FoodSeg103 (not including background)
12num_labels = 103
13
14image_input = ds[4]["image"]
15ground_truth = ds[4]["label"]
16image_prompt = ds[29]["image"]
17mask_prompt = ds[29]["label"]
18
19inputs = image_processor(
20 images=image_input,
21 prompt_images=image_prompt,
22 prompt_masks=mask_prompt,
23 num_labels=num_labels,
24 return_tensors="pt"
25)
26
27with torch.no_grad():
28 outputs = model(**inputs)
29
30target_sizes = [image_input.size[::-1]]
31mask = image_processor.post_process_semantic_segmentation(outputs, target_sizes, num_labels=num_labels)[0]1@misc{wang2023seggpt,
2 title={SegGPT: Segmenting Everything In Context},
3 author={Xinlong Wang and Xiaosong Zhang and Yue Cao and Wen Wang and Chunhua Shen and Tiejun Huang},
4 year={2023},
5 eprint={2304.03284},
6 archivePrefix={arXiv},
7 primaryClass={cs.CV}
8}