Update app.py

app.py

@@ -12,33 +12,75 @@ def format_label(label):
   """
   return label[label.find(" ")+1:-1]
 
-def check_species_status(species_name):
-  status = ''
-
-  return status
-
 def info(species_name):
   status = check_species_status(species_name)
   if status == '':
     info = ''
   return info
 
-def get_vi_name(en_name):
+def get_name(en_name, lan):
   """
   Return name in Vietnamese
   """
-  print(en_name)
-  return data[data['en_name'] == en_name]['vi_name'].to_list()[0]
+  return data[data[f'name_en'] == en_name][f'name_{lan}'].to_list()[0]
+
+def get_fun_fact(en_name, lan):
+  """
+  Return fun fact of the species
+  """
+  return data[data[f'name_en'] == en_name][f'fun_fact_{lan}'].to_list()[0]
 
 def get_law(en_name):
-  cites = data[data['en_name'] == en_name]['CITES'].to_list()[0]
-  nd06 = data[data['en_name'] == en_name]['ND06'].to_list()[0]
+  cites = data[data['name_en'] == en_name]['CITES'].to_list()[0]
+  nd06 = data[data['name_en'] == en_name]['ND06'].to_list()[0]
   return cites, nd06
 
-def get_habitat(en_name):
-  return data[data['en_name'] == en_name]['habitat'].to_list()[0]
+def get_habitat(en_name, lan):
+  return data[data['name_en'] == en_name][f'habitat_{lan}'].to_list()[0]
+
+def get_conservation_status(en_name, lan):
+  status = data[data['name_en'] == en_name]['conservation_status'].to_list()[0]
+  status_dict = {
+      'NE': {
+          'vi': "Không được đánh giá",
+          'en': "Not Evaluated"
+      },
+      'DD': {
+          'vi': "Thiếu dữ liệu",
+          'en': "Data Deficient"
+      },
+      'LC': {
+          'vi': "Ít quan tâm",
+          'en': "Least Concern"
+      },
+      'NT': {
+          'vi': "Sắp bị đe dọa",
+          'en': "Near Threatened"
+      },
+      'VU': {
+          'vi': "Sắp nguy cấp",
+          'en': "Vulnerable"
+      },
+      'EN': {
+          'vi': "Nguy cấp",
+          'en': "Endangered"
+      },
+      'CR': {
+          'vi': "Cực kỳ nguy cấp",
+          'en': "Critically Endangered"
+      },
+      'EW': {
+          'vi': "Tuyệt chủng trong tự nhiên",
+          'en': "Extinct in the Wild"
+      },
+      'EX': {
+          'vi': "Tuyệt chủng",
+          'en': "Extinct"
+      },
+  }
+  return status_dict[status][lan]
   
-def predict(image):
+def predict(image, language):
 
   # Load the model
   model = load_model('keras_model.h5')
@@ -64,16 +106,54 @@ def predict(image):
   pred = model.predict(data)
   pred = pred.tolist()
 
+  sorted_index = np.argsort(pred).tolist()[0]
+
+  print(sorted_index)
+
+  display_index = []
+  for i in sorted_index[::-1]:
+    if pred[0][i] > 0.01:
+      display_index.append(i)
+
   with open('labels.txt','r') as f:
     labels = f.readlines()
   
-  en_name = format_label(labels[pred.index(max(pred))]).strip()
+  en_name = format_label(labels[sorted_index[-1]]).strip()
+
+  lan = ''
+  if language == "Tiếng Việt":
+    lan = 'vi'
+  else:
+    lan = 'en'
 
-  result = {get_vi_name(format_label(labels[i])): round(pred[0][i],2) for i in range(len(pred[0]))}
+  result = {get_name(format_label(labels[i]), lan): round(pred[0][i],2) for i in display_index}
   cites, nd06 = get_law(en_name)
   info = ""
-  if str(nd06) != "":
-    info += f'CITES: {cites}, NĐ06: {nd06} \n \n'
+
+  fun_fact = ""
+  if lan == 'vi':
+    fun_fact += "💡 **Thông tin lý thú** \n\n"
+  if lan == 'en':
+    fun_fact += "💡 **Fun fact** \n\n"
+
+  fun_fact += f"{get_fun_fact(en_name, lan)}."
+
+  status = ""
+  if lan == 'vi':
+    status += "**🐢 Tình trạng bảo tồn** \n\n"
+  if lan == 'en':
+    status += "**🐢 Conservation status** \n\n"
+
+  status += f"{get_conservation_status(en_name, lan)}"
+  
+  law = ""
+  if lan == 'vi':
+    law += "**👮‍♀️ Luật** \n\n"
+  if lan == 'en':
+    law += "**👮‍♀️ Law** \n\n"
+  law += f'CITES: {cites}, NĐ06: {nd06}'
+
+  if lan == 'vi' and str(nd06) != "":
     info += "Đây là loài được pháp luật bảo vệ. Mọi hành vi buôn bán, nuôi nhốt không có \
     [giấy phép](https://thuvienphapluat.vn/van-ban/Tai-nguyen-Moi-truong/Nghi-dinh-06-2019-ND-CP-quan-ly-thuc-vat-rung-dong-vat-rung-nguy-cap-quy-hiem-405883.aspx) đều vi phạm pháp luật \n"
     info += "- Nếu bạn vô tình bắt gặp loài này bị buôn bán mà không có giấy phép, \
@@ -82,21 +162,30 @@ def predict(image):
     info += "- Nếu bạn đang nuôi thì nên giao nộp cho cơ quan chức năng để trả về tự nhiên. Tham khảo đơn vị tiếp nhận DVHD ở địa phương \
     bạn tại [đây](https://drive.google.com/file/d/1K2ZWcHKGEsNudh_LtHgHJOXlVw-GQ6AZ/view). \n"
     info += f"- Nếu bạn bắt gặp trong vườn nhà thì có thể xem xét thả chúng về môi trường sống. Hãy đảm bảo nơi bạn thả là\
-     **{get_habitat(en_name)}**."
+     **{get_habitat(en_name, lan)}**. \n\n"
+    info += f"Liên hệ Chương trình Bảo tồn Rùa châu Á ([ATP](https://asianturtleprogram.org/vi/)) qua email: info@asianturtleprogram.org hoặc [Facebook](https://www.facebook.com/search/top?q=asian%20turtle%20program) để được hướng dẫn trong các trường hợp cụ thể"
+
+  if lan == 'en' and str(nd06) != "":
+    info += f'CITES: {cites}, NĐ06: {nd06} \n \n'
+    info += "This species is protected by law. All acts of trafficking and captive breeding without \
+    [authority permisison](https://thuvienphapluat.vn/van-ban/Tai-nguyen-Moi-truong/Nghi-dinh-06-2019-ND-CP-quan-ly-thuc-vat-rung-dong-vat-rung-nguy-cap-quy-hiem-405883.aspx) are illegal \n"
+    info += "- If you happen to come across this species being traded without a license, \
+    do not buy for any purpose (e.g., to release). \
+    Instead, you should report the suspicious act to ENV via their hotline **1800-1522**. \n"
+    info += "- If you are raising them, you should hand them over to the authorities to return to the wild. Find your local wild-animal receiver \
+    [here](https://drive.google.com/file/d/1K2ZWcHKGEsNudh_LtHgHJOXlVw-GQ6AZ/view). \n"
+    info += f"- If you find them in your garden, consider releasing them back into their habitat. Please make sure the place you release is \
+     **{get_habitat(en_name, lan)}**. \n\n"
+    info += f"Contact Asian Turtle Program ([ATP](https://asianturtleprogram.org/)) via email (info@asianturtleprogram.org) or [Facebook](https://www.facebook.com/search/top?q=asian%20turtle%20program) for instructions in specific cases."
   
-  return result, info
+  return result, fun_fact, status, law, info
   
   
-description="""
-VNTurtle nhận diện các loài rùa Việt Nam. Mô hình mẫu này có thể nhận diện 10 loại rùa thường xuất hiện ở VN gồm **5** loài bản địa
-**(1) Rùa núi viền**, **(2) Rùa núi vàng**, **(3) Rùa ba gờ**, **(4) Rùa răng**, và **(5) Rùa sa nhân**, 
-và **5** loài ngoại lai **(1) Rùa Sulcata**, **(2) Rùa bản đồ**, **(3) Rùa cá sấu**, **(4) Rùa tai đỏ**, và **(5) Rùa ninja**
-"""
 article = "© Hình ảnh minh hoạ được cung cấp bởi ATP"
 examples = [ ['test1.jpg'], ['test2.jpg'], ['test3.jpg']  ]
 gr.Interface(fn=predict, 
-             inputs=gr.Image(type="pil", label="Input Image"),
-             outputs=[gr.Label(label="Predictions"), gr.Markdown()],
+             inputs=[gr.Image(type="pil", label="Input Image"), gr.Radio(value="Tiếng Việt", choices=["Tiếng Việt", "English"], label="Language")],
+             outputs=[gr.Label(label="Predictions"), gr.Markdown(), gr.Markdown(), gr.Markdown(), gr.Markdown()],
              live=True,
              title='VNTurtle - Toy Model',
              description=description,