<jupyter_start><jupyter_text>---  Prepare dataset<jupyter_code>TRAIN_RATIO = 0.9
def is_specialfile(path,exts):
    _, file_extension = os.path.splitext(path)
    return file_extension.lower() in exts

img_extensions=['.jpg','.jpeg','.png']
def is_image(path):
    return is_specialfile(path,img_extensions)
from facial_analysis import FacialImageProcessing
imgProcessing=FacialImageProcessing()
def process_image(src_file,dest_file):
    try:
        #Image.open(src_file).crop(bbox).save(dest_file)
        #Image.open(src_file).thumbnail(INPUT_SIZE, Image.ANTIALIAS).save(dest_file)
        im = Image.open(src_file)
        im=im.resize(TARGET_SIZE, Image.LANCZOS)
        im.save(dest_file)
        #print(src_file,dest_file)
    except IOError:
        print ("cannot create thumbnail for '%s'" % src_file)


def process_face_image(src_file,dest_file):
    try:
        draw = cv2.imread(src_file)
        img=cv2.cvtColor(draw,cv2.COLOR_BGR2RGB)
        bounding_boxes, points = imgProcessing.detect_faces(img)
        best_bb=[]
        best_square=0
        for b in bounding_boxes:
            b=[int(bi) for bi in b]
            #print(b,img.shape)
            x1,y1,x2,y2=b[0:4]
            if x2>x1 and y2>y1:
                sq=(x2-x1)*(y2-y1)
                if sq>best_square:
                    best_square=sq
                    best_bb=b
        
        if len(best_bb)!=0:
            img_h,img_w,_=img.shape
            face_x,face_y=best_bb[0],best_bb[1]
            face_w,face_h=(best_bb[2]-best_bb[0]),(best_bb[3]-best_bb[1])
            dw,dh=20,40#max(int(face_w*0.05),10),max(int(face_h*0.05),10)
            #sz=max(face_w+2*dw,face_h+2*dh)
            #dw,dh=(sz-face_w)//2,(sz-face_h)//2

            box = (max(0,face_x-dw), max(0,face_y-dh), min(img_w,face_x+face_w+dw), min(img_h,face_y+face_h+dh))
            
            face_img=draw[box[1]:box[3],box[0]:box[2],:]
            face_img = cv2.resize(face_img, TARGET_SIZE)
            cv2.imwrite(dest_file,face_img)
        else:
            print('No faces found for ', src_file)
        
        #print(src_file,dest_file)
    except IOError:
        print ("cannot create facial image for '%s'" % src_file)

for person in log_progress(os.listdir(ALL_DATA_DIR)):
    #print(person)
    person_dir=os.path.join(ALL_DATA_DIR,person)
    files=[f for f in os.listdir(person_dir) if is_image(f)]
    shuffle(files)
    train_files=files[: int(len(files) * TRAIN_RATIO)]
    train_dir=os.path.join(TRAIN_DATA_DIR,person)
    if not os.path.exists(train_dir):
        os.makedirs(train_dir)
    for f in train_files:
        #process_image(os.path.join(person_dir,f),os.path.join(train_dir,f))
        process_face_image(os.path.join(person_dir,f),os.path.join(train_dir,f))

    val_files=files[int(len(files) * TRAIN_RATIO):]
    val_dir=os.path.join(VAL_DATA_DIR,person)
    if not os.path.exists(val_dir):
        os.makedirs(val_dir)
    for f in val_files:
        #process_image(os.path.join(person_dir,f),os.path.join(val_dir,f))
        process_face_image(os.path.join(person_dir,f),os.path.join(val_dir,f))<jupyter_output><empty_output><jupyter_text>---  Train<jupyter_code>import glob
from itertools import chain
import os
import random
import zipfile

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from PIL import Image
from sklearn.model_selection import train_test_split
from torch.optim.lr_scheduler import StepLR
from torch.utils.data import DataLoader, Dataset
from torchvision import datasets, transforms
from tqdm.notebook import tqdm

print(f"Torch: {torch.__version__}")
# Training settings
batch_size = 48#64 #48# 32 #8 #16 #
epochs = 40
lr = 3e-5
gamma = 0.7
seed = 42
device = 'cuda'
use_cuda = torch.cuda.is_available()
print(use_cuda)
#net_description='affectnet_'+net_description
train_dir,test_dir=TRAIN_DATA_DIR,VAL_DATA_DIR

print(train_dir,test_dir)
IMG_SIZE_ORIG=256
IMG_SIZE=224 # 300 # 80 #
train_transforms = transforms.Compose(
    [
        transforms.Resize((IMG_SIZE_ORIG,IMG_SIZE_ORIG)),
        #transforms.RandomAffine(degrees=10, translate=(0.1, 0.1), scale=(0.8,1.2), shear=0.2),
        transforms.RandomResizedCrop(IMG_SIZE),
        #transforms.Resize((IMG_SIZE,IMG_SIZE)),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])
    ]
)

test_transforms = transforms.Compose(
    [
        transforms.Resize((IMG_SIZE,IMG_SIZE)),
        #transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])
    ]
)
kwargs = {'num_workers': 4, 'pin_memory': True} if use_cuda else {}

train_dataset = datasets.ImageFolder(root=train_dir, transform=train_transforms)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True, **kwargs)
test_dataset = datasets.ImageFolder(root=test_dir, transform=test_transforms)
test_loader  = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False, **kwargs) 

