Update app.py

app.py

@@ -94,50 +94,40 @@ categories = labels + ["all"]
 # function to get a rendom image and all data from DocLayNet
 def generate_annotated_image(dataset_name, split, domain, category):
 
-  def get_dataset(dataset_name, split, domain, category):
-
-    # error message
-    msg_error = ""
-
-    # get dataset
-    if dataset_name == "small": example = dataset_small
-    else: example = dataset_base
-
-    # get split
-    if split == "all":
-      example = concatenate_datasets([example["train"], example["validation"], example["test"]])
-    else:
-      example = example[split]
-
-    # get domain
-    domain_name = domains_names[domains.index(domain)]
-    if domain_name != "all":
-      example = example.filter(lambda example: example["doc_category"] == domain_name)
-      if len(example) == 0:
-        msg_error = f'There is no image with at least one annotated bounding box that matches to your parameters ("{domain}" domain / "DocLayNet {dataset_name}" dataset splitted into "{split}").'
-        example = dict()
-        return example, msg_error
-
-    # get category
-    idx_list = list()
-    if category != "all":
-      for idx, categories_list in zip(example["id"], example["categories"]):
-        if category in categories_list:
-          idx_list.append(idx)
-      example = example.select(idx_list)
-      if len(example) == 0:
-        msg_error = f'There is no image with at least one annotated bounding box that matches to your parameters (category: "{category}" / domain: "{domain}" / dataset: "DocLayNet {dataset_name}" / split: "{split}").'
-        example = dict()
-        return example, msg_error
-
-    return example, msg_error
-
-  # get results
-  example, msg_error = get_dataset(dataset_name, split, domain, category)
+  # error message
+  msg_error = ""
+
+  # get dataset
+  if dataset_name == "small": example = dataset_small
+  else: example = dataset_base
+
+  # get split
+  if split == "all":
+    example = concatenate_datasets([example["train"], example["validation"], example["test"]])
+  else:
+    example = example[split]
+
+  # get domain
+  domain_name = domains_names[domains.index(domain)]
+  if domain_name != "all":
+    example = example.filter(lambda example: example["doc_category"] == domain_name)
+    if len(example) == 0:
+      msg_error = f'There is no image with at least one annotated bounding box that matches to your parameters (domain: "{domain}" / dataset: "DocLayNet {dataset_name}" / split: "{split}").'
+      example = dict()
+
+  # get category
+  idx_list = list()
+  if category != "all":
+    for idx, categories_list in zip(example["id"], example["categories"]):
+      if category in categories_list:
+        idx_list.append(idx)
+    example = example.select(idx_list)
+    if len(example) == 0:
+      msg_error = f'There is no image with at least one annotated bounding box that matches to your parameters (category: "{category}" / domain: "{domain}" / dataset: "DocLayNet {dataset_name}" / split: "{split}").'
+      example = dict()
 
   if len(msg_error) > 0:
     return msg_error, images_wo_content, images_wo_content, df_paragraphs_wo_content, df_lines_wo_content
-
   else:
     # get random image & PDF data
     image_files = example["image"]
@@ -310,4 +300,4 @@ with gr.Blocks(title="DocLayNet image viewer", css=".gradio-container") as demo:
         cache_examples=True,
     )
 
-demo.launch(share=True)
+demo.launch()