DmitryRyumin's picture
Summary
d5e9efc
raw
history blame
1.66 kB
"""
File: model.py
Author: Elena Ryumina and Dmitry Ryumin
Description: This module provides functions for loading and processing a pre-trained deep learning model
for facial expression recognition.
License: MIT License
"""
import torch
import requests
from PIL import Image
from torchvision import transforms
# Importing necessary components for the Gradio app
from app.config import config_data
def load_model(model_url, model_path):
try:
# Загрузка модели
with requests.get(model_url, stream=True) as response:
with open(model_path, "wb") as file:
for chunk in response.iter_content(chunk_size=8192):
file.write(chunk)
return torch.jit.load(model_path).eval()
except Exception as e:
print(f"Error loading model: {e}")
return None
# Загрузите модель
pth_model = load_model(config_data.model_url, config_data.model_path)
def pth_processing(fp):
class PreprocessInput(torch.nn.Module):
def init(self):
super(PreprocessInput, self).init()
def forward(self, x):
x = x.to(torch.float32)
x = torch.flip(x, dims=(0,))
x[0, :, :] -= 91.4953
x[1, :, :] -= 103.8827
x[2, :, :] -= 131.0912
return x
def get_img_torch(img, target_size=(224, 224)):
transform = transforms.Compose([transforms.PILToTensor(), PreprocessInput()])
img = img.resize(target_size, Image.Resampling.NEAREST)
img = transform(img)
img = torch.unsqueeze(img, 0)
return img
return get_img_torch(fp)