Spaces:
Build error
Build error
| import io | |
| import os | |
| import zipfile | |
| import gdown | |
| import h5py | |
| import numpy as np | |
| import requests | |
| from matplotlib.pyplot import imread, imsave | |
| from scipy.io import loadmat | |
| from scipy.ndimage import gaussian_filter | |
| def download_salicon(data_path): | |
| """Downloads the SALICON dataset. Three folders are then created that | |
| contain the stimuli, binary fixation maps, and blurred saliency | |
| distributions respectively. | |
| Args: | |
| data_path (str): Defines the path where the dataset will be | |
| downloaded and extracted to. | |
| .. seealso:: The code for downloading files from google drive is based | |
| on the solution provided at [https://bit.ly/2JSVgMQ]. | |
| """ | |
| print(">> Downloading SALICON dataset...", end="", flush=True) | |
| default_path = data_path + "salicon/" | |
| fixations_path = default_path + "fixations/" | |
| saliency_path = default_path + "saliency/" | |
| os.makedirs(fixations_path, exist_ok=True) | |
| os.makedirs(saliency_path, exist_ok=True) | |
| ids = ["1g8j-hTT-51IG1UFwP0xTGhLdgIUCW5e5", | |
| "1P-jeZXCsjoKO79OhFUgnj6FGcyvmLDPj", | |
| "1PnO7szbdub1559LfjYHMy65EDC4VhJC8"] | |
| urls = ["https://drive.google.com/uc?id=" + | |
| i + "&export=download" for i in ids] | |
| save_paths = [default_path, fixations_path, saliency_path] | |
| for count, url in enumerate(urls): | |
| gdown.download(url, data_path + "tmp.zip", quiet=True) | |
| with zipfile.ZipFile(data_path + "tmp.zip", "r") as zip_ref: | |
| for file in zip_ref.namelist(): | |
| if "test" not in file: | |
| zip_ref.extract(file, save_paths[count]) | |
| os.rename(default_path + "images", default_path + "stimuli") | |
| os.remove(data_path + "tmp.zip") | |
| print("done!", flush=True) | |
| def download_mit1003(data_path): | |
| """Downloads the MIT1003 dataset. Three folders are then created that | |
| contain the stimuli, binary fixation maps, and blurred saliency | |
| distributions respectively. | |
| Args: | |
| data_path (str): Defines the path where the dataset will be | |
| downloaded and extracted to. | |
| """ | |
| print(">> Downloading MIT1003 dataset...", end="", flush=True) | |
| default_path = data_path + "mit1003/" | |
| stimuli_path = default_path + "stimuli/" | |
| fixations_path = default_path + "fixations/" | |
| saliency_path = default_path + "saliency/" | |
| os.makedirs(stimuli_path, exist_ok=True) | |
| os.makedirs(fixations_path, exist_ok=True) | |
| os.makedirs(saliency_path, exist_ok=True) | |
| url = "https://people.csail.mit.edu/tjudd/WherePeopleLook/ALLSTIMULI.zip" | |
| with open(data_path + "tmp.zip", "wb") as f: | |
| f.write(requests.get(url).content) | |
| with zipfile.ZipFile(data_path + "tmp.zip", "r") as zip_ref: | |
| for file in zip_ref.namelist(): | |
| if file.endswith(".jpeg"): | |
| file_name = os.path.split(file)[1] | |
| file_path = stimuli_path + file_name | |
| with open(file_path, "wb") as stimulus: | |
| stimulus.write(zip_ref.read(file)) | |
| url = "https://people.csail.mit.edu/tjudd/WherePeopleLook/ALLFIXATIONMAPS.zip" | |
| with open(data_path + "tmp.zip", "wb") as f: | |
| f.write(requests.get(url).content) | |
| with zipfile.ZipFile(data_path + "tmp.zip", "r") as zip_ref: | |
| for file in zip_ref.namelist(): | |
| file_name = os.path.split(file)[1] | |
| if file.endswith("Pts.jpg"): | |
| file_path = fixations_path + file_name | |
| # this file is mistakenly included in the dataset and can be ignored | |
| if file_name == "i05june05_static_street_boston_p1010764fixPts.jpg": | |
| continue | |
| with open(file_path, "wb") as fixations: | |
| fixations.write(zip_ref.read(file)) | |
| elif file.endswith("Map.jpg"): | |
| file_path = saliency_path + file_name | |
| with open(file_path, "wb") as saliency: | |
| saliency.write(zip_ref.read(file)) | |
| os.remove(data_path + "tmp.zip") | |
| print("done!", flush=True) | |
| def download_cat2000(data_path): | |
| """Downloads the CAT2000 dataset. Three folders are then created that | |
| contain the stimuli, binary fixation maps, and blurred saliency | |
| distributions respectively. | |
| Args: | |
| data_path (str): Defines the path where the dataset will be | |
| downloaded and extracted to. | |
| """ | |
| print(">> Downloading CAT2000 dataset...", end="", flush=True) | |
| default_path = data_path + "cat2000/" | |
| os.makedirs(data_path, exist_ok=True) | |
| url = "http://saliency.mit.edu/trainSet.zip" | |
| with open(data_path + "tmp.zip", "wb") as f: | |
| f.write(requests.get(url).content) | |
| with zipfile.ZipFile(data_path + "tmp.zip", "r") as zip_ref: | |
| for file in zip_ref.namelist(): | |
| if not("Output" in file or "allFixData" in file): | |
| zip_ref.extract(file, data_path) | |
| os.rename(data_path + "trainSet/", default_path) | |
| os.rename(default_path + "Stimuli", default_path + "stimuli") | |
| os.rename(default_path + "FIXATIONLOCS", default_path + "fixations") | |
| os.rename(default_path + "FIXATIONMAPS", default_path + "saliency") | |
| os.remove(data_path + "tmp.zip") | |
| print("done!", flush=True) | |
| def download_dutomron(data_path): | |
| """Downloads the DUT-OMRON dataset. Three folders are then created that | |
| contain the stimuli, binary fixation maps, and blurred saliency | |
| distributions respectively. | |
| Args: | |
| data_path (str): Defines the path where the dataset will be | |
| downloaded and extracted to. | |
| """ | |
| print(">> Downloading DUTOMRON dataset...", end="", flush=True) | |
| default_path = data_path + "dutomron/" | |
| stimuli_path = default_path + "stimuli/" | |
| fixations_path = default_path + "fixations/" | |
| saliency_path = default_path + "saliency/" | |
| os.makedirs(stimuli_path, exist_ok=True) | |
| os.makedirs(fixations_path, exist_ok=True) | |
| os.makedirs(saliency_path, exist_ok=True) | |
| url = "http://saliencydetection.net/dut-omron/download/DUT-OMRON-image.zip" | |
| with open(data_path + "tmp.zip", "wb") as f: | |
| f.write(requests.get(url).content) | |
| with zipfile.ZipFile(data_path + "tmp.zip", "r") as zip_ref: | |
| for file in zip_ref.namelist(): | |
| if file.endswith(".jpg") and "._" not in file: | |
| file_name = os.path.basename(file) | |
| file_path = stimuli_path + file_name | |
| with open(file_path, "wb") as stimulus: | |
| stimulus.write(zip_ref.read(file)) | |
| url = "http://saliencydetection.net/dut-omron/download/DUT-OMRON-eye-fixations.zip" | |
| with open(data_path + "tmp.zip", "wb") as f: | |
| f.write(requests.get(url).content) | |
| with zipfile.ZipFile(data_path + "tmp.zip", "r") as zip_ref: | |
| for file in zip_ref.namelist(): | |
| if file.endswith(".mat") and "._" not in file: | |
| file_name = os.path.basename(file) | |
| file_name = os.path.splitext(file_name)[0] + ".png" | |
| loaded_zip = io.BytesIO(zip_ref.read(file)) | |
| fixations = loadmat(loaded_zip)["s"] | |
| sorted_idx = fixations[:, 2].argsort() | |
| fixations = fixations[sorted_idx] | |
| size = fixations[0, :2] | |
| fixations_map = np.zeros((size[1], size[0])) | |
| fixations_map[fixations[1:, 1], | |
| fixations[1:, 0]] = 1 | |
| saliency_map = gaussian_filter(fixations_map, 16) | |
| imsave(saliency_path + file_name, saliency_map, cmap="gray") | |
| imsave(fixations_path + file_name, fixations_map, cmap="gray") | |
| os.remove(data_path + "tmp.zip") | |
| print("done!", flush=True) | |
| def download_pascals(data_path): | |
| """Downloads the PASCAL-S dataset. Three folders are then created that | |
| contain the stimuli, binary fixation maps, and blurred saliency | |
| distributions respectively. | |
| Args: | |
| data_path (str): Defines the path where the dataset will be | |
| downloaded and extracted to. | |
| """ | |
| print(">> Downloading PASCALS dataset...", end="", flush=True) | |
| default_path = data_path + "pascals/" | |
| stimuli_path = default_path + "stimuli/" | |
| fixations_path = default_path + "fixations/" | |
| saliency_path = default_path + "saliency/" | |
| os.makedirs(stimuli_path, exist_ok=True) | |
| os.makedirs(fixations_path, exist_ok=True) | |
| os.makedirs(saliency_path, exist_ok=True) | |
| url = "http://cbs.ic.gatech.edu/salobj/download/salObj.zip" | |
| with open(data_path + "tmp.zip", "wb") as f: | |
| f.write(requests.get(url).content) | |
| with zipfile.ZipFile(data_path + "tmp.zip", "r") as zip_ref: | |
| for file in zip_ref.namelist(): | |
| file_name = os.path.basename(file) | |
| if file.endswith(".jpg") and "imgs/pascal" in file: | |
| file_path = stimuli_path + file_name | |
| with open(file_path, "wb") as stimulus: | |
| stimulus.write(zip_ref.read(file)) | |
| elif file.endswith(".png") and "pascal/humanFix" in file: | |
| file_path = saliency_path + file_name | |
| with open(file_path, "wb") as saliency: | |
| saliency.write(zip_ref.read(file)) | |
| elif "pascalFix.mat" in file: | |
| loaded_zip = io.BytesIO(zip_ref.read(file)) | |
| with h5py.File(loaded_zip, "r") as f: | |
| fixations = np.array(f.get("fixCell"))[0] | |
| fixations_list = [] | |
| for reference in fixations: | |
| obj = np.array(f[reference]) | |
| obj = np.stack((obj[0], obj[1]), axis=-1) | |
| fixations_list.append(obj) | |
| elif "pascalSize.mat" in file: | |
| loaded_zip = io.BytesIO(zip_ref.read(file)) | |
| with h5py.File(loaded_zip, "r") as f: | |
| sizes = np.array(f.get("sizeData")) | |
| sizes = np.transpose(sizes, (1, 0)) | |
| for idx, value in enumerate(fixations_list): | |
| size = [int(x) for x in sizes[idx]] | |
| fixations_map = np.zeros(size) | |
| for fixation in value: | |
| fixations_map[int(fixation[0]) - 1, | |
| int(fixation[1]) - 1] = 1 | |
| file_name = str(idx + 1) + ".png" | |
| file_path = fixations_path + file_name | |
| imsave(file_path, fixations_map, cmap="gray") | |
| os.remove(data_path + "tmp.zip") | |
| print("done!", flush=True) | |
| def download_osie(data_path): | |
| """Downloads the OSIE dataset. Three folders are then created that | |
| contain the stimuli, binary fixation maps, and blurred saliency | |
| distributions respectively. | |
| Args: | |
| data_path (str): Defines the path where the dataset will be | |
| downloaded and extracted to. | |
| """ | |
| print(">> Downloading OSIE dataset...", end="", flush=True) | |
| default_path = data_path + "osie/" | |
| stimuli_path = default_path + "stimuli/" | |
| fixations_path = default_path + "fixations/" | |
| saliency_path = default_path + "saliency/" | |
| os.makedirs(stimuli_path, exist_ok=True) | |
| os.makedirs(fixations_path, exist_ok=True) | |
| os.makedirs(saliency_path, exist_ok=True) | |
| url = "https://github.com/NUS-VIP/predicting-human-gaze-beyond-pixels/archive/master.zip" | |
| with open(data_path + "tmp.zip", "wb") as f: | |
| f.write(requests.get(url).content) | |
| with zipfile.ZipFile(data_path + "tmp.zip", "r") as zip_ref: | |
| for file in zip_ref.namelist(): | |
| file_name = os.path.basename(file) | |
| if file.endswith(".jpg") and "data/stimuli" in file: | |
| file_path = stimuli_path + file_name | |
| with open(file_path, "wb") as stimulus: | |
| stimulus.write(zip_ref.read(file)) | |
| elif file_name == "fixations.mat": | |
| loaded_zip = io.BytesIO(zip_ref.read(file)) | |
| loaded_mat = loadmat(loaded_zip)["fixations"] | |
| for idx, value in enumerate(loaded_mat): | |
| subjects = value[0][0][0][1] | |
| fixations_map = np.zeros((600, 800)) | |
| for subject in subjects: | |
| x_vals = subject[0][0][0][0][0] | |
| y_vals = subject[0][0][0][1][0] | |
| fixations = np.stack((y_vals, x_vals), axis=-1) | |
| fixations = fixations.astype(int) | |
| fixations_map[fixations[:, 0], | |
| fixations[:, 1]] = 1 | |
| file_name = str(1001 + idx) + ".png" | |
| saliency_map = gaussian_filter(fixations_map, 16) | |
| imsave(saliency_path + file_name, saliency_map, cmap="gray") | |
| imsave(fixations_path + file_name, fixations_map, cmap="gray") | |
| os.remove(data_path + "tmp.zip") | |
| print("done!", flush=True) | |
| def download_fiwi(data_path): | |
| """Downloads the FIWI dataset. Three folders are then created that | |
| contain the stimuli, binary fixation maps, and blurred saliency | |
| distributions respectively. | |
| Args: | |
| data_path (str): Defines the path where the dataset will be | |
| downloaded and extracted to. | |
| """ | |
| print(">> Downloading FIWI dataset...", end="", flush=True) | |
| default_path = data_path + "fiwi/" | |
| stimuli_path = default_path + "stimuli/" | |
| fixations_path = default_path + "fixations/" | |
| saliency_path = default_path + "saliency/" | |
| os.makedirs(stimuli_path, exist_ok=True) | |
| os.makedirs(fixations_path, exist_ok=True) | |
| os.makedirs(saliency_path, exist_ok=True) | |
| url = "https://www.dropbox.com/s/30nxg2uwd1wpb80/webpage_dataset.zip?dl=1" | |
| with open(data_path + "tmp.zip", "wb") as f: | |
| f.write(requests.get(url).content) | |
| with zipfile.ZipFile(data_path + "tmp.zip", "r") as zip_ref: | |
| for file in zip_ref.namelist(): | |
| file_name = os.path.basename(file) | |
| if file.endswith(".png") and "stimuli" in file: | |
| file_path = stimuli_path + file_name | |
| with open(file_path, "wb") as stimulus: | |
| stimulus.write(zip_ref.read(file)) | |
| elif file.endswith(".png") and "all5" in file: | |
| loaded_zip = io.BytesIO(zip_ref.read(file)) | |
| fixations = imread(loaded_zip) | |
| saliency = gaussian_filter(fixations, 30) | |
| imsave(saliency_path + file_name, saliency, cmap="gray") | |
| imsave(fixations_path + file_name, fixations, cmap="gray") | |
| os.remove(data_path + "tmp.zip") | |
| print("done!", flush=True) | |
| def download_pretrained_weights(data_path, key): | |
| """Downloads the pre-trained weights for the VGG16 model when | |
| training or the MSI-Net when testing on new data instances. | |
| Args: | |
| data_path (str): Defines the path where the weights will be | |
| downloaded and extracted to. | |
| key (str): Describes the type of model for which the weights will | |
| be downloaded. This contains the device and dataset. | |
| .. seealso:: The code for downloading files from google drive is based | |
| on the solution provided at [https://bit.ly/2JSVgMQ]. | |
| """ | |
| print(">> Downloading pre-trained weights...", end="", flush=True) | |
| os.makedirs(data_path, exist_ok=True) | |
| ids = { | |
| "vgg16_hybrid": "1ff0va472Xs1bvidCwRlW3Ctf7Hbyyn7p", | |
| "model_salicon_cpu": "1Xy9C72pcA8DO4CY0rc6B7wsuE9L9DDZY", | |
| "model_salicon_gpu": "1Th7fqVYx25ePMZz4LYsjNQWgAu8tJqwL", | |
| "model_mit1003_cpu": "1jsESjYtsTvkMqKftA4rdstfB7mSYw5Ec", | |
| "model_mit1003_gpu": "1P_tWxBl3igZlzcHGp5H3T3kzsOskWeG6", | |
| "model_cat2000_cpu": "1XxaEx7xxD6rHasQTa-VY7T7eVpGhMxuV", | |
| "model_cat2000_gpu": "1T6ChEGB6Mf02gKXrENjdeD6XXJkE_jHh", | |
| "model_dutomron_cpu": "14tuRZpKi8LMDKRHNVUylu6RuAaXLjHTa", | |
| "model_dutomron_gpu": "15LG_M45fpYC1pTwnwmArNTZw_Z3BOIA-", | |
| "model_pascals_cpu": "1af9IvBqFamKWx64Ror6ALivuKNioOVIf", | |
| "model_pascals_gpu": "1C-T-RQzX2SaiY9Nw1HmaSx6syyCt01Z0", | |
| "model_osie_cpu": "1JD1tvAqZGxj_gEGmIfoxb9dTe5HOaHj1", | |
| "model_osie_gpu": "1g8UPr1hGpUdOSWerRb751pZqiWBOZOCh", | |
| "model_fiwi_cpu": "19qj9nAjd5gVHLB71oRn_YfYDw5n4Uf2X", | |
| "model_fiwi_gpu": "12OpIMIi2IyDVaxkE2d37XO9uUsSYf1Ec" | |
| } | |
| url = "https://drive.google.com/uc?id=" + ids[key] + "&export=download" | |
| gdown.download(url, data_path + "tmp.zip", quiet=True) | |
| with zipfile.ZipFile(data_path + "tmp.zip", "r") as zip_ref: | |
| for file in zip_ref.namelist(): | |
| zip_ref.extract(file, data_path) | |
| os.remove(data_path + "tmp.zip") | |
| print("done!", flush=True) | |