File size: 2,460 Bytes
032e687 |
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 |
import os
import json
import numpy as np
import copy
from PIL import Image
from lmdeploy.vl.constants import IMAGE_TOKEN
from pycocotools import mask as mask_utils
import torch.nn.functional as F
import torch
import cv2
from projects.mllm_labeling.image_blending_fn import contour_rendering
def get_masks_from_anns(annotations):
captions = []
masks = []
for annotation in annotations:
rle = annotation['object_anno']['segmentation']
mask = mask_utils.decode(rle)
masks.append(mask)
captions.append(annotation['caption'])
masks = np.stack(masks, axis=0)
return masks, captions
image_folder = './1215_demos/mask_outs/out/'
save_dir = './1215_demos/overall_demos/'
image_files = []
for file_name in os.listdir(image_folder):
if 'out' not in file_name and '.jpg' in file_name:
image_files.append(file_name)
annotation_folder = '1215_demos/mllm_object_cap/'
anno_files = os.listdir(annotation_folder)
annotations = []
for anno_file in anno_files:
with open(os.path.join(annotation_folder, anno_file), 'r') as f:
annotations += json.load(f)
image2anno_dict = {}
for annotation in annotations:
image_id = annotation['image_id']
if image_id not in image2anno_dict.keys():
image2anno_dict[image_id] = [annotation]
else:
image2anno_dict[image_id].append(annotation)
for i, image_name in enumerate(image2anno_dict.keys()):
txt_strs = ''
print('====================================================')
image_path = os.path.join(image_folder, image_name)
image = Image.open(image_path).convert('RGB')
image_annotations = image2anno_dict[image_name]
masks, captions = get_masks_from_anns(image_annotations)
anno_ids = []
for i_obj, caption in enumerate(captions):
txt_strs += '\n\n' + f"{i_obj}: " + caption + '\n\n'
anno_ids.append(i_obj)
print(caption)
print('+++++++++++++++++++++++++++')
image_shape = image.size
masks = torch.Tensor(masks).unsqueeze(0)
masks = F.interpolate(
masks,
size=(image_shape[1], image_shape[0]),
mode='nearest').squeeze(0)
masks = masks.numpy().astype(np.uint8)
image = np.array(image)
contour_rendering(image, masks, mask_ids=anno_ids)
image = Image.fromarray(image)
image.save(os.path.join(save_dir, f"{i}.png"))
with open(os.path.join(save_dir, f"{i}.txt"), 'w') as f:
f.write(txt_strs)
|