Spaces:
Sleeping
Sleeping
Delete vision.py
Browse files
vision.py
DELETED
@@ -1,50 +0,0 @@
|
|
1 |
-
|
2 |
-
# models/vision.py -- Working
|
3 |
-
|
4 |
-
from transformers import pipeline
|
5 |
-
from PIL import Image
|
6 |
-
|
7 |
-
class VisionModel:
|
8 |
-
def __init__(
|
9 |
-
self,
|
10 |
-
model_name: str = "valentinafeve/yolos-fashionpedia",
|
11 |
-
threshold: float = 0.7
|
12 |
-
):
|
13 |
-
self.pipe = pipeline("object-detection", model=model_name)
|
14 |
-
self.threshold = threshold
|
15 |
-
|
16 |
-
def detect(self, image: Image.Image):
|
17 |
-
# 1) Ensure RGB
|
18 |
-
if image.mode != "RGB":
|
19 |
-
image = image.convert("RGB")
|
20 |
-
|
21 |
-
# 2) Run detection
|
22 |
-
results = self.pipe(image)
|
23 |
-
|
24 |
-
# 3) Process & filter
|
25 |
-
processed = []
|
26 |
-
for r in results:
|
27 |
-
score = float(r["score"])
|
28 |
-
if score < self.threshold:
|
29 |
-
continue
|
30 |
-
|
31 |
-
# r["box"] is a dict: {"xmin":..., "ymin":..., "xmax":..., "ymax":...}
|
32 |
-
box = r["box"]
|
33 |
-
coords = [
|
34 |
-
float(box["xmin"]),
|
35 |
-
float(box["ymin"]),
|
36 |
-
float(box["xmax"]),
|
37 |
-
float(box["ymax"]),
|
38 |
-
]
|
39 |
-
|
40 |
-
processed.append({
|
41 |
-
"label": r["label"],
|
42 |
-
"score": score,
|
43 |
-
"box": coords
|
44 |
-
})
|
45 |
-
|
46 |
-
return processed
|
47 |
-
|
48 |
-
|
49 |
-
|
50 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|