ek-5 commited on
Commit
aeb7d8a
·
verified ·
1 Parent(s): 0cff19b

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +53 -36
app.py CHANGED
@@ -1,4 +1,4 @@
1
- import os
2
  import torch
3
  import io
4
  from fastapi import FastAPI, File, UploadFile
@@ -7,31 +7,41 @@ from ultralytics import YOLO
7
  from PIL import Image
8
  import uvicorn
9
 
10
- app = FastAPI(title="YOLO + GIT Large: Visual Analysis (Color & Shape)")
 
11
 
12
  device = "cuda" if torch.cuda.is_available() else "cpu"
13
  MY_MODEL_PATH = 'best.pt'
14
 
15
- # تحميل الموديلات
 
 
16
  try:
17
  detection_model = YOLO(MY_MODEL_PATH)
18
- print("✅ YOLO Model Loaded")
19
- except:
 
20
  detection_model = YOLO("yolov8n.pt")
21
- print("⚠️ Using Default YOLOv8n")
22
 
23
- processor = AutoProcessor.from_pretrained("microsoft/git-large")
24
- caption_model = AutoModelForCausalLM.from_pretrained("microsoft/git-large").to(device)
 
 
 
25
 
26
  @app.get("/")
27
  def home():
28
- return {"message": "Server is running. Use /docs to test /analyze endpoint."}
 
 
29
 
30
  @app.post("/analyze")
31
  async def analyze_image(file: UploadFile = File(...)):
 
32
  data = await file.read()
33
  original_image = Image.open(io.BytesIO(data)).convert("RGB")
34
 
 
35
  results = detection_model(original_image, conf=0.25)
36
  integrated_results = []
37
 
@@ -40,46 +50,53 @@ async def analyze_image(file: UploadFile = File(...)):
40
  label = r.names[int(box.cls)]
41
  coords = box.xyxy[0].tolist()
42
 
43
- # قص العنصر مع هامش (Padding) 15 بكسل لرؤية الزوايا والأطراف بدقة
44
- pad = 15
45
- cropped_img = original_image.crop((
46
- max(0, coords[0]-pad), max(0, coords[1]-pad),
47
- min(original_image.width, coords[2]+pad), min(original_image.height, coords[3]+pad)
48
- ))
49
-
50
- # --- التعديل هنا: برومبت يركز على الصفات البصرية ---
51
- # بدأنا الجملة بصفات "اللون والشكل" ليقوم الموديل بإكمال الوصف
52
- prompt = f"a photo of a {label}. the specific color and shape of this {label} are"
53
 
54
- inputs = processor(images=cropped_img, text=prompt, return_tensors="pt").to(device)
 
 
 
 
55
 
56
  generated_ids = caption_model.generate(
57
  pixel_values=inputs.pixel_values,
58
- input_ids=inputs.input_ids,
59
- max_new_tokens=40, # عدد كلمات كافٍ للوصف
60
- num_beams=5,
61
- repetition_penalty=1.3,
62
- do_sample=False # نستخدم Beam Search هنا لدقة أعلى في الألوان
63
  )
64
 
65
- full_desc = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
66
-
67
- # تنظيف النتيجة لاستخراج الوصف فقط بعد البرومبت
68
- if prompt in full_desc:
69
- visual_details = full_desc.split(prompt)[-1].strip()
70
- else:
71
- visual_details = full_desc.replace(f"a photo of a {label}", "").strip()
72
 
73
  integrated_results.append({
74
  "object_id": i + 1,
75
  "label": label,
76
- "visual_description": f"The {label} has {visual_details}"
 
77
  })
78
 
 
79
  if not integrated_results:
80
- return {"message": "No objects detected."}
 
 
 
 
 
 
81
 
82
- return {"results": integrated_results}
 
 
 
83
 
 
84
  if __name__ == "__main__":
85
- uvicorn.run(app, host="0.0.0.0", port=7860)
 
1
+ mport os
2
  import torch
3
  import io
4
  from fastapi import FastAPI, File, UploadFile
 
7
  from PIL import Image
8
  import uvicorn
9
 