print(len(train_dataset), len(test_dataset))
num_classes=len(train_dataset.classes)
print(num_classes)
# loss function
criterion = nn.CrossEntropyLoss()
import sys
sys.path.append('/home/avsavchenko/src/distr/sam/')
from sam import SAM
def test(model):
    epoch_val_accuracy = 0
    epoch_val_loss = 0
    model.eval()
    with torch.no_grad():
        for data, label in test_loader:
            data = data.to(device)
            label = label.to(device)

            val_output = model(data)
            val_loss = criterion(val_output, label)

            acc = (val_output.argmax(dim=1) == label).float().mean()
            epoch_val_accuracy += acc / len(test_loader)
            epoch_val_loss += val_loss / len(test_loader)
    return epoch_val_accuracy,epoch_val_loss
import copy
def train(model,n_epochs=epochs, learningrate=lr, use_sam=False):
    # optimizer
    if use_sam:
        optimizer = SAM(filter(lambda p: p.requires_grad, model.parameters()), optim.Adam, lr=learningrate)
    else:
        optimizer=optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=learningrate)
    # scheduler
    #scheduler = StepLR(optimizer, step_size=1, gamma=gamma)
    best_acc=0
    best_model=None
    for epoch in range(n_epochs):
        epoch_loss = 0
        epoch_accuracy = 0
        model.train()
        for data, label in tqdm(train_loader):
            data = data.to(device)
            label = label.to(device)

            output = model(data)
            loss = criterion(output, label)

            if use_sam:
                #optimizer.zero_grad()
                loss.backward()
                optimizer.first_step(zero_grad=True)
  
                # second forward-backward pass
                output = model(data)
                loss = criterion(output, label)
                loss.backward()
                optimizer.second_step(zero_grad=True)
            else:
                optimizer.zero_grad()
                loss.backward()
                optimizer.step()

            acc = (output.argmax(dim=1) == label).float().mean()
            epoch_accuracy += acc / len(train_loader)
            epoch_loss += loss / len(train_loader)

        epoch_val_accuracy,epoch_val_loss=test(model)
        print(
            f"Epoch : {epoch+1} - loss : {epoch_loss:.4f} - acc: {epoch_accuracy:.4f} - val_loss : {epoch_val_loss:.4f} - val_acc: {epoch_val_accuracy:.4f}\n"
        )
        if best_acc<epoch_val_accuracy:
            best_acc=epoch_val_accuracy
            best_model=copy.deepcopy(model.state_dict())
        #scheduler.step()
    
    if best_model is not None:
        model.load_state_dict(best_model)
        print(f"Best acc:{best_acc}")
        epoch_val_accuracy,epoch_val_loss=test(model)
        print(
            f"val_loss : {epoch_val_loss:.4f} - val_acc: {epoch_val_accuracy:.4f}\n"
        )
    else:
        print(f"No best model Best acc:{best_acc}")
from torchvision.models import resnet101,mobilenet_v2
import timm
#model=resnet101(pretrained=True)
#model=mobilenet_v2(pretrained=True)
#model=torch.hub.load('rwightman/gen-efficientnet-pytorch', 'efficientnet_b0', pretrained=True)
model=timm.create_model('rexnet_150', pretrained=True) #'vit_base_patch16_224' 'tf_efficientnet_b4_ns'
#model=timm.create_model('tf_efficientnet_b2_ns', pretrained=True,features_only=True) 
print(model)
#model.classifier=nn.Linear(in_features=1536, out_features=num_classes) #1792 #1536 #1280 #1408
model.head.fc=nn.Linear(in_features=1920, out_features=num_classes)
model=model.to(device)
print(model)
if True:
    img = torch.randn(1, 3, IMG_SIZE, IMG_SIZE).to(device)
    model=model.to(device)
    model.eval()
    f=model.forward(img)
    print(f.shape)
    model.train()
#adapted from https://pytorch.org/tutorials/beginner/finetuning_torchvision_models_tutorial.html
def set_parameter_requires_grad(model, requires_grad):
    for param in model.parameters():
        param.requires_grad = requires_grad
set_parameter_requires_grad(model, requires_grad=False)
#set_parameter_requires_grad(model.classifier, requires_grad=True)
set_parameter_requires_grad(model.head.fc, requires_grad=True)
train(model,1,0.001,use_sam=True)
set_parameter_requires_grad(model, requires_grad=True)
train(model,10,1e-4,use_sam=True)
epoch_val_accuracy,epoch_val_loss=test(model)
print(f"val_loss : {epoch_val_loss:.4f} - val_acc: {epoch_val_accuracy:.4f}\n")
#val_loss : 0.2196 - val_acc: 0.9564
#PATH = "vggface2_enet3_new.pt"
PATH = "vggface2_rexnet_150.pt"
torch.save(model, PATH)
model = torch.load(PATH)
model.eval()
if False:
    class Identity(nn.Module):
        def __init__(self):
            super(Identity, self).__init__()

        def forward(self, x):
            return x

#model.classifier=torch.nn.Identity()
model.head.fc=torch.nn.Identity()
torch.save(model.state_dict(), 'state_'+PATH)
torch.save(model, 'features_'+PATH)<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.1225211074023169, 'fraction_numerical': 0.022776359709405065, 'mean_word_length': 3.3238539898132426, 'pattern_counts': {'":': 0, '<': 10, '<?xml version=': 0, '>': 12, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '165461', 'license': 'Apache-2.0', 'n_tokens_mistral': 3754, 'n_tokens_neox': 3613, 'n_words': 564, 'path': 'src/train_faces_torch.ipynb', 'repo_name': 'Xsneo/face-emotion-recognition', 'size': '10186'}