Spaces:
Running
on
Zero
Running
on
Zero
Update app.py
Browse files
app.py
CHANGED
@@ -194,47 +194,54 @@ async def predict_single_dog(image):
|
|
194 |
|
195 |
# return dogs
|
196 |
|
197 |
-
async def detect_multiple_dogs(image, conf_threshold=0.35, iou_threshold=0.55
|
198 |
results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
|
199 |
dogs = []
|
200 |
boxes = []
|
201 |
-
image_area = image.width * image.height
|
202 |
|
203 |
for box in results.boxes:
|
204 |
if box.cls == 16: # COCO dataset class for dog is 16
|
205 |
xyxy = box.xyxy[0].tolist()
|
206 |
confidence = box.conf.item()
|
207 |
-
|
208 |
-
# 計算檢測框面積
|
209 |
-
box_area = (xyxy[2] - xyxy[0]) * (xyxy[3] - xyxy[1])
|
210 |
-
|
211 |
-
# 如果檢測框面積小於圖片面積的一定比例,則忽略
|
212 |
-
if box_area / image_area < min_area_ratio:
|
213 |
-
continue
|
214 |
-
|
215 |
boxes.append((xyxy, confidence))
|
216 |
|
217 |
if not boxes:
|
218 |
dogs.append((image, 1.0, [0, 0, image.width, image.height]))
|
219 |
else:
|
220 |
-
# 使用非最大抑制
|
221 |
nms_boxes = non_max_suppression(boxes, iou_threshold)
|
222 |
|
223 |
-
|
224 |
-
|
225 |
-
|
226 |
-
for box, confidence in merged_boxes:
|
227 |
-
x1, y1, x2, y2 = box
|
228 |
-
w, h = x2 - x1, y2 - y1
|
229 |
-
x1 = max(0, x1 - w * 0.05)
|
230 |
-
y1 = max(0, y1 - h * 0.05)
|
231 |
-
x2 = min(image.width, x2 + w * 0.05)
|
232 |
-
y2 = min(image.height, y2 + h * 0.05)
|
233 |
cropped_image = image.crop((x1, y1, x2, y2))
|
234 |
-
|
|
|
|
|
235 |
|
236 |
return dogs
|
237 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
238 |
def merge_overlapping_boxes(boxes, overlap_threshold):
|
239 |
merged = []
|
240 |
while boxes:
|
@@ -369,7 +376,7 @@ async def predict(image):
|
|
369 |
explanations.append(dog_explanation)
|
370 |
buttons.extend([f"Dog {i+1}: More about {breed}" for breed in topk_breeds[:3]])
|
371 |
else:
|
372 |
-
explanations.append("The image is unclear or the breed is not in the dataset. Please upload a clearer image.")
|
373 |
|
374 |
|
375 |
final_explanation = "\n\n".join(explanations)
|
|
|
194 |
|
195 |
# return dogs
|
196 |
|
197 |
+
async def detect_multiple_dogs(image, conf_threshold=0.35, iou_threshold=0.55):
|
198 |
results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
|
199 |
dogs = []
|
200 |
boxes = []
|
|
|
201 |
|
202 |
for box in results.boxes:
|
203 |
if box.cls == 16: # COCO dataset class for dog is 16
|
204 |
xyxy = box.xyxy[0].tolist()
|
205 |
confidence = box.conf.item()
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
206 |
boxes.append((xyxy, confidence))
|
207 |
|
208 |
if not boxes:
|
209 |
dogs.append((image, 1.0, [0, 0, image.width, image.height]))
|
210 |
else:
|
|
|
211 |
nms_boxes = non_max_suppression(boxes, iou_threshold)
|
212 |
|
213 |
+
for box, confidence in nms_boxes:
|
214 |
+
x1, y1, x2, y2 = [int(coord) for coord in box]
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
215 |
cropped_image = image.crop((x1, y1, x2, y2))
|
216 |
+
|
217 |
+
if is_valid_dog(cropped_image):
|
218 |
+
dogs.append((cropped_image, confidence, [x1, y1, x2, y2]))
|
219 |
|
220 |
return dogs
|
221 |
|
222 |
+
def is_valid_dog(image):
|
223 |
+
# 將PIL Image轉換為numpy陣列
|
224 |
+
img_array = np.array(image)
|
225 |
+
|
226 |
+
# 1. 簡單的紋理檢測
|
227 |
+
gray = np.mean(img_array, axis=2)
|
228 |
+
texture = np.std(gray)
|
229 |
+
|
230 |
+
# 2. 顏色分布檢測
|
231 |
+
img_rgb = img_array.reshape(-1, 3)
|
232 |
+
kmeans = KMeans(n_clusters=3, n_init=10)
|
233 |
+
kmeans.fit(img_rgb)
|
234 |
+
colors = kmeans.cluster_centers_
|
235 |
+
color_variety = np.std(colors)
|
236 |
+
|
237 |
+
# 3. 形狀檢測(簡化版,檢查長寬比)
|
238 |
+
aspect_ratio = image.width / image.height
|
239 |
+
|
240 |
+
# 根據特徵綜合判斷
|
241 |
+
if texture > 10 and color_variety > 30 and 0.5 < aspect_ratio < 2:
|
242 |
+
return True
|
243 |
+
return False
|
244 |
+
|
245 |
def merge_overlapping_boxes(boxes, overlap_threshold):
|
246 |
merged = []
|
247 |
while boxes:
|
|
|
376 |
explanations.append(dog_explanation)
|
377 |
buttons.extend([f"Dog {i+1}: More about {breed}" for breed in topk_breeds[:3]])
|
378 |
else:
|
379 |
+
explanations.append(f"{i+1} The image is unclear or the breed is not in the dataset. Please upload a clearer image.")
|
380 |
|
381 |
|
382 |
final_explanation = "\n\n".join(explanations)
|