File size: 3,030 Bytes
2385a75 |
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 |
"""
Example usage script for the Skin Type Classification model on Hugging Face.
"""
from transformers import AutoModelForImageClassification, AutoImageProcessor
from PIL import Image
import torch
import requests
from io import BytesIO
def load_model(model_name="your-username/skin-type-classifier"):
"""Load the model and processor from Hugging Face."""
model = AutoModelForImageClassification.from_pretrained(model_name)
processor = AutoImageProcessor.from_pretrained(model_name)
return model, processor
def predict_skin_type(image_path_or_url, model, processor):
"""
Predict skin type from an image.
Args:
image_path_or_url: Path to local image or URL
model: The loaded model
processor: The loaded processor
Returns:
dict: Prediction results with class and confidence
"""
# Load image
if image_path_or_url.startswith(('http://', 'https://')):
response = requests.get(image_path_or_url)
image = Image.open(BytesIO(response.content))
else:
image = Image.open(image_path_or_url)
# Convert to RGB if needed
if image.mode != 'RGB':
image = image.convert('RGB')
# Process image
inputs = processor(images=image, return_tensors="pt")
# Make prediction
with torch.no_grad():
outputs = model(**inputs)
predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
predicted_class_idx = predictions.argmax().item()
confidence = predictions[0][predicted_class_idx].item()
# Map to class names
class_names = {0: "dry", 1: "oily"}
predicted_class = class_names[predicted_class_idx]
return {
"predicted_class": predicted_class,
"confidence": confidence,
"all_scores": {
"dry": predictions[0][0].item(),
"oily": predictions[0][1].item()
}
}
def main():
"""Example usage of the skin type classification model."""
print("🔬 Loading Skin Type Classification Model...")
# Load model and processor
model, processor = load_model()
print("✅ Model loaded successfully!")
# Example with local image (replace with your image path)
try:
image_path = "example_skin_image.jpg" # Replace with actual image path
result = predict_skin_type(image_path, model, processor)
print(f"\n📊 Prediction Results:")
print(f"Predicted Skin Type: {result['predicted_class']}")
print(f"Confidence: {result['confidence']:.2%}")
print(f"All Scores: {result['all_scores']}")
except FileNotFoundError:
print("ℹ️ Please provide a valid image path to test the model")
# Example usage patterns
print("\n💡 Usage Examples:")
print("1. Local image: predict_skin_type('path/to/image.jpg', model, processor)")
print("2. URL image: predict_skin_type('https://example.com/image.jpg', model, processor)")
if __name__ == "__main__":
main()
|