Update app.py

app.py

@@ -1,65 +1,148 @@
+# import gradio as gr
+# from transformers import AutoProcessor, AutoModelForImageTextToText
+# from PIL import Image
+# import re
+
+# # Load SmolDocling model & processor once
+# processor = AutoProcessor.from_pretrained("ds4sd/SmolDocling-256M-preview")
+# model = AutoModelForImageTextToText.from_pretrained("ds4sd/SmolDocling-256M-preview")
+
+# def extract_fcel_values_from_image(image, prompt_text):
+#     """Run SmolDocling on an image and return numeric values inside <fcel> tags."""
+#     # Prepare prompt for the model
+#     messages = [
+#         {"role": "user", "content": [{"type": "image"}, {"type": "text", "text": prompt_text}]}
+#     ]
+#     prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
+#     inputs = processor(text=prompt, images=[image], return_tensors="pt")
+
+#     # Generate output
+#     outputs = model.generate(**inputs, max_new_tokens=2048)
+#     prompt_length = inputs.input_ids.shape[1]
+#     generated = outputs[:, prompt_length:]
+#     result = processor.batch_decode(generated, skip_special_tokens=False)[0]
+#     clean_text = result.replace("<end_of_utterance>", "").strip()
+
+#     # Extract only <fcel> values
+#     values = re.findall(r"<fcel>([\d.]+)", clean_text)
+#     values = [float(v) for v in values]  # convert to floats
+
+#     return values, clean_text
+
+# def compare_images(image1, image2, prompt_text):
+#     # Extract fcel values from both images
+#     values1, raw1 = extract_fcel_values_from_image(image1, prompt_text)
+#     values2, raw2 = extract_fcel_values_from_image(image2, prompt_text)
+
+#     # Calculate accuracy
+#     if len(values1) == len(values2) and values1 == values2:
+#         accuracy = 100.0
+#     else:
+#         matches = sum(1 for a, b in zip(values1, values2) if a == b)
+#         total = max(len(values1), len(values2))
+#         accuracy = (matches / total) * 100 if total > 0 else 0
+
+#     return {
+#         # "Extracted Values 1": values1,
+#         # "Extracted Values 2": values2,
+#         "Accuracy (%)": accuracy
+#     }
+
+# # Gradio UI
+# demo = gr.Interface(
+#     fn=compare_images,
+#     inputs=[
+#         gr.Image(type="pil", label="Upload First Table Image"),
+#         gr.Image(type="pil", label="Upload Second Table Image"),
+#         gr.Textbox(lines=1, placeholder="Enter prompt (e.g. Extract table as OTSL)", label="Prompt")
+#     ],
+#     outputs="json",
+#     title="Table Data Accuracy Checker (SmolDocling)",
+#     description="Uploads two table images, extracts only <fcel> values from OTSL output, and compares them for accuracy."
+# )
+
+# demo.launch()
+
 import gradio as gr
-from transformers import AutoProcessor, AutoModelForImageTextToText
-from PIL import Image
+from transformers import AutoProcessor, AutoModelForVision2Seq, TextIteratorStreamer
+from transformers.image_utils import load_image
+from threading import Thread
+import torch
+import html
 import re
+from PIL import Image, ImageOps
 
-# Load SmolDocling model & processor once
+# Load model & processor once at startup
 processor = AutoProcessor.from_pretrained("ds4sd/SmolDocling-256M-preview")
