dream

1
#!/usr/bin/env python
2

3
import json
4
import numpy as np
5
import pandas as pd
6
import wget
7
from time import time
8

9

10
collect_movies_based_on_rating = False
11
collect_movies_based_on_numvotes = True
12

13
for url in [
14
    "https://datasets.imdbws.com/name.basics.tsv.gz",
15
    "https://datasets.imdbws.com/title.ratings.tsv.gz",
16
    "https://datasets.imdbws.com/title.akas.tsv.gz",
17
    "https://datasets.imdbws.com/title.basics.tsv.gz",
18
    "https://datasets.imdbws.com/title.crew.tsv.gz",
19
    "https://datasets.imdbws.com/title.episode.tsv.gz",
20
    "https://datasets.imdbws.com/title.principals.tsv.gz",
21
    "https://datasets.imdbws.com/title.episode.tsv.gz",
22
]:
23
    filename = wget.download(url)
24

25
# # Choose imdb-ids of most popular movies
26

27
fpath = "./title.ratings.tsv.gz"
28
df_ratings = pd.read_table(fpath, low_memory=False)
29

30
df = pd.read_table(
31
    "./title.basics.tsv.gz", low_memory=False, na_values={"startYear": ["\\N"], "endYear": ["\\N"], "isAdult": ["\\N"]}
32
)
33

34
df = df.merge(df_ratings, on="tconst")
35

36
# fill start year values
37
df["startYear"] = df["startYear"].fillna(value=df["startYear"])
38
df["startYear"] = df["startYear"].fillna(value=0)
39
df["endYear"] = df["endYear"].fillna(value=df["startYear"])
40
df["isAdult"] = df["isAdult"].fillna(value=0)
41

42
df = df.astype(dtype={"startYear": np.int32, "endYear": np.int32, "isAdult": np.int32})
43

44
target = ["movie", "tvMovie", "tvSeries", "tvMiniSeries"]
45
ind_drop = df[~df["titleType"].isin(target)].index
46
df = df.drop(ind_drop)
47

48
if collect_movies_based_on_rating:
49
    movies_ids = []
50
    movies_ids.extend(df.loc[(df["startYear"] <= 1990) & (df["averageRating"] > 8), "tconst"].values)
51
    movies_ids.extend(
52
        df.loc[(df["startYear"] > 1990) & (df["startYear"] <= 2005) & (df["averageRating"] > 7), "tconst"].values
53
    )
54
    movies_ids.extend(
55
        df.loc[(df["startYear"] > 2005) & (df["startYear"] <= 2015) & (df["averageRating"] > 6), "tconst"].values
56
    )
57
    movies_ids.extend(
58
        df.loc[(df["startYear"] > 2015) & (df["startYear"] <= 2021) & (df["averageRating"] > 5), "tconst"].values
59
    )
60

61
    with open("imdb_ids.txt", "w") as f:
62
        for movie_id in movies_ids:
63
            f.write(str(movie_id) + "\n")
64

65
if collect_movies_based_on_numvotes:
66
    movies_ids = []
67
    movies_ids.extend(df.loc[df.loc[:, "numVotes"] > 1000, "tconst"].values)
68
    movies_ids.extend(df.loc[df.loc[:, "averageRating"] > 6.0, "tconst"].values)
69
    movies_ids = list(set(movies_ids))
70

71
    with open("imdb_ids.txt", "w") as f:
72
        for movie_id in movies_ids:
73
            f.write(str(movie_id) + "\n")
74

75
with open("imdb_ids.txt", "r") as f:
76
    all_movies_ids = f.read().splitlines()
77

78
all_movies_ids = list(set(all_movies_ids))
79
print(f"Total number of considered movies: {len(all_movies_ids)}")
80

81
# # Collect titles and ratings
82

83
t0 = time()
84
fpath = "./title.ratings.tsv.gz"
85
df_ratings = pd.read_table(fpath, low_memory=False)
86

87
ind_drop = df_ratings[~df_ratings["tconst"].isin(all_movies_ids)].index
88
df_ratings = df_ratings.drop(ind_drop)
89
assert df_ratings.shape[0] == len(all_movies_ids), print("Number of samples less than number of movies")
90

91
fpath = "./title.basics.tsv.gz"
92

93
df = pd.read_table(fpath, low_memory=False, na_values={"startYear": ["\\N"], "endYear": ["\\N"], "isAdult": ["\\N"]})
94

95
ind_drop = df[~df["tconst"].isin(all_movies_ids)].index
96
df = df.drop(ind_drop)
97

98
df = df.merge(df_ratings, on="tconst")
99

100
df.rename(
101
    columns={
102
        "originalTitle": "original title",
103
        "primaryTitle": "title",
104
        "genres": "genre",
105
        "averageRating": "imdb_rating",
106
        "tconst": "imdb_id",
107
    },
108
    inplace=True,
109
)
110

111
df.drop_duplicates(inplace=True)
112

113
df["titleType"] = df["titleType"].apply(lambda x: "Series" if "Series" in x else "")
114
df["genre"] = [",".join([x, y]) if y != "" else x for x, y in zip(df["genre"], df["titleType"])]
115
df["genre"] = df["genre"].apply(lambda x: x if x != "\\N" else "")
116
df["genre"] = df["genre"].apply(lambda x: x.split(","))
117

