hank1229 commited on
Commit
f750671
·
verified ·
1 Parent(s): 390638a

Delete vision.py

Browse files
Files changed (1) hide show
  1. vision.py +0 -50
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
-