Update app.py

app.py

@@ -1,33 +1,19 @@
-import gradio as gr
-from transformers import CLIPProcessor, CLIPModel
-from PIL import Image
-import torch
+from transformers import pipeline
 
-# Load the model and processor
-model_id = "DGurgurov/clip-vit-base-patch32-oxford-pets"
-model = CLIPModel.from_pretrained(model_id)
-processor_id = "openai/clip-vit-base-patch32"
-processor = CLIPProcessor.from_pretrained(processor_id)
+# Define the pipeline for image classification
+pipe = pipeline("image-classification", model="DGurgurov/clip-vit-base-patch32-oxford-pets")
 
-# Define the inference function
+# Define the predict function using the pipeline
 def predict(image):
-    # Preprocess the image
-    image = Image.fromarray(image.astype('uint8'))  # Convert numpy array to PIL Image
-    inputs = processor(images=image, return_tensors="pt")
-    
-    # Perform inference
-    outputs = model(**inputs)
-    logits_per_image = outputs.logits_per_image
-    probs = torch.nn.functional.softmax(logits_per_image, dim=1)
-    
-    # Get top 5 class predictions
-    labels = processor.tokenizer.convert_ids_to_tokens(outputs.logits_per_image.argmax(dim=-1))
-    top_5 = labels[torch.topk(probs[0], 5).indices]
-    
-    return {f"Class {i}": label for i, label in enumerate(top_5)}
+    # Perform inference using the pipeline
+    results = pipe(image)
+    return {f"Class {i}": result['label'] for i, result in enumerate(results)}
+
+# Now you can use this predict function in your Gradio interface
+import gradio as gr
 
 # Define Gradio interface
-image = gr.components.Image(type="pil")
+image = gr.components.Image()
 label = gr.components.Label(num_top_classes=5)
 
 interface = gr.Interface(