/
image.py
39 lines (30 loc) · 1.36 KB
/
image.py
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
import random
import os
from PIL import Image,ImageFilter,ImageDraw
import numpy as np
import h5py
from PIL import ImageStat
import cv2
import torchvision, torch
from torchvision import datasets, transforms
def load_data(img_path,train = True):
# gt_path = img_path.replace('.jpg','.h5').replace('images','ground_truth') # TRANCOS V3
gt_path = img_path.replace('.jpg','.h5').replace('images','ground_truth_h5') # SHANGHAI
img = Image.open(img_path).convert('RGB')
gt_file = h5py.File(gt_path, 'r')
target = np.asarray(gt_file['density'])
# if train: # ???
# crop_size = (img.size[0]/2,img.size[1]/2)
# if random.randint(0,9)<= -1: # ??
# dx = int(random.randint(0,1)*img.size[0]*1./2)
# dy = int(random.randint(0,1)*img.size[1]*1./2)
# else:
# dx = int(random.random()*img.size[0]*1./2)
# dy = int(random.random()*img.size[1]*1./2)
# img = img.crop((dx,dy,crop_size[0]+dx,crop_size[1]+dy))
# target = target[dy:crop_size[1]+dy,dx:crop_size[0]+dx]
# if random.random()>0.8:
# target = np.fliplr(target)
# img = img.transpose(Image.FLIP_LEFT_RIGHT)
target = cv2.resize(target,(int(target.shape[1]/8),int(target.shape[0]/8)),interpolation = cv2.INTER_CUBIC)*64
return img,target