118
df.fillna({"startYear": 0, "endYear": 0}, inplace=True)
119
df["startYear"] = df["startYear"].astype("int")
120
df["endYear"] = df["endYear"].astype("int")
121
df.drop(["titleType", "isAdult", "runtimeMinutes"], axis=1, inplace=True)
122
assert df.shape[0] == len(all_movies_ids), print("Number of samples less than number of movies")
123

124
print(f"Total time: {time() - t0}")
125

126
# # Collect names of actors etc
127

128
t0 = time()
129
fpath = "./title.principals.tsv.gz"
130

131
df_principals = pd.read_table(fpath)
132
df_principals = df_principals.loc[:, ["tconst", "nconst", "ordering", "category", "characters"]]
133
df_principals.rename(columns={"tconst": "imdb_id"}, inplace=True)
134
print(df_principals.head())
135

136
ind_drop = df_principals[~df_principals["imdb_id"].isin(all_movies_ids)].index
137
df_principals = df_principals.drop(ind_drop)
138
print(df_principals.head())
139

140
ind_drop = df_principals[~df_principals["ordering"].isin([1, 2, 3, 4, 5, 6])].index
141
df_principals = df_principals.drop(ind_drop)
142
print(df_principals.head())
143

144
df_principals["category"] = df_principals["category"].apply(lambda x: x if x != "actress" else "actor")
145
target_profs = ["director", "producer", "actor", "writer"]
146
ind_drop = df_principals[~df_principals["category"].isin(target_profs)].index
147
df_principals = df_principals.drop(ind_drop)
148
print(df_principals.head())
149

150
fpath = "./name.basics.tsv.gz"
151

152
df_names = pd.read_table(fpath)
153
df_names = df_names.loc[:, ["primaryName", "nconst"]]
154
print(df_names.head())
155

156
df_principals = df_principals.merge(df_names, on="nconst")
157
print(df_principals["characters"])
158

159
special_char = df_principals.loc[4, "characters"]
160
df_principals["characters"] = df_principals["characters"].apply(
161
    lambda x: [] if x == special_char or len(x) == 0 else json.loads(x)
162
)
163
print(df_principals.head())
164

165
print(f"Total time: {time() - t0}")
166

167
# # Collect persons
168
t0 = time()
169

170

171
def collect_movie_persons(x):
172
    return pd.Series(
173
        {
174
            f"{role}s": x.loc[x.sort_values(by=["ordering"])["category"] == prof, name].values.tolist()
175
            for prof, role, name in zip(
176
                ["director", "producer", "actor", "writer", "actor"],
177
                ["director", "producer", "actor", "writer", "character"],
178
                ["primaryName", "primaryName", "primaryName", "primaryName", "characters"],
179
            )
180
        }
181
    )
182

183

184
df_principals = pd.DataFrame(df_principals.groupby("imdb_id").apply(collect_movie_persons))
185
print(df_principals.head())
186
df_principals["characters"] = df_principals["characters"].apply(lambda x: sum(x, []) if isinstance(x, list) else [])
187
print(df_principals.head())
188

189
df.set_index("imdb_id", inplace=True)
190
df = df.join(df_principals, on="imdb_id")
191
df.reset_index(inplace=True)
192

193
df.fillna(value={f"{prof}s": "" for prof in target_profs}, inplace=True)
194

195
assert df.shape[0] == len(all_movies_ids), print("Number of samples less than number of movies")
196

197
print(f"Total time: {time() - t0}")
198

199
# # Collect alternative titles
200

201
t0 = time()
202

203
fpath = "./title.akas.tsv.gz"
204

205
df_akas = pd.read_table(fpath, low_memory=False)
206
df_akas = df_akas.loc[df_akas["region"] == "US", :]
207
df_akas.rename(columns={"titleId": "imdb_id"}, inplace=True)
208

209
ind_drop = df_akas[~df_akas["imdb_id"].isin(all_movies_ids)].index
210
df_akas = df_akas.drop(ind_drop)
211

212
grouped_data = df_akas.groupby("imdb_id")["title"].apply(lambda x: "::".join(x))
213
df_titles = pd.DataFrame(grouped_data)
214
df_titles.rename(columns={"title": "all_titles"}, inplace=True)
215

216
df.set_index("imdb_id", inplace=True)
217
df = df.join(df_titles, on="imdb_id")
218
df.reset_index(inplace=True)
219
df.fillna(value={"all_titles": ""}, inplace=True)
220

221
assert df.shape[0] == len(all_movies_ids), print("Number of samples less than number of movies")
222

223
print(f"Total time: {time() - t0}")
224

225
database = df.to_dict("records")
226
for el in database:
227
    el["genre"] = el["genre"] if el["genre"] != "" else None
228
    el["startYear"] = el["startYear"] if el["startYear"] != 0 else None
229
    el["endYear"] = el["endYear"] if el["endYear"] != 0 else None
230
    el["all_titles"] = el["all_titles"].split("::") if el["all_titles"] != "" else []
231
    for prof in ["director", "producer", "actor", "writer"]:
232
        el[f"{prof}s"] = list(el[f"{prof}s"]) if list(el[f"{prof}s"]) != "" else []
233

234
with open("database_most_popular_main_info.json", "w") as f:
235
    json.dump(database, f, indent=2)
236