|
|
import json |
|
|
import os |
|
|
from PIL import Image |
|
|
import numpy as np |
|
|
from pycocotools.mask import encode, decode, frPyObjects |
|
|
from tqdm import tqdm |
|
|
import copy |
|
|
|
|
|
if __name__ == '__main__': |
|
|
|
|
|
new_img_id = 0 |
|
|
nerf_dataset = [] |
|
|
first_frame_annotation_relpath = "teacup/Annotations/DSCF0672.png" |
|
|
first_frame_img_relpath = "teacup/JPEGImages/DSCF0672.JPG" |
|
|
first_frame_annotation_path = "/home/yuqian_fu/Projects/sam2/teacup/Annotations/DSCF0672.png" |
|
|
first_frame_annotation_img = Image.open(first_frame_annotation_path) |
|
|
first_frame_annotation = np.array(first_frame_annotation_img) |
|
|
height, width = first_frame_annotation.shape |
|
|
unique_instances = np.unique(first_frame_annotation) |
|
|
unique_instances = unique_instances[unique_instances != 0] |
|
|
coco_format_annotations = [] |
|
|
|
|
|
for instance_value in unique_instances: |
|
|
binary_mask = (first_frame_annotation == instance_value).astype(np.uint8) |
|
|
segmentation = encode(np.asfortranarray(binary_mask)) |
|
|
segmentation = { |
|
|
'counts': segmentation['counts'].decode('ascii'), |
|
|
'size': segmentation['size'], |
|
|
} |
|
|
area = binary_mask.sum().astype(float) |
|
|
coco_format_annotations.append( |
|
|
{ |
|
|
'segmentation': segmentation, |
|
|
'area': area, |
|
|
'category_id': instance_value.astype(float), |
|
|
} |
|
|
) |
|
|
|
|
|
|
|
|
sample_img_relpath = "teacup/JPEGImages/DSCF0732.JPG" |
|
|
image_info = { |
|
|
'file_name': sample_img_relpath, |
|
|
'height': height, |
|
|
'width': width, |
|
|
} |
|
|
|
|
|
sample_annotation_path = "/home/yuqian_fu/Projects/sam2/teacup/Annotations/DSCF0732.png" |
|
|
sample_annotation = np.array(Image.open(sample_annotation_path)) |
|
|
sample_unique_instances = np.unique(sample_annotation) |
|
|
sample_unique_instances = sample_unique_instances[sample_unique_instances != 0] |
|
|
anns = [] |
|
|
for instance_value in sample_unique_instances: |
|
|
assert instance_value in unique_instances, 'Found new target not in the first frame' |
|
|
binary_mask = (sample_annotation == instance_value).astype(np.uint8) |
|
|
segmentation = encode(np.asfortranarray(binary_mask)) |
|
|
segmentation = { |
|
|
'counts': segmentation['counts'].decode('ascii'), |
|
|
'size': segmentation['size'], |
|
|
} |
|
|
area = binary_mask.sum().astype(float) |
|
|
anns.append( |
|
|
{ |
|
|
'segmentation': segmentation, |
|
|
'area': area, |
|
|
'category_id': instance_value.astype(float), |
|
|
} |
|
|
) |
|
|
first_frame_anns = copy.deepcopy(coco_format_annotations) |
|
|
if len(anns) < len(first_frame_anns): |
|
|
first_frame_anns = [ann for ann in first_frame_anns if ann['category_id'] in sample_unique_instances] |
|
|
assert len(anns) == len(first_frame_anns) |
|
|
sample = { |
|
|
'image': sample_img_relpath, |
|
|
'image_info': image_info, |
|
|
'anns': anns, |
|
|
'first_frame_image': first_frame_img_relpath, |
|
|
'first_frame_anns': first_frame_anns, |
|
|
'new_img_id': new_img_id, |
|
|
'video_name': "teacup", |
|
|
} |
|
|
nerf_dataset.append(sample) |
|
|
|
|
|
|
|
|
save_path = "/home/yuqian_fu/Projects/sam2/predicted_mask/nerf6.json" |
|
|
with open(save_path, 'w') as f: |
|
|
json.dump(nerf_dataset, f) |
|
|
print(f'Save at {save_path}. Total sample: {len(nerf_dataset)}') |
