DeathDaDev commited on
Commit
7078ca9
1 Parent(s): a7d74ed

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +33 -0
app.py ADDED
@@ -0,0 +1,33 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import AutoModelForImageClassification, AutoProcessor
3
+ import torch
4
+
5
+ # Load the model and processor
6
+ model_name = "DeathDaDev/Materializer"
7
+ processor = AutoProcessor.from_pretrained(model_name)
8
+ model = AutoModelForImageClassification.from_pretrained(model_name)
9
+
10
+ # Define the prediction function
11
+ def classify_image(image):
12
+ # Preprocess the image
13
+ inputs = processor(images=image, return_tensors="pt")
14
+
15
+ # Perform inference
16
+ with torch.no_grad():
17
+ logits = model(**inputs).logits
18
+
19
+ # Get the predicted class
20
+ predicted_class_idx = logits.argmax(-1).item()
21
+ return model.config.id2label[predicted_class_idx]
22
+
23
+ # Create the Gradio interface
24
+ iface = gr.Interface(
25
+ fn=classify_image,
26
+ inputs=gr.inputs.Image(type="pil"),
27
+ outputs=gr.outputs.Label(num_top_classes=3),
28
+ title="Image Classification with Materializer",
29
+ description="Upload an image to classify it using the Materializer model."
30
+ )
31
+
32
+ # Launch the interface
33
+ iface.launch()