Upload era_data.py

era_data.py

@@ -0,0 +1,311 @@
+import pandas as pd
+from glob import glob
+from torch.utils.data import Dataset
+import os
+from PIL import Image
+import numpy as np
+import cv2
+
+
+def get_IDS(IMG_DIR='output/images_preprocessed', era=False, CATALOGUE_FN='output/cdli_catalogue_data.csv'):
+    img_fns = glob(os.path.join(IMG_DIR, '*.png'))
+    IDS = [os.path.basename(fn).rstrip('.png') for fn in img_fns]
+    if era:
+        IDS = list(set(IDS) & set(pd.read_csv(
+            CATALOGUE_FN, usecols=['id_text', 'era'], dtype={'id_text': object}
+        ).dropna(subset=['era']).set_index('id_text').to_dict()['era'].keys()))
+    return IDS
+
+def pad_zeros(x):
+    x_new = str(x)
+    return (6-len(x_new))*'0'+x_new
+
+class TabletEraDataset(Dataset):
+    
+    ERA_INDICES = {
+        'early_bronze': 0,
+        'mid_late_bronze': 1,
+        'iron': 2
+    }
+    
+    def __init__(self, CATALOGUE_FN='output/cdli_catalogue_data.csv', IMG_DIR='output/images_preprocessed', IDS=None):
+        self.id2era = pd.read_csv(
+            CATALOGUE_FN, usecols=['id_text', 'era'], dtype={'id_text': object}
+        ).dropna(subset=['era']).set_index('id_text').to_dict()['era']
+        
+        self.img_fns = glob(os.path.join(IMG_DIR, '*.png'))
+        self.IDS = [os.path.basename(fn).rstrip('.png') for fn in self.img_fns]
+        
+        if IDS is not None:
+            print(f'Filtering {len(self.IDS)} IDS down to provided {len(IDS)}...')
+            IDS_set = set(IDS)
+            indices = [i for i, ID in enumerate(self.IDS) if ID in IDS_set]
+            self.img_fns = [self.img_fns[i] for i in indices]
+            self.IDS = [self.IDS[i] for i in indices]
+        
+    def __len__(self):
+        return len(self.IDS)
+        
+    def __getitem__(self, idx):
+        fn = self.img_fns[idx]
+        ID = self.IDS[idx]
+        era = self.id2era[ID]
+        img = np.asarray(Image.open(fn))
+        return img.astype(np.float32) / 255, self.ERA_INDICES[era]
+
+
+class TabletPeriodDataset(Dataset):
+    
+    # based on (normed) periods with at least 100 photos:
+    PERIOD_INDICES = {
+        
+        'other': 0,
+        'Ur III': 1,
+        'Neo-Assyrian': 2,
+        'Old Babylonian': 3,
+        'Middle Babylonian': 4,
+        'Neo-Babylonian': 5,
+        'Old Akkadian': 6,
+        'Achaemenid': 7,
+        'Early Old Babylonian': 8,
+        'ED IIIb': 9,
+        'Middle Assyrian': 10,
+        'Old Assyrian': 11,
+        'Uruk III': 12,
+        'Proto-Elamite': 13,
+        'Lagash II': 14,
+        'Ebla': 15,
+        'ED IIIa': 16,
+        'Hellenistic': 17,
+        'ED I-II': 18,
+        'Middle Elamite': 19,
+        'Hittite': 20,
+        'Uruk IV': 21
+    }
+
+    PROVENIENCE_INDICES = {
+        'Nineveh': 1,
+         'Nippur': 2,
+         'unknown': 3,
+         'Umma': 4,
+         'Puzris-Dagan': 5,
+         'Girsu': 6,
+         'Ur': 7,
+         'Uruk': 8,
+         'Kanesh': 9,
+         'Assur': 10,
+         'Adab': 11,
+         'Garsana': 12,
+         'Gasur/Nuzi': 13,
+         'Susa': 14,
+         'Sippar-Yahrurum': 15,
+         'Larsa': 16,
+         'Nerebtum': 17,
+         'mod. Babylonia': 18,
+         'Parsa': 19,
+         'Kish': 20,
+         'Kalhu': 21,
+         'Tuttul': 22,
+         'Suruppak': 23,
+         'Babili': 24,
+         'Ebla': 25,
+         'mod. Beydar': 26,
+         'Akhetaten': 27,
+         'Esnunna': 28,
+         'Borsippa': 29,
+         'Kar-Tukulti-Ninurta': 30,
+         'mod. Jemdet Nasr': 31,
+         'mod. northern Babylonia': 32,
+         'Alalakh': 33,
+         'Hattusa': 34,
+         'Isin': 35,
+         'Elbonia': 36,
+         'Sibaniba': 37,
+         'Tutub': 38,
+         'Pi-Kasi': 39,
+         'Irisagrig': 40,
+         'Ansan': 41,
+         'Dilbat': 42,
+         'Zabalam': 43,
+         'mod. Mugdan/ Umm al-Jir': 44,
+         'Marad': 45,
+         'Eridu': 46,
+         'Seleucia': 47,
+         'mod. Abu Halawa': 48,
+         'Dur-Untas': 49,
+         'Nagar': 50,
+         'Lagaba': 51,
+         'Asnakkum': 52,
+         'Dur-Kurigalzu': 53,
+         'mod. Tell Sabaa': 54,
+         'mod. Abu Jawan': 55,
+         'mod. Tell Fakhariyah': 56,
+         'Dur-Abi-esuh': 57,
+         'Ugarit': 58,
+         'mod. Diqdiqqah': 59,
+         'Tarbisu': 60,
+         'Lagash': 61,
+         'Kisurra': 62,
+         'Elammu': 63,
+         'Du-Enlila': 64,
+         'Kutha': 65,
+         'mod. Umm el-Hafriyat': 66,
+         'Dur-Sarrukin': 67,
+         'Bad-Tibira': 68,
+         'Bit-zerija': 69,
+         'Kilizu': 70,
+         'mod. Pasargadae': 71,
+         'Abdju': 72,
+         'Surmes': 73,
+         'mod. Qatibat': 74,
+         'Tigunanum': 75,
+         'mod. Tell al-Lahm': 76,
+         'mod. Mesopotamia': 77,
+         'Subat-Enlil': 78,
+         'mod. Konar Sandal': 79,
+         'Gissi': 80,
+         'Agamatanu': 81,
+         'Aqa': 82,
+         'Kapri-sa-naqidati': 83,
+         'Esura': 84,
+         'Nahalla': 85,
+         'Bit-Sahtu': 86,
+         'mod. Sepphoris': 87,
+         'Dusabar': 88,
+         'mod. Tell Sifr': 89,
+         'Nasir': 90,
+         'Kumu': 91,
+         'Kazallu': 92,
+         'Kapru': 93,
+         'Hurruba': 94,
+         'mod. Deh-e-no, Iran': 95,
+         "mod. Za'aleh": 96,
+         'mod. Tepe Farukhabad': 97,
+         'Hursagkalama': 98,
+         'Carchemish': 99,
+         'mod. Ben Shemen, Israel': 100,
+         'Kutalla': 101,
+         'Der': 102,
+         'Imgur-Enlil': 103,
+         'mod. Hillah': 104,
+         'mod. Uhudu': 105,
+         'mod. Mahmudiyah': 106,
+         'Terqa': 107,
+         'Arrapha': 108,
+         'mod. Tell en-Nasbeh': 109,
+         'mod. Kalah Shergat': 110,
+         'Kar-Nabu': 111,
+         'Harran': 112,
+         'mod. Til-Buri': 113,
+         'Shuruppak': 114,
+         'mod. Abu Salabikh': 115,
+         "Ma'allanate": 116,
+         'Kar-Mullissu': 117,
+         'mod. Naqs-i-Rustam': 118
+    }
+    
+    GENRE_INDICES = {
+        
+        'Administrative': 1,
+        'Letter': 2,
+        'Legal': 3,
+        'Royal/Monumental': 4,
+        'Literary': 5,
+        'Lexical': 6,
+        'Omen': 7,
+        'uncertain': 8,
+        'Administrative ?': 1,
+        'School': 9,
+        'Mathematical': 10,
+        'Prayer/Incantation': 11,
+        'Lexical ?': 6,
+        'Scientific': 12,
+        'Ritual': 13,
+        'Letter ?': 2,
+        'Literary ?': 5,
+        'fake (modern)': 14,
+        'Lexical; Literary': 6,
+        'Legal ?': 3,
+        'Literary; Mathematical': 5,
+        'Astronomical': 15,
+        'Lexical; Mathematical': 6,
+        'School ?': 9,
+        'Mathematical ?': 10,
+        'Royal/Monumental ?': 4,
+        'Private/Votive': 16,
+        'fake (modern) ?': 14,
+        'Other (see subgenre)': 8,
+        'Historical': 2,
+        'Literary; Lexical': 5,
+        'Lexical; Literary; Mathematical': 6,
+        'Literary; Administrative': 5,
+        'Literary; Letter': 5,
+        'Scientific ?': 12,
+        'Royal/Monumental; Literary': 4,
+        'Private/Votive ?': 16,
+        'School; Literary': 9,
+        'Prayer/Incantation ?': 11,
+        'Ritual ?': 13,
+        'Lexical; School': 6
+    }
+    
+    def __init__(self, CATALOGUE_FN='output/cdli_catalogue_data.csv', IMG_DIR='output/images', IDS=None, mask=False):
+        
+        df = pd.read_csv(
+            CATALOGUE_FN, usecols=['id_text', 'era', 'period_normed', 'provenience_normed', 'genre'], dtype={'id_text': object}
+        ).dropna(subset=['era'])
+        
+        df["id_text"] = df.id_text.apply(lambda x: pad_zeros(x))
+        df = df[df['period_normed'].isin(TabletPeriodDataset.PERIOD_INDICES.keys())]
+        
+        self.id2period = df.set_index('id_text').to_dict()['period_normed']
+        self.id2provenience = df.set_index('id_text').to_dict()['provenience_normed']
+        self.id2genre = df.set_index('id_text').to_dict()['genre']
+        self.genre = df.set_index('id_text').to_dict()['genre']
+        self.img_fns = glob(os.path.join(IMG_DIR, '*.png'))
+        self.IDS = [os.path.basename(fn).rstrip('.png') for fn in self.img_fns]
+        
+        if IDS is not None:
+            print(f'Filtering {len(self.IDS)} IDS down to provided {len(IDS)}...')
+            IDS_set = set(IDS)
+            indices = [i for i, ID in enumerate(self.IDS) if ID in IDS_set]
+            self.img_fns = [self.img_fns[i] for i in indices]
+            self.IDS = [self.IDS[i] for i in indices]
+        
+        self.mask = mask
+        
+    def __len__(self):
+        return len(self.IDS)
+        
+    def __getitem__(self, idx):
+        fn = self.img_fns[idx]
+        ID = self.IDS[idx]
+        try:
+            period = self.id2period[ID]
+        except KeyError as ke:
+            #print('Key Not Found in Period Dictionary:', ke)
+            period = 0
+
+        try:
+            genre = self.id2genre[ID]
+        except KeyError as ke:
+            #print('Key Not Found in Period Dictionary:', ke)
+            genre = 8 # other/uncertain
+
+        try:
+            provenience = self.id2provenience[ID]
+        except KeyError as ke:
+            #print('Key Not Found in Period Dictionary:', ke)
+            provenience = 3 # unknown
+        
+        img = np.asarray(Image.open(fn))
+        alpha = 3 # Contrast control (1.0-3.0)
+        beta = 0 # Brightness control (0-100)
+        adjusted = cv2.convertScaleAbs(img, alpha=alpha, beta=beta)
+        img = img.astype(np.float32) / 255
+        img = cv2.GaussianBlur(img, (11,11), 0)
+        if self.mask:
+            img = (img > 0.125).astype(np.float32) ### 0.25 was great for most besides the really dark ones
+        
+        
+        return ID, img, self.PERIOD_INDICES.get(period, 0), self.GENRE_INDICES.get(genre, 8), self.PROVENIENCE_INDICES.get(provenience, 3) # 0: other