ITE / app.py
Soooma's picture
Approach 2 is working
cf2a1b9 verified
import torch
from transformers import CLIPProcessor, CLIPModel
from PIL import Image
import gradio as gr
device = "cuda" if torch.cuda.is_available() else "cpu"
# Load the model
model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device)
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
# Load normal image for reference
normal_image = Image.open("normal_sample.jpg")
with torch.no_grad():
inputs = processor(images=normal_image, return_tensors="pt").to(device)
normal_features = model.get_image_features(**inputs)
normal_features = normal_features / normal_features.norm(p=2, dim=-1, keepdim=True)
def detect_anomaly(img):
with torch.no_grad():
inputs = processor(images=img, return_tensors="pt").to(device)
test_features = model.get_image_features(**inputs)
test_features = test_features / test_features.norm(p=2, dim=-1, keepdim=True)
similarity = (test_features @ normal_features.T).item()
if similarity < 0.8: # threshold example
result = "Anomaly Detected"
else:
result = "Normal"
return f"Similarity: {similarity:.2f} | {result}"
gr.Interface(fn=detect_anomaly,
inputs=gr.Image(),
outputs="text").launch()