Soooma commited on
Commit
27fb6da
·
verified ·
1 Parent(s): 86253f3

Add the code to app file

Browse files
Files changed (1) hide show
  1. app.py +46 -0
app.py CHANGED
@@ -0,0 +1,46 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ from transformers import AutoModel, AutoImageProcessor
3
+ from PIL import Image
4
+ import gradio as gr
5
+ import numpy as np
6
+ from sklearn.metrics.pairwise import cosine_similarity
7
+
8
+ device = "cuda" if torch.cuda.is_available() else "cpu"
9
+
10
+ # Load DINOv2 model
11
+ model_name = "facebook/dinov2-base"
12
+ model = AutoModel.from_pretrained(model_name).to(device)
13
+ processor = AutoImageProcessor.from_pretrained(model_name)
14
+
15
+ # Load OK reference image (built-in normal)
16
+ ok_image = Image.open("https://huggingface.co/datasets/Soooma/sofc/OK1.jpg")
17
+
18
+ # Precompute OK feature embedding
19
+ with torch.no_grad():
20
+ ok_input = processor(images=ok_image, return_tensors="pt").to(device)
21
+ ok_feat = model(**ok_input).last_hidden_state.mean(dim=1)
22
+ ok_feat = ok_feat.cpu().numpy()
23
+
24
+ def detect_anomaly(image):
25
+ if image is None:
26
+ return "No image uploaded."
27
+
28
+ with torch.no_grad():
29
+ inputs = processor(images=image, return_tensors="pt").to(device)
30
+ feat = model(**inputs).last_hidden_state.mean(dim=1)
31
+ feat = feat.cpu().numpy()
32
+
33
+ similarity = cosine_similarity(feat, ok_feat)[0][0]
34
+
35
+ if similarity < 0.90:
36
+ return f"Anomaly Detected | Similarity: {similarity:.3f}"
37
+ else:
38
+ return f"Normal | Similarity: {similarity:.3f}"
39
+
40
+ gr.Interface(
41
+ fn=detect_anomaly,
42
+ inputs=gr.Image(type="pil"),
43
+ outputs="text",
44
+ title="Anomaly Detector (DINOv2)",
45
+ description="Upload an image of a stack. The model compares it to a known OK sample using DINOv2 features."
46
+ ).launch()