DawnC commited on
Commit
2db8b4f
1 Parent(s): 729f9dc

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +121 -29
app.py CHANGED
@@ -13,7 +13,7 @@ os.system('pip install ultralytics')
13
  from ultralytics import YOLO
14
 
15
  # 下載YOLOv5預訓練模型
16
- model_yolo = YOLO('yolov5s.pt') # 使用 YOLOv5 預訓練模型
17
 
18
 
19
  dog_breeds = ["Afghan_Hound", "African_Hunting_Dog", "Airedale", "American_Staffordshire_Terrier",
@@ -166,6 +166,103 @@ def get_akc_breeds_link():
166
  # return f"An error occurred: {e}", gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
167
 
168
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
169
 
170
  def predict(image):
171
  if image is None:
@@ -177,22 +274,16 @@ def predict(image):
177
  image = Image.fromarray(image)
178
 
179
  # 使用 YOLO 偵測狗
180
- results = model_yolo(image)
181
- boxes = results[0].boxes # 提取 YOLO 偵測結果中的邊界框
182
-
183
- if len(boxes) == 0:
184
- return "The image is too unclear or the dog breed is not in the dataset. Please upload a clearer image of the dog.", gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
185
-
186
- # 檢查 YOLO 偵測到的邊界框���量
187
- print(f"Detected {len(boxes)} dogs in the image.")
188
 
 
189
  explanations = []
190
  visible_buttons = []
 
191
 
192
- # 處理每個偵測到的狗
193
- for i, box in enumerate(boxes):
194
- x1, y1, x2, y2 = map(int, box.xyxy[0])
195
- cropped_image = image.crop((x1, y1, x2, y2)) # 裁剪出每隻狗的區域
196
  image_tensor = preprocess_image(cropped_image)
197
 
198
  with torch.no_grad():
@@ -206,32 +297,36 @@ def predict(image):
206
  topk_breeds = [dog_breeds[idx.item()] for idx in topk_indices[0]]
207
  topk_probs_percent = [f"{prob.item() * 100:.2f}%" for prob in topk_probs[0]]
208
 
209
- # 根據置信度決定返回結果
 
 
 
 
 
210
  if top1_prob >= 0.5:
211
  breed = topk_breeds[0]
212
  description = get_dog_description(breed)
213
- explanations.append(f"Detected a dog: **{breed}** with {topk_probs_percent[0]} confidence.\n{format_description(description, breed)}")
 
214
  elif 0.2 <= top1_prob < 0.5:
215
  explanation = (
216
- f"Detected a dog with moderate confidence. Here are the top 3 possible breeds:\n"
217
- f"1. **{topk_breeds[0]}** ({topk_probs_percent[0]} confidence)\n"
218
- f"2. **{topk_breeds[1]}** ({topk_probs_percent[1]} confidence)\n"
219
- f"3. **{topk_breeds[2]}** ({topk_probs_percent[2]} confidence)\n"
220
  )
221
  explanations.append(explanation)
222
  visible_buttons.extend([f"More about {topk_breeds[0]}", f"More about {topk_breeds[1]}", f"More about {topk_breeds[2]}"])
223
  else:
224
- explanations.append("The image is too unclear or the breed is not in the dataset. Please upload a clearer image.")
225
 
226
- # 將結果匯總後返回
227
  final_explanation = "\n\n".join(explanations)
228
- return final_explanation, gr.update(visible=len(visible_buttons) >= 1), gr.update(visible=len(visible_buttons) >= 2), gr.update(visible=len(visible_buttons) >= 3)
229
 
230
  except Exception as e:
231
  return f"An error occurred: {e}", gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
232
 
233
 
234
-
235
  def format_description(description, breed):
236
  if isinstance(description, dict):
237
  formatted_description = "\n\n".join([f"**{key}**: {value}" for key, value in description.items()])
@@ -249,11 +344,6 @@ def format_description(description, breed):
249
 
250
  return formatted_description
251
 
252
- def show_details(breed):
253
- breed_name = breed.split("More about ")[-1]
254
- description = get_dog_description(breed_name)
255
- return format_description(description, breed_name)
256
-
257
  with gr.Blocks(css="""
258
  .container {
259
  max-width: 900px;
@@ -290,6 +380,7 @@ with gr.Blocks(css="""
290
 
291
  with gr.Row():
292
  input_image = gr.Image(label="Upload a dog image", type="numpy")
 
293
  output = gr.Markdown(label="Prediction Results")
294
 
295
  with gr.Row():
@@ -297,7 +388,7 @@ with gr.Blocks(css="""
297
  btn2 = gr.Button("View More 2", visible=False)
298
  btn3 = gr.Button("View More 3", visible=False)
299
 
300
- input_image.change(predict, inputs=input_image, outputs=[output, btn1, btn2, btn3])
301
 
302
  btn1.click(show_details, inputs=btn1, outputs=output)
303
  btn2.click(show_details, inputs=btn2, outputs=output)
@@ -314,3 +405,4 @@ with gr.Blocks(css="""
314
  if __name__ == "__main__":
315
  iface.launch()
316
 
 
 
13
  from ultralytics import YOLO
14
 
15
  # 下載YOLOv5預訓練模型
16
+ model_yolo = YOLO('yolov8n.pt') # 使用 YOLOv5 預訓練模型
17
 
18
 
19
  dog_breeds = ["Afghan_Hound", "African_Hunting_Dog", "Airedale", "American_Staffordshire_Terrier",
 
166
  # return f"An error occurred: {e}", gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
167
 
168
 
169
+ # def format_description(description, breed):
170
+ # if isinstance(description, dict):
171
+ # formatted_description = "\n\n".join([f"**{key}**: {value}" for key, value in description.items()])
172
+ # else:
173
+ # formatted_description = description
174
+
175
+ # akc_link = get_akc_breeds_link()
176
+ # formatted_description += f"\n\n**Want to learn more about dog breeds?** [Visit the AKC dog breeds page]({akc_link}) and search for {breed} to find detailed information."
177
+
178
+ # disclaimer = ("\n\n*Disclaimer: The external link provided leads to the American Kennel Club (AKC) dog breeds page. "
179
+ # "You may need to search for the specific breed on that page. "
180
+ # "I am not responsible for the content on external sites. "
181
+ # "Please refer to the AKC's terms of use and privacy policy.*")
182
+ # formatted_description += disclaimer
183
+
184
+ # return formatted_description
185
+
186
+ # def show_details(breed):
187
+ # breed_name = breed.split("More about ")[-1]
188
+ # description = get_dog_description(breed_name)
189
+ # return format_description(description, breed_name)
190
+
191
+ # with gr.Blocks(css="""
192
+ # .container {
193
+ # max-width: 900px;
194
+ # margin: 0 auto;
195
+ # padding: 20px;
196
+ # background-color: rgba(255, 255, 255, 0.9);
197
+ # border-radius: 15px;
198
+ # box-shadow: 0 0 20px rgba(0, 0, 0, 0.1);
199
+ # }
200
+ # .gr-form { display: flex; flex-direction: column; align-items: center; }
201
+ # .gr-box { width: 100%; max-width: 500px; }
202
+ # .output-markdown, .output-image {
203
+ # margin-top: 20px;
204
+ # padding: 15px;
205
+ # background-color: #f5f5f5;
206
+ # border-radius: 10px;
207
+ # }
208
+ # .examples {
209
+ # display: flex;
210
+ # justify-content: center;
211
+ # flex-wrap: wrap;
212
+ # gap: 10px;
213
+ # margin-top: 20px;
214
+ # }
215
+ # .examples img {
216
+ # width: 100px;
217
+ # height: 100px;
218
+ # object-fit: cover;
219
+ # }
220
+ # """) as iface:
221
+
222
+ # gr.HTML("<h1 style='font-family:Roboto; font-weight:bold; color:#2C3E50; text-align:center;'>🐶 Dog Breed Classifier 🔍</h1>")
223
+ # gr.HTML("<p style='font-family:Open Sans; color:#34495E; text-align:center;'>Upload a picture of a dog, and the model will predict its breed, provide detailed information, and include an extra information link!</p>")
224
+
225
+ # with gr.Row():
226
+ # input_image = gr.Image(label="Upload a dog image", type="numpy")
227
+ # output = gr.Markdown(label="Prediction Results")
228
+
229
+ # with gr.Row():
230
+ # btn1 = gr.Button("View More 1", visible=False)
231
+ # btn2 = gr.Button("View More 2", visible=False)
232
+ # btn3 = gr.Button("View More 3", visible=False)
233
+
234
+ # input_image.change(predict, inputs=input_image, outputs=[output, btn1, btn2, btn3])
235
+
236
+ # btn1.click(show_details, inputs=btn1, outputs=output)
237
+ # btn2.click(show_details, inputs=btn2, outputs=output)
238
+ # btn3.click(show_details, inputs=btn3, outputs=output)
239
+
240
+ # gr.Examples(
241
+ # examples=['Border_Collie.jpg', 'Golden_Retriever.jpeg', 'Saint_Bernard.jpeg', 'French_Bulldog.jpeg', 'Samoyed.jpg'],
242
+ # inputs=input_image
243
+ # )
244
+
245
+ # gr.HTML('For more details on this project and other work, feel free to visit my GitHub <a href="https://github.com/Eric-Chung-0511/Learning-Record/tree/main/Data%20Science%20Projects/Dog%20Breed%20Classifier">Dog Breed Classifier</a>')
246
+
247
+ # # launch the program
248
+ # if __name__ == "__main__":
249
+ # iface.launch()
250
+
251
+
252
+ # 使用 YOLOv8 進行狗偵測
253
+ def detect_dogs(image):
254
+ results = yolo_model.predict(image)
255
+ dogs = []
256
+
257
+ for result in results:
258
+ for box in result.boxes:
259
+ if box.cls == 16: # COCO 資料集中的狗類別是16
260
+ xyxy = box.xyxy
261
+ confidence = box.conf
262
+ cropped_image = image.crop((xyxy[0], xyxy[1], xyxy[2], xyxy[3]))
263
+ dogs.append((cropped_image, confidence, xyxy))
264
+
265
+ return dogs
266
 
267
  def predict(image):
268
  if image is None:
 
274
  image = Image.fromarray(image)
275
 
276
  # 使用 YOLO 偵測狗
277
+ dogs = detect_dogs(image)
278
+ if len(dogs) == 0:
279
+ return "No dogs detected or the image is too unclear. Please upload a clearer image.", gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
 
 
 
 
 
280
 
281
+ # 開始處理每一隻狗
282
  explanations = []
283
  visible_buttons = []
284
+ annotated_image = image.copy()
285
 
286
+ for i, (cropped_image, confidence, box) in enumerate(dogs):
 
 
 
287
  image_tensor = preprocess_image(cropped_image)
288
 
289
  with torch.no_grad():
 
297
  topk_breeds = [dog_breeds[idx.item()] for idx in topk_indices[0]]
298
  topk_probs_percent = [f"{prob.item() * 100:.2f}%" for prob in topk_probs[0]]
299
 
300
+ # 標註狗的邊界框
301
+ draw = ImageDraw.Draw(annotated_image)
302
+ draw.rectangle(box.tolist(), outline="red", width=3)
303
+ draw.text((box[0], box[1]), f"Dog {i+1}", fill="red")
304
+
305
+ # 信心度大於 50%,顯示詳細品種資訊
306
  if top1_prob >= 0.5:
307
  breed = topk_breeds[0]
308
  description = get_dog_description(breed)
309
+ explanations.append(f"Dog {i+1}: **{breed}**\n{format_description(description, breed)}")
310
+ # 信心度 20%-49%,顯示 Top 3 品種
311
  elif 0.2 <= top1_prob < 0.5:
312
  explanation = (
313
+ f"Dog {i+1}: Detected with moderate confidence. Here are the top 3 possible breeds:\n"
314
+ f"1. **{topk_breeds[0]}** ({topk_probs_percent[0]})\n"
315
+ f"2. **{topk_breeds[1]}** ({topk_probs_percent[1]})\n"
316
+ f"3. **{topk_breeds[2]}** ({topk_probs_percent[2]})\n"
317
  )
318
  explanations.append(explanation)
319
  visible_buttons.extend([f"More about {topk_breeds[0]}", f"More about {topk_breeds[1]}", f"More about {topk_breeds[2]}"])
320
  else:
321
+ explanations.append(f"Dog {i+1}: The image is unclear or the breed is not in the dataset.")
322
 
 
323
  final_explanation = "\n\n".join(explanations)
324
+ return annotated_image, final_explanation, gr.update(visible=len(visible_buttons) >= 1), gr.update(visible=len(visible_buttons) >= 2), gr.update(visible=len(visible_buttons) >= 3)
325
 
326
  except Exception as e:
327
  return f"An error occurred: {e}", gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
328
 
329
 
 
330
  def format_description(description, breed):
331
  if isinstance(description, dict):
332
  formatted_description = "\n\n".join([f"**{key}**: {value}" for key, value in description.items()])
 
344
 
345
  return formatted_description
346
 
 
 
 
 
 
347
  with gr.Blocks(css="""
348
  .container {
349
  max-width: 900px;
 
380
 
381
  with gr.Row():
382
  input_image = gr.Image(label="Upload a dog image", type="numpy")
383
+ output_image = gr.Image(label="Annotated Image")
384
  output = gr.Markdown(label="Prediction Results")
385
 
386
  with gr.Row():
 
388
  btn2 = gr.Button("View More 2", visible=False)
389
  btn3 = gr.Button("View More 3", visible=False)
390
 
391
+ input_image.change(predict, inputs=input_image, outputs=[output_image, output, btn1, btn2, btn3])
392
 
393
  btn1.click(show_details, inputs=btn1, outputs=output)
394
  btn2.click(show_details, inputs=btn2, outputs=output)
 
405
  if __name__ == "__main__":
406
  iface.launch()
407
 
408
+