10
+ # --- 1. إعداد التطبيق والموديلات ---
11
+ app = FastAPI(title="YOLO + GIT Large: Final Visual Description API")
12
 
13
  device = "cuda" if torch.cuda.is_available() else "cpu"
14
  MY_MODEL_PATH = 'best.pt'
15
 
16
+ print(f"🔄 جاري التحميل على جهاز: {device}...")
17
+
18
+ # تحميل YOLO الخاص بكِ
19
  try:
20
  detection_model = YOLO(MY_MODEL_PATH)
21
+ print("✅ YOLO Model: Loaded successfully")
22
+ except Exception as e:
23
+ print(f"⚠️ YOLO Warning: Using default yolov8n.pt - {e}")
24
  detection_model = YOLO("yolov8n.pt")
 
25
 
26
+ # تحميل موديل الوصف GIT-Large
27
+ model_name = "microsoft/git-large"
28
+ processor = AutoProcessor.from_pretrained(model_name)
29
+ caption_model = AutoModelForCausalLM.from_pretrained(model_name).to(device)
30
+ print(f"✅ Caption Model: {model_name} Loaded")
31
 
32
  @app.get("/")
33
  def home():
34
+ return {"status": "Online", "instruction": "Use /docs to test the /analyze endpoint"}
35
+
36
+ # --- 2. وظيفة المعالجة والتحليل ---
37
 
38
  @app.post("/analyze")
39
  async def analyze_image(file: UploadFile = File(...)):
40
+ # قراءة الصورة
41
  data = await file.read()
42
  original_image = Image.open(io.BytesIO(data)).convert("RGB")
43
 
44
+ # كشف الأجسام باستخدام YOLO
45
  results = detection_model(original_image, conf=0.25)
46
  integrated_results = []
47
 
 
50
  label = r.names[int(box.cls)]
51
  coords = box.xyxy[0].tolist()
52
 
53
+ # قص العنصر مع هامش (Padding) 20 بكسل لرؤية الشكل واللون بوضوح
54
+ pad = 20
55
+ left = max(0, coords[0] - pad)
56
+ top = max(0, coords[1] - pad)
57
+ right = min(original_image.width, coords[2] + pad)
58
+ bottom = min(original_image.height, coords[3] + pad)
 
 
 
 
59
 
60
+ cropped_img = original_image.crop((left, top, right, bottom))
61
+
62
+ # --- استراتيجية الوصف الحر (بدون برومبت نصي مقيد) ---
63
+ # نترك الموديل يحلل الصورة بصرياً فقط
64
+ inputs = processor(images=cropped_img, return_tensors="pt").to(device)
65
 
66
  generated_ids = caption_model.generate(
67
  pixel_values=inputs.pixel_values,
68
+ max_length=60, # طول كافٍ لوصف اللون والشكل
69
+ min_length=12, # إجبار الموديل على التفصيل وعدم الاختصار
70
+ num_beams=5, # جودة عالية في اختيار الكلمات
71
+ repetition_penalty=1.5,
72
+ early_stopping=True
73
  )
74
 
75
+ # فك التشفير للوصف الناتج
76
+ description = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
 
 
 
 
 
77
 
78
  integrated_results.append({
79
  "object_id": i + 1,
80
  "label": label,
81
+ "confidence": f"{float(box.conf[0]):.2f}",
82
+ "visual_description": f"Detected {label}: {description.strip()}"
83
  })
84
 
85
+ # في حال لم يتم كشف أي شيء
86
  if not integrated_results:
87
+ inputs = processor(images=original_image, return_tensors="pt").to(device)
88
+ out = caption_model.generate(pixel_values=inputs.pixel_values, max_length=50)
89
+ general_desc = processor.batch_decode(out, skip_special_tokens=True)[0]
90
+ return {
91
+ "message": "No specific objects detected by YOLO.",
92
+ "general_scene_description": general_desc
93
+ }
94
 
95
+ return {
96
+ "detected_count": len(integrated_results),
97
+ "results": integrated_results
98
+ }
99
 
100
+ # --- 3. تشغيل السيرفر ---
101
  if __name__ == "__main__":
102
+ uvicorn.run(app, host="0.0.0.0", port=7860)