Spaces:
Sleeping
Sleeping
import os | |
import streamlit as st | |
import cv2 | |
import numpy as np | |
from skimage.io import imread | |
from skimage.transform import resize | |
from skimage.feature import hog | |
import torch | |
import torch.nn as nn | |
import torchvision.models as models | |
import torchvision.transforms as transforms | |
from PIL import Image | |
import pickle | |
# Load the pretrained resnet-50 model | |
resnet = models.resnet50(pretrained=True) | |
resnet = nn.Sequential(*list(resnet.children())[:-1]) | |
resnet.eval() | |
# Load the demo model | |
with open('demo.pkl', 'rb') as f: | |
demo = pickle.load(f) | |
# defining some of the important functions that are needed | |
# defining feature extraction for resnet-50 | |
def extract_features(image, model): | |
image_np = np.array(image) | |
preprocess = transforms.Compose([ | |
transforms.ToPILImage(), | |
transforms.Resize(256), | |
transforms.CenterCrop(224), | |
transforms.ToTensor(), | |
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), | |
]) | |
image = preprocess(image_np) | |
image = image.unsqueeze(0) | |
with torch.no_grad(): | |
features = model(image) | |
features = features.squeeze(0) | |
return features | |
# hog calculation | |
def compute_hog(img): | |
resized_img = resize(img, (128*4, 64*4)) | |
fd, hog_image = hog(resized_img, orientations=9, pixels_per_cell=(8, 8), | |
cells_per_block=(2, 2), visualize=True,channel_axis=-1) | |
return fd | |
def get_pixel(img, center, x, y): | |
new_value = 0 | |
try: | |
if img[x][y] >= center: | |
new_value = 1 | |
except: | |
pass | |
return new_value | |
# calculate lbp | |
def lbp_calculated_pixel(img, x, y): | |
center = img[x][y] | |
val_ar = [] | |
val_ar.append(get_pixel(img, center, x-1, y+1)) | |
val_ar.append(get_pixel(img, center, x, y+1)) | |
val_ar.append(get_pixel(img, center, x+1, y+1)) | |
val_ar.append(get_pixel(img, center, x+1, y)) | |
val_ar.append(get_pixel(img, center, x+1, y-1)) | |
val_ar.append(get_pixel(img, center, x, y-1)) | |
val_ar.append(get_pixel(img, center, x-1, y-1)) | |
val_ar.append(get_pixel(img, center, x-1, y)) | |
power_val = [1, 2, 4, 8, 16, 32, 64, 128] | |
val = 0 | |
for i in range(len(val_ar)): | |
val += val_ar[i] * power_val[i] | |
return val | |
def calcLBP(img): | |
height, width, channel = img.shape | |
img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) | |
img_lbp = np.zeros((height, width,3), np.uint8) | |
for i in range(0, height): | |
for j in range(0, width): | |
img_lbp[i, j] = lbp_calculated_pixel(img_gray, i, j) | |
hist_lbp = cv2.calcHist([img_lbp], [0], None, [256], [0, 256]) | |
return hist_lbp.flatten() | |
# Function to infer the class of an image | |
def infer(image_path): | |
image = imread(image_path) # reading the image | |
hog_feature = compute_hog(image) # hog features | |
lbp_feature = calcLBP(image) # lbp features | |
cnn_feature = extract_features(image, resnet).numpy() #cnn features | |
hog_feature = hog_feature.reshape(-1) # reshaping | |
lbp_feature = lbp_feature.reshape(-1) #reshaping | |
cnn_feature = cnn_feature.flatten() # cnn features | |
combined_feature = np.concatenate((hog_feature, lbp_feature, cnn_feature)) # combining all the features | |
prediction = demo.predict([combined_feature]) | |
return prediction[0] | |
# Streamlit code | |
st.title("Face Identification") | |
uploaded_file = st.file_uploader("Choose an image...", type=["png"]) | |
if uploaded_file is not None: | |
with open("temp_image.png", "wb") as f: | |
f.write(uploaded_file.getvalue()) | |
image = Image.open("temp_image.png") | |
st.image(image, caption='Uploaded Image', use_column_width=True) | |
st.write("") | |
st.write("Classifying...") | |
prediction = infer("temp_image.png") | |
st.write("Prediction:", prediction) | |