-model = AutoModelForImageTextToText.from_pretrained("ds4sd/SmolDocling-256M-preview")
-
-def extract_fcel_values_from_image(image, prompt_text):
-    """Run SmolDocling on an image and return numeric values inside <fcel> tags."""
-    # Prepare prompt for the model
-    messages = [
-        {"role": "user", "content": [{"type": "image"}, {"type": "text", "text": prompt_text}]}
-    ]
-    prompt = processor.apply_chat_template(messages, add_generation_prompt=True)
-    inputs = processor(text=prompt, images=[image], return_tensors="pt")
-
-    # Generate output
-    outputs = model.generate(**inputs, max_new_tokens=2048)
-    prompt_length = inputs.input_ids.shape[1]
-    generated = outputs[:, prompt_length:]
-    result = processor.batch_decode(generated, skip_special_tokens=False)[0]
-    clean_text = result.replace("<end_of_utterance>", "").strip()
-
-    # Extract only <fcel> values
-    values = re.findall(r"<fcel>([\d.]+)", clean_text)
-    values = [float(v) for v in values]  # convert to floats
-
-    return values, clean_text
-
-def compare_images(image1, image2, prompt_text):
-    # Extract fcel values from both images
-    values1, raw1 = extract_fcel_values_from_image(image1, prompt_text)
-    values2, raw2 = extract_fcel_values_from_image(image2, prompt_text)
-
-    # Calculate accuracy
-    if len(values1) == len(values2) and values1 == values2:
-        accuracy = 100.0
-    else:
-        matches = sum(1 for a, b in zip(values1, values2) if a == b)
-        total = max(len(values1), len(values2))
-        accuracy = (matches / total) * 100 if total > 0 else 0
-
-    return {
-        # "Extracted Values 1": values1,
-        # "Extracted Values 2": values2,
-        "Accuracy (%)": accuracy
+model = AutoModelForVision2Seq.from_pretrained("ds4sd/SmolDocling-256M-preview", torch_dtype=torch.bfloat16).to("cuda")
+
+def add_random_padding(image, min_percent=0.1, max_percent=0.10):
+    image = image.convert("RGB")
+    width, height = image.size
+    pad_w_percent = random.uniform(min_percent, max_percent)
+    pad_h_percent = random.uniform(min_percent, max_percent)
+    pad_w = int(width * pad_w_percent)
+    pad_h = int(height * pad_h_percent)
+    corner_pixel = image.getpixel((0, 0))  # Top-left corner
+    padded_image = ImageOps.expand(image, border=(pad_w, pad_h, pad_w, pad_h), fill=corner_pixel)
+    return padded_image
+
+def extract_table(image_file):
+    # Load image
+    image = load_image(image_file)
+    
+    # Optionally add padding if needed for model robustness (optional)
+    image = add_random_padding(image)
+
+    # Fixed prompt to extract table only (modify if needed)
+    text = "Convert this table to OTSL."
+
+    # Build the message structure for processor
+    resulting_messages = [{
+        "role": "user",
+        "content": [{"type": "image"}] + [{"type": "text", "text": text}]
+    }]
+    
+    prompt = processor.apply_chat_template(resulting_messages, add_generation_prompt=True)
+    inputs = processor(text=prompt, images=[image], return_tensors="pt").to('cuda')
+
+    generation_args = {
+        "input_ids": inputs.input_ids,
+        "pixel_values": inputs.pixel_values,
+        "attention_mask": inputs.attention_mask,
+        "max_new_tokens": 8192,
+        "num_return_sequences": 1,
     }
 
+    streamer = TextIteratorStreamer(processor, skip_prompt=True, skip_special_tokens=True)
+    generation_args["streamer"] = streamer
+
+    thread = Thread(target=model.generate, kwargs=generation_args)
+    thread.start()
+
+    output_text = ""
+    for new_text in streamer:
+        output_text += new_text
+
+    # Clean and return output
+    cleaned_output = output_text.replace("<end_of_utterance>", "").strip()
+    
+    # Optionally convert <chart> tags to <otsl> if present
+    if "<chart>" in cleaned_output:
+        cleaned_output = cleaned_output.replace("<chart>", "<otsl>").replace("</chart>", "</otsl>")
+        cleaned_output = re.sub(r'(<loc_500>)(?!.*<loc_500>)<[^>]+>', r'\1', cleaned_output)
+
+    return cleaned_output or "No table found or unable to extract."
+
 # Gradio UI
 demo = gr.Interface(
-    fn=compare_images,
-    inputs=[
-        gr.Image(type="pil", label="Upload First Table Image"),
-        gr.Image(type="pil", label="Upload Second Table Image"),
-        gr.Textbox(lines=1, placeholder="Enter prompt (e.g. Extract table as OTSL)", label="Prompt")
-    ],
-    outputs="json",
-    title="Table Data Accuracy Checker (SmolDocling)",
-    description="Uploads two table images, extracts only <fcel> values from OTSL output, and compares them for accuracy."
+    fn=extract_table,
+    inputs=gr.Image(type="filepath", label="Upload Table Image"),
+    outputs=gr.Textbox(label="Extracted Table (OTSL Format)"),
+    title="Table Extraction from Image using SmolDocling-256M",
+    description="Upload an image containing a table. The model will extract the table and output it in OTSL format."
 )
 
-demo.launch()
+demo.launch(debug=True)
+