Spaces:
Running
Running
File size: 3,774 Bytes
89a78a7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 |
import os
import requests
from flask import Flask, render_template, request, jsonify
from werkzeug.utils import secure_filename
import torch
import torchvision.transforms as T
from PIL import Image
import torch.nn.functional as F # For softmax
app = Flask(__name__)
# Define device
device = "cuda" if torch.cuda.is_available() else "cpu"
# Model and transformation setup
def download_model_if_not_exists(url, model_path):
"""Download model from Hugging Face repository if it doesn't exist locally."""
if not os.path.exists(model_path):
print("Model not found locally, downloading from Hugging Face...")
response = requests.get(url)
if response.status_code == 200:
with open(model_path, 'wb') as f:
f.write(response.content)
print(f"Model downloaded and saved to {model_path}")
else:
print("Failed to download model. Please check the URL.")
else:
print("Model already exists locally.")
def load_model(model_path):
"""Load model from the given path."""
model = torch.load(model_path, map_location=torch.device('cpu'))
model.eval() # Set model to evaluation mode
model.to(device)
return model
def preprocess_image(image_path):
transform = T.Compose([
T.Resize((224, 224)), # Resize image to 224x224
T.ToTensor(), # Convert image to Tensor
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) # Normalize
])
image = Image.open(image_path).convert("RGB") # Open and convert image to RGB
return transform(image).unsqueeze(0) # Add batch dimension
def get_probabilities(logits):
"""Apply softmax to get probabilities."""
probabilities = F.softmax(logits, dim=1)
percentages = probabilities * 100
return percentages
def predict(image_path, model, class_names):
"""Make prediction using the trained model."""
image_tensor = preprocess_image(image_path).to(device)
model.eval()
with torch.inference_mode(): # Disable gradient calculations
outputs = model(image_tensor)
percentages = get_probabilities(outputs)
_, predicted_class = torch.max(outputs, 1) # Get the index of the highest logit
predicted_label = class_names[predicted_class.item()]
return predicted_label, percentages
# Define class names
class_names = ['Heart', 'Oblong', 'Oval', 'Round', 'Square']
# Path to the model file
model_path = r"model_85_nn_.pth" # Update this with the correct model path
model_url = "https://huggingface.co/fahd9999/model_85_nn_/resolve/main/model_85_nn_.pth?download=true"
# Download the model only if it doesn't exist locally
download_model_if_not_exists(model_url, model_path)
# Load the model
model = load_model(model_path)
# API to render the index page
@app.route('/')
def index():
return render_template('index.html')
# API to handle image upload and prediction
@app.route('/predict', methods=['POST'])
def predict_face_shape():
if 'file' not in request.files:
return jsonify({'error': 'No file part'})
file = request.files['file']
if file.filename == '':
return jsonify({'error': 'No selected file'})
if file:
os.makedirs('uploads',exist_ok=True)
filename = secure_filename(file.filename)
file_path = os.path.join('uploads', filename)
file.save(file_path)
predicted_label, percentages = predict(file_path, model, class_names)
result = {class_names[i]: percentages[0, i].item() for i in range(len(class_names))}
sorted_result = dict(sorted(result.items(), key=lambda item: item[1], reverse=True))
print(sorted_result)
return jsonify(sorted_result)
if __name__ == '__main__':
app.run(debug=False)
|