Spaces:
Running
Running
| import PIL | |
| from PIL import Image, ImageDraw ,ImageFont | |
| from matplotlib import pyplot as plt | |
| import torchvision.transforms as T | |
| import os | |
| import torch | |
| import yaml | |
| def show_torch_img(img): | |
| img = to_np_image(img) | |
| plt.imshow(img) | |
| plt.axis("off") | |
| def to_np_image(all_images): | |
| all_images = (all_images.permute(0, 2, 3, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8).cpu().numpy()[0] | |
| return all_images | |
| def tensor_to_pil(tensor_imgs): | |
| if type(tensor_imgs) == list: | |
| tensor_imgs = torch.cat(tensor_imgs) | |
| tensor_imgs = (tensor_imgs / 2 + 0.5).clamp(0, 1) | |
| to_pil = T.ToPILImage() | |
| pil_imgs = [to_pil(img) for img in tensor_imgs] | |
| return pil_imgs | |
| def pil_to_tensor(pil_imgs): | |
| to_torch = T.ToTensor() | |
| if type(pil_imgs) == PIL.Image.Image: | |
| tensor_imgs = to_torch(pil_imgs).unsqueeze(0)*2-1 | |
| elif type(pil_imgs) == list: | |
| tensor_imgs = torch.cat([to_torch(pil_imgs).unsqueeze(0)*2-1 for img in pil_imgs]).to(device) | |
| else: | |
| raise Exception("Input need to be PIL.Image or list of PIL.Image") | |
| return tensor_imgs | |
| ## TODO implement this | |
| # n = 10 | |
| # num_rows = 4 | |
| # num_col = n // num_rows | |
| # num_col = num_col + 1 if n % num_rows else num_col | |
| # num_col | |
| def add_margin(pil_img, top = 0, right = 0, bottom = 0, | |
| left = 0, color = (255,255,255)): | |
| width, height = pil_img.size | |
| new_width = width + right + left | |
| new_height = height + top + bottom | |
| result = Image.new(pil_img.mode, (new_width, new_height), color) | |
| result.paste(pil_img, (left, top)) | |
| return result | |
| def image_grid(imgs, rows = 1, cols = None, | |
| size = None, | |
| titles = None, text_pos = (0, 0)): | |
| if type(imgs) == list and type(imgs[0]) == torch.Tensor: | |
| imgs = torch.cat(imgs) | |
| if type(imgs) == torch.Tensor: | |
| imgs = tensor_to_pil(imgs) | |
| if not size is None: | |
| imgs = [img.resize((size,size)) for img in imgs] | |
| if cols is None: | |
| cols = len(imgs) | |
| assert len(imgs) >= rows*cols | |
| top=20 | |
| w, h = imgs[0].size | |
| delta = 0 | |
| if len(imgs)> 1 and not imgs[1].size[1] == h: | |
| delta = top | |
| h = imgs[1].size[1] | |
| if not titles is None: | |
| font = ImageFont.truetype("/usr/share/fonts/truetype/freefont/FreeMono.ttf", | |
| size = 20, encoding="unic") | |
| h = top + h | |
| grid = Image.new('RGB', size=(cols*w, rows*h+delta)) | |
| for i, img in enumerate(imgs): | |
| if not titles is None: | |
| img = add_margin(img, top = top, bottom = 0,left=0) | |
| draw = ImageDraw.Draw(img) | |
| draw.text(text_pos, titles[i],(0,0,0), | |
| font = font) | |
| if not delta == 0 and i > 0: | |
| grid.paste(img, box=(i%cols*w, i//cols*h+delta)) | |
| else: | |
| grid.paste(img, box=(i%cols*w, i//cols*h)) | |
| return grid | |
| """ | |
| input_folder - dataset folder | |
| """ | |
| def load_dataset(input_folder): | |
| # full_file_names = glob.glob(input_folder) | |
| # class_names = [x[0] for x in os.walk(input_folder)] | |
| class_names = next(os.walk(input_folder))[1] | |
| class_names[:] = [d for d in class_names if not d[0] == '.'] | |
| file_names=[] | |
| for class_name in class_names: | |
| cur_path = os.path.join(input_folder, class_name) | |
| filenames = next(os.walk(cur_path), (None, None, []))[2] | |
| filenames = [f for f in filenames if not f[0] == '.'] | |
| file_names.append(filenames) | |
| return class_names, file_names | |
| def dataset_from_yaml(yaml_location): | |
| with open(yaml_location, 'r') as stream: | |
| data_loaded = yaml.safe_load(stream) | |
| return data_loaded |