text_plain_output_1.png
from PIL import Image

from sklearn.model_selection import train_test_split

from torch.utils.data import DataLoader

from tqdm import tqdm

import json

import os

import torch

import torchvision.transforms as T

import torch
from torch.utils.data import DataLoader
import torchvision
import torchvision.transforms as T
from torchvision.io import read_image, ImageReadMode
from sklearn.model_selection import train_test_split
import os
import json
from tqdm import tqdm
from PIL import Image
import matplotlib.pyplot as plt
import random
Image.MAX_IMAGE_PIXELS = None

ds_path = '/kaggle/input/raw-mmimdb/mmimdb/dataset/'
ids = list(set([x.split('.')[0] for x in os.listdir(ds_path)]))

img = f'{ds_path}{ids[0]}.jpeg'
meta = f'{ds_path}{ids[0]}.json'
Image.open(img)

with open(meta, 'r') as f:
    metadata = json.load(f)

horror_ids = []
drama_ids = []
for id in tqdm(ids):
    metadata_file = f'{ds_path}{id}.json'
    with open(metadata_file, 'r') as f:
        metadata = json.load(f)
    genres = metadata['genres']
    if 'Horror' in genres and 'Drama' not in genres:
        horror_ids.append(id)
    if 'Drama' in genres and 'Horror' not in genres:
        drama_ids.append(id)

horror_ids = horror_ids[:2000]
drama_ids = drama_ids[:2000]

ids = horror_ids + drama_ids
labels = [0] * len(horror_ids) + [1] * len(drama_ids)

class ImageDataset(torch.utils.data.Dataset):
    def __init__(self, image_folder, ids, labels, transform=None):
        self.image_folder = image_folder
        self.ids = ids
        self.labels = labels
        self.transform = transform
    def __len__(self):
        return len(self.ids)
    def __getitem__(self, idx):
        image_path = self.image_folder + self.ids[idx] + '.jpeg'
        image = Image.open(image_path).convert('RGB')
        label = self.labels[idx]
        if self.transform:
            image = self.transform(image)
        return (image, label)

X_train, X_valid, y_train, y_valid = train_test_split(ids, labels, test_size=0.4, random_state=42, stratify=labels)
X_valid, X_test, y_valid, y_test = train_test_split(X_valid, y_valid, test_size=0.5, random_state=42, stratify=y_valid)
img_height, img_width = (512, 512)
transform = T.Compose([T.ToTensor(), T.Resize((img_height, img_width)), T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])])
train_ds = ImageDataset(ds_path, X_train, y_train, transform=transform)
valid_ds = ImageDataset(ds_path, X_valid, y_valid, transform=transform)
test_ds = ImageDataset(ds_path, X_test, y_test, transform=transform)
batch_size = 10
device = 'cuda' if torch.cuda.is_available() else 'cpu'
loader_args = dict(batch_size=batch_size, num_workers=os.cpu_count(), pin_memory=True, shuffle=True)
train_dl = DataLoader(train_ds, **loader_args)
valid_dl = DataLoader(valid_ds, **loader_args)
test_dl = DataLoader(test_ds, **loader_args)
print('Sizes of the datasets: ', len(train_ds), len(valid_ds), len(test_ds))