Spaces:
Runtime error
Runtime error
File size: 5,125 Bytes
fb6c2da |
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 135 136 137 138 |
import torch
import json
import cv2
import torch
import os
from basicsr.utils import img2tensor, tensor2img
import random
class dataset_coco():
def __init__(self, path_json, root_path, image_size, mode='train'):
super(dataset_coco, self).__init__()
with open(path_json, 'r', encoding='utf-8') as fp:
data = json.load(fp)
data = data['images']
self.paths = []
self.root_path = root_path
for file in data:
input_path = file['filepath']
if mode == 'train':
if 'val' not in input_path:
self.paths.append(file)
else:
if 'val' in input_path:
self.paths.append(file)
def __getitem__(self, idx):
file = self.paths[idx]
input_path = file['filepath']
input_name = file['filename']
path = os.path.join(self.root_path, input_path, input_name)
im = cv2.imread(path)
im = cv2.resize(im, (512,512))
im = img2tensor(im, bgr2rgb=True, float32=True)/255.
sentences = file['sentences']
sentence = sentences[int(random.random()*len(sentences))]['raw'].strip('.')
return {'im':im, 'sentence':sentence}
def __len__(self):
return len(self.paths)
class dataset_coco_mask():
def __init__(self, path_json, root_path_im, root_path_mask, image_size):
super(dataset_coco_mask, self).__init__()
with open(path_json, 'r', encoding='utf-8') as fp:
data = json.load(fp)
data = data['annotations']
self.files = []
self.root_path_im = root_path_im
self.root_path_mask = root_path_mask
for file in data:
name = "%012d.png"%file['image_id']
self.files.append({'name':name, 'sentence':file['caption']})
def __getitem__(self, idx):
file = self.files[idx]
name = file['name']
# print(os.path.join(self.root_path_im, name))
im = cv2.imread(os.path.join(self.root_path_im, name.replace('.png','.jpg')))
im = cv2.resize(im, (512,512))
im = img2tensor(im, bgr2rgb=True, float32=True)/255.
mask = cv2.imread(os.path.join(self.root_path_mask, name))#[:,:,0]
mask = cv2.resize(mask, (512,512))
mask = img2tensor(mask, bgr2rgb=True, float32=True)[0].unsqueeze(0)#/255.
sentence = file['sentence']
return {'im':im, 'mask':mask, 'sentence':sentence}
def __len__(self):
return len(self.files)
class dataset_coco_mask_color():
def __init__(self, path_json, root_path_im, root_path_mask, image_size):
super(dataset_coco_mask_color, self).__init__()
with open(path_json, 'r', encoding='utf-8') as fp:
data = json.load(fp)
data = data['annotations']
self.files = []
self.root_path_im = root_path_im
self.root_path_mask = root_path_mask
for file in data:
name = "%012d.png"%file['image_id']
self.files.append({'name':name, 'sentence':file['caption']})
def __getitem__(self, idx):
file = self.files[idx]
name = file['name']
# print(os.path.join(self.root_path_im, name))
im = cv2.imread(os.path.join(self.root_path_im, name.replace('.png','.jpg')))
im = cv2.resize(im, (512,512))
im = img2tensor(im, bgr2rgb=True, float32=True)/255.
mask = cv2.imread(os.path.join(self.root_path_mask, name))#[:,:,0]
mask = cv2.resize(mask, (512,512))
mask = img2tensor(mask, bgr2rgb=True, float32=True)/255.#[0].unsqueeze(0)#/255.
sentence = file['sentence']
return {'im':im, 'mask':mask, 'sentence':sentence}
def __len__(self):
return len(self.files)
class dataset_coco_mask_color_sig():
def __init__(self, path_json, root_path_im, root_path_mask, image_size):
super(dataset_coco_mask_color_sig, self).__init__()
with open(path_json, 'r', encoding='utf-8') as fp:
data = json.load(fp)
data = data['annotations']
self.files = []
self.root_path_im = root_path_im
self.root_path_mask = root_path_mask
reg = {}
for file in data:
name = "%012d.png"%file['image_id']
if name in reg:
continue
self.files.append({'name':name, 'sentence':file['caption']})
reg[name] = name
def __getitem__(self, idx):
file = self.files[idx]
name = file['name']
# print(os.path.join(self.root_path_im, name))
im = cv2.imread(os.path.join(self.root_path_im, name.replace('.png','.jpg')))
im = cv2.resize(im, (512,512))
im = img2tensor(im, bgr2rgb=True, float32=True)/255.
mask = cv2.imread(os.path.join(self.root_path_mask, name))#[:,:,0]
mask = cv2.resize(mask, (512,512))
mask = img2tensor(mask, bgr2rgb=True, float32=True)/255.#[0].unsqueeze(0)#/255.
sentence = file['sentence']
return {'im':im, 'mask':mask, 'sentence':sentence, 'name': name}
def __len__(self):
return len(self.files) |