souvikmaji22's picture
Upload 3 files
8dafa03 verified
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)