File size: 1,465 Bytes
b3fb8ff 1b8fb45 091081b 3d6d264 1b8fb45 4fae8fb 1b8fb45 4fae8fb 70044ba 4fae8fb 1b8fb45 4fae8fb 1b8fb45 |
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 |
import os
from ultralytics import YOLO
from PIL import Image
import gradio as gr
from huggingface_hub import snapshot_download
REPO_ID = "jiaxinnnnn/Interior-Style-Classification"
WEIGHTS_FILENAME = "Best_Accuracy.pt"
def load_model(repo_id: str) -> YOLO:
"""
Downloads the model repo from Hugging Face Hub into the Space cache
and loads the YOLO model from the weights file.
"""
download_dir = snapshot_download(repo_id=repo_id)
weights_path = os.path.join(download_dir, WEIGHTS_FILENAME)
if not os.path.exists(weights_path):
raise FileNotFoundError(
f"Cannot find {WEIGHTS_FILENAME} inside downloaded repo: {download_dir}\n"
f"Files found: {os.listdir(download_dir)}"
)
return YOLO(weights_path, task="classify")
classification_model = load_model(REPO_ID)
def predict(pilimg: Image.Image):
results = classification_model.predict(pilimg)
probs = results[0].probs
class_id = probs.top1
confidence = probs.top1conf.item()
class_name = classification_model.names[class_id]
return {
"style": class_name.title(),
"confidence": confidence
}
demo = gr.Interface(
fn=predict,
inputs=gr.Image(type="pil", label="Upload an interior image"),
outputs=gr.JSON(label="Prediction"),
title="Interior Style Classification (YOLOv8)",
description="Upload an image and the model will classify the interior design style."
)
demo.launch()
|