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 matplotlib.pyplot as plt

import os

import torch

import torchvision

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)

model = torchvision.models.resnet50(weights=torchvision.models.resnet.ResNet50_Weights.IMAGENET1K_V1)
model

for param in model.parameters():
    param.requires_grad = False
model.fc = torch.nn.Linear(model.fc.in_features, 2)
loss_fn = torch.nn.CrossEntropyLoss()
model.to(device)
params_to_update = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.Adam(params_to_update, lr=0.0003)

epochs = 7
train_losses = []
valid_losses = []
for epoch in range(epochs):
    running_loss = 0.0
    running_corrects = 0
    model.train()
    with tqdm(total=len(train_ds), desc=f'Train epoch {epoch}/{epochs}', unit='img') as pb:
        for X, y in train_dl:
            X, y = (X.to(device), y.to(device))
            pred = model(X)
            loss = loss_fn(pred, y)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            running_loss += loss.item() * X.size(0)
            running_corrects += torch.sum(pred.argmax(1) == y.data)
            pb.update(X.shape[0])
            pb.set_postfix(**{'loss (batch)': loss.item()})
    epoch_loss = running_loss / len(train_ds)
    train_losses.append(epoch_loss)
    epoch_acc = running_corrects / len(train_ds)
    running_loss = 0.0
    model.eval()
    with torch.no_grad():
        corrects = 0
        with tqdm(total=len(valid_ds), desc=f'Validation epoch {epoch}/{epochs - 1}', unit='img') as pb:
            for X, y in valid_dl:
                X, y = (X.to(device), y.to(device))
                outputs = model(X)
                _, predicted = torch.max(outputs, 1)
                loss = loss_fn(outputs, y)
                running_loss += loss.item() * X.size(0)
                corrects += (predicted == y).sum().item()
                pb.update(X.shape[0])
    epoch_loss = running_loss / len(valid_ds)
    accuracy = corrects / len(valid_ds)
    valid_losses.append(epoch_loss)

plt.plot(range(epochs), train_losses, label='train')
plt.plot(range(epochs), valid_losses, label='valid')
plt.title('Train and test loss')
plt.legend()
plt.show()