Spaces:
Runtime error
Runtime error
| import pathlib | |
| import gradio as gr | |
| from loguru import logger | |
| from transformers import AutoFeatureExtractor, AutoModelForImageClassification | |
| logger.info("starting gradio app") | |
| CURRENT_DIR = pathlib.Path(__file__).resolve().parent | |
| APP_NAME = "Mona Lisa Detection" | |
| logger.debug("loading processor and model.") | |
| processor = AutoFeatureExtractor.from_pretrained( | |
| "drift-ai/autotrain-mona-lisa-detection-38345101350", use_auth_token=True | |
| ) | |
| model = AutoModelForImageClassification.from_pretrained( | |
| "drift-ai/autotrain-mona-lisa-detection-38345101350", use_auth_token=True | |
| ) | |
| logger.debug("loading processor and model succeeded.") | |
| def process_image(image, model=model, processor=processor): | |
| logger.info("Making a prediction ...") | |
| inputs = processor(images=image, return_tensors="pt") | |
| outputs = model(**inputs) | |
| logits = outputs.logits | |
| predicted_class_idx = logits.argmax(-1).item() | |
| label = {1: "Not Mona Lisa", 0: "Mona Lisa"} | |
| predictions = logits.softmax(dim=-1).tolist() | |
| result = {label[predicted_class_idx]: predictions[0][predicted_class_idx]} | |
| print("Predicted class:", result) | |
| logger.info("Prediction finished.") | |
| return result | |
| examples = [ | |
| "mona-lisa-1.jpg", | |
| "mona-lisa-2.jpg", | |
| "mona-lisa-3.jpg", | |
| "not-mona-lisa-1.jpg", | |
| "not-mona-lisa-2.jpg", | |
| "not-mona-lisa-3.jpg", | |
| ] | |
| if __name__ == "__main__": | |
| title = """ | |
| Mona Lisa Detection. | |
| """ | |
| app = gr.Interface( | |
| fn=process_image, | |
| inputs=[ | |
| gr.inputs.Image(type="pil", label="Image"), | |
| ], | |
| outputs=gr.Label(label="Predictions:", show_label=True), | |
| examples=examples, | |
| examples_per_page=32, | |
| title=title, | |
| enable_queue=True, | |
| ).launch() | |