Update app.py

app.py

@@ -28,6 +28,25 @@ class_names = {
     14: ('swimming pool', 'استخر شنا')
 }
 
+# رنگ‌ها برای هر کلاس
+colors = {
+    0: (255, 0, 0),  # Red
+    1: (0, 255, 0),  # Green
+    2: (0, 0, 255),  # Blue
+    3: (255, 255, 0),  # Yellow
+    4: (255, 0, 255),  # Magenta
+    5: (0, 255, 255),  # Cyan
+    6: (128, 0, 128),  # Purple
+    7: (255, 165, 0),  # Orange
+    8: (0, 128, 0),  # Dark Green
+    9: (128, 128, 0),  # Olive
+    10: (128, 0, 0),  # Maroon
+    11: (0, 128, 128),  # Teal
+    12: (0, 0, 128),  # Navy
+    13: (75, 0, 130),  # Indigo
+    14: (199, 21, 133)  # MediumVioletRed
+}
+
 # تابع برای تشخیص اشیاء در تصاویر
 def detect_and_draw_image(input_image):
     # تبدیل تصویر PIL به آرایه NumPy
@@ -37,13 +56,8 @@ def detect_and_draw_image(input_image):
     results = model.predict(source=input_image_np, conf=0.3)
     obb_results = results[0].obb
 
-    # تبدیل تصویر به PIL برای رسم
-    draw = ImageDraw.Draw(input_image)
-    # بارگذاری یک فونت برای نوشتن متن
-    try:
-        font = ImageFont.truetype("arial.ttf", size=15)
-    except IOError:
-        font = ImageFont.load_default()
+    # بارگذاری تصویر اصلی به صورت OpenCV برای رسم جعبه‌ها
+    image_np = np.array(input_image.convert('RGB'))[:, :, ::-1]  # تبدیل PIL به OpenCV
 
     counts = {}
     for obb, conf, cls in zip(obb_results.data.cpu().numpy(), obb_results.conf.cpu().numpy(), obb_results.cls.cpu().numpy()):
@@ -60,9 +74,14 @@ def detect_and_draw_image(input_image):
         rect = ((x_center, y_center), (width, height), rotation * 180.0 / np.pi)  # تبدیل رادیان به درجه
         box = cv2.boxPoints(rect)
         box = np.int0(box)
-        draw.polygon([tuple(point) for point in box], outline="red")
-        # نوشتن برچسب و اطمینان
-        draw.text((x_center, y_center - 10), f"{label_en}: {confidence:.2f}", fill="red", font=font)
+        color = colors.get(class_id, (0, 255, 0))  # استفاده از رنگ مشخص برای هر کلاس
+        cv2.drawContours(image_np, [box], 0, color, 2)
+        cv2.putText(image_np, f'{label_en}: {confidence:.2f}', (int(x_center), int(y_center)), 
+                    cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 1, cv2.LINE_AA)
+
+    # تبدیل تصویر به RGB برای Gradio
+    image_rgb = cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB)
+    output_image = Image.fromarray(image_rgb)
 
     # ایجاد DataFrame برای نمایش نتایج
     df = pd.DataFrame({
@@ -71,7 +90,7 @@ def detect_and_draw_image(input_image):
         'Object Count': list(counts.values())
     })
 
-    return input_image, df
+    return output_image, df
 
 # رابط کاربری تصویر
 image_interface = gr.Interface(
@@ -84,4 +103,4 @@ image_interface = gr.Interface(
 )
 
 # اجرای برنامه با استفاده از رابط کاربری
-image_interface.launch(debug=True, share=True)
+image_interface.launch(debug=True, share=True)