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) | |