Create app.py

app.py

@@ -0,0 +1,344 @@
+import gradio as gr
+from transformers.image_utils import load_image
+from threading import Thread
+import time
+import torch
+import spaces
+import cv2
+import numpy as np
+from PIL import Image
+import re
+from transformers import (
+    Qwen2VLForConditionalGeneration,
+    AutoProcessor,
+    TextIteratorStreamer,
+)
+from transformers import Qwen2_5_VLForConditionalGeneration
+
+# ---------------------------
+# Helper Functions
+# ---------------------------
+def progress_bar_html(label: str, primary_color: str = "#4B0082", secondary_color: str = "#9370DB") -> str:
+    """
+    Returns an HTML snippet for a thin animated progress bar with a label.
+    Colors can be customized; default colors are used for Qwen2VL/Aya‑Vision.
+    """
+    return f'''
+<div style="display: flex; align-items: center;">
+    <span style="margin-right: 10px; font-size: 14px;">{label}</span>
+    <div style="width: 110px; height: 5px; background-color: {secondary_color}; border-radius: 2px; overflow: hidden;">
+        <div style="width: 100%; height: 100%; background-color: {primary_color}; animation: loading 1.5s linear infinite;"></div>
+    </div>
+</div>
+<style>
+@keyframes loading {{
+    0% {{ transform: translateX(-100%); }}
+    100% {{ transform: translateX(100%); }}
+}}
+</style>
+    '''
+
+def downsample_video(video_path):
+    """
+    Downsamples a video file by extracting 10 evenly spaced frames.
+    Returns a list of tuples (PIL.Image, timestamp).
+    """
+    vidcap = cv2.VideoCapture(video_path)
+    total_frames = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))
+    fps = vidcap.get(cv2.CAP_PROP_FPS)
+    frames = []
+    if total_frames <= 0 or fps <= 0:
+        vidcap.release()
+        return frames
+    # Determine 10 evenly spaced frame indices.
+    frame_indices = np.linspace(0, total_frames - 1, 10, dtype=int)
+    for i in frame_indices:
+        vidcap.set(cv2.CAP_PROP_POS_FRAMES, i)
+        success, image = vidcap.read()
+        if success:
+            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+            pil_image = Image.fromarray(image)
+            timestamp = round(i / fps, 2)
+            frames.append((pil_image, timestamp))
+    vidcap.release()
+    return frames
+
+def extract_medicine_names(text):
+    """
+    Extracts medicine names from OCR text output.
+    Uses a combination of pattern matching and formatting to identify medications.
+    Returns a formatted list of medicines found.
+    """
+    # Common medicine patterns (extended to catch more formats)
+    lines = text.split('\n')
+    medicines = []
+    
+    # Look for patterns typical in prescriptions
+    for line in lines:
+        # Clean and standardize the line
+        clean_line = line.strip()
+        
+        # Skip very short lines, headers, or non-relevant text
+        if len(clean_line) < 3 or re.search(r'(prescription|rx|patient|name|date|doctor|hospital|clinic|address)', clean_line.lower()):
+            continue
+            
+        # Medicine names often appear at the beginning of lines, with dosage info following
+        # Look for tablet/capsule/mg indicators - strong indicators of medication
+        if re.search(r'(tab|tablet|cap|capsule|mg|ml|injection|syrup|solution|suspension|ointment|cream|gel|patch|suppository|inhaler|drops)', clean_line.lower()):
+            # Extract the likely medicine name - the part before the dosage/form or the entire line if it's short
+            medicine_match = re.split(r'(\d+\s*mg|\d+\s*ml|\d+\s*tab|\d+\s*cap)', clean_line, 1)[0].strip()
+            if medicine_match and len(medicine_match) > 2:
+                medicines.append(medicine_match)
+        
+        # Check for brand names or generic medication patterns
+        elif re.match(r'^[A-Z][a-z]+\s*[A-Z0-9]', clean_line) or re.match(r'^[A-Z][a-z]+', clean_line):
+            # Likely a medicine name starting with a capital letter
+            medicine_parts = re.split(r'(\d+|\s+\d+\s*times|\s+\d+\s*times\s+daily)', clean_line, 1)
+            if medicine_parts and len(medicine_parts[0]) > 2:
+                medicines.append(medicine_parts[0].strip())
+    
+    # Remove duplicates while preserving order
+    unique_medicines = []
+    for med in medicines:
+        if med not in unique_medicines:
+            unique_medicines.append(med)
+    
+    return unique_medicines
+
+# Model and Processor Setup
+# Qwen2VL OCR (default branch)
+QV_MODEL_ID = "prithivMLmods/Qwen2-VL-OCR-2B-Instruct"  # [or] prithivMLmods/Qwen2-VL-OCR2-2B-Instruct
+qwen_processor = AutoProcessor.from_pretrained(QV_MODEL_ID, trust_remote_code=True)
+qwen_model = Qwen2VLForConditionalGeneration.from_pretrained(
+    QV_MODEL_ID,
+    trust_remote_code=True,
+    torch_dtype=torch.float16
+).to("cuda").eval()
+
+# RolmOCR branch (@RolmOCR)
+ROLMOCR_MODEL_ID = "reducto/RolmOCR" 
+rolmocr_processor = AutoProcessor.from_pretrained(ROLMOCR_MODEL_ID, trust_remote_code=True)
+rolmocr_model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+    ROLMOCR_MODEL_ID,
+    trust_remote_code=True,
+    torch_dtype=torch.bfloat16
+).to("cuda").eval()
+
+# Main Inference Function
+@spaces.GPU
+def model_inference(input_dict, history):
+    text = input_dict["text"].strip()
+    files = input_dict.get("files", [])
+    
+    # Check for prescription-specific command
+    if text.lower().startswith("@prescription") or text.lower().startswith("@med"):
+        # Specific mode for medicine extraction
+        if not files:
+            yield "Error: Please upload a prescription image to extract medicine names."
+            return
+            
+        # Use RolmOCR for better text extraction from prescriptions
+        images = [load_image(image) for image in files[:1]]  # Taking just the first image for processing
+        
+        messages = [{
+            "role": "user",
+            "content": [
+                {"type": "image", "image": images[0]},
+                {"type": "text", "text": "Extract all text from this medical prescription image, focus on medicine names, dosages, and instructions."},
+            ],
+        }]
+        
+        prompt_full = rolmocr_processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+        inputs = rolmocr_processor(
+            text=[prompt_full],
+            images=images,
+            return_tensors="pt",
+            padding=True,
+        ).to("cuda")
+        
+        # First, get the complete OCR text
+        streamer = TextIteratorStreamer(rolmocr_processor, skip_prompt=True, skip_special_tokens=True)
+        generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=1024)
+        thread = Thread(target=rolmocr_model.generate, kwargs=generation_kwargs)
+        thread.start()
+        
+        ocr_text = ""
+        yield progress_bar_html("Processing Prescription with Medicine Extractor")
+        
+        for new_text in streamer:
+            ocr_text += new_text
+            ocr_text = ocr_text.replace("<|im_end|>", "")
+            time.sleep(0.01)
+        
+        # After getting full OCR text, extract medicine names
+        medicines = extract_medicine_names(ocr_text)
+        
+        # Format the results nicely
+        result = "## Extracted Medicine Names\n\n"
+        if medicines:
+            for i, med in enumerate(medicines, 1):
+                result += f"{i}. {med}\n"
+        else:
+            result += "No medicine names detected in the prescription.\n\n"
+            
+        result += "\n\n## Full OCR Text\n\n```\n" + ocr_text + "\n```"
+        yield result
+        return
+
+    # RolmOCR Inference (@RolmOCR)
+    if text.lower().startswith("@rolmocr"):
+        # Remove the tag from the query.
+        text_prompt = text[len("@rolmocr"):].strip()
+        # Check if a video is provided for inference.
+        if files and isinstance(files[0], str) and files[0].lower().endswith((".mp4", ".avi", ".mov")):
+            video_path = files[0]
+            frames = downsample_video(video_path)
+            if not frames:
+                yield "Error: Could not extract frames from the video."
+                return
+            # Build the message: prompt followed by each frame with its timestamp.
+            content_list = [{"type": "text", "text": text_prompt}]
+            for image, timestamp in frames:
+                content_list.append({"type": "text", "text": f"Frame {timestamp}:"})
+                content_list.append({"type": "image", "image": image})
+            messages = [{"role": "user", "content": content_list}]
+            # For video, extract images only.
+            video_images = [image for image, _ in frames]
+            prompt_full = rolmocr_processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+            inputs = rolmocr_processor(
+                text=[prompt_full],
+                images=video_images,
+                return_tensors="pt",
+                padding=True,
+            ).to("cuda")
+        else:
+            # Assume image(s) or text query.
+            if len(files) > 1:
+                images = [load_image(image) for image in files]
+            elif len(files) == 1:
+                images = [load_image(files[0])]
+            else:
+                images = []
+            if text_prompt == "" and not images:
+                yield "Error: Please input a text query and/or provide an image for the @RolmOCR feature."
+                return
+            messages = [{
+                "role": "user",
+                "content": [
+                    *[{"type": "image", "image": image} for image in images],
+                    {"type": "text", "text": text_prompt},
+                ],
+            }]
+            prompt_full = rolmocr_processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+            inputs = rolmocr_processor(
+                text=[prompt_full],
+                images=images if images else None,
+                return_tensors="pt",
+                padding=True,
+            ).to("cuda")
+        streamer = TextIteratorStreamer(rolmocr_processor, skip_prompt=True, skip_special_tokens=True)
+        generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=1024)
+        thread = Thread(target=rolmocr_model.generate, kwargs=generation_kwargs)
+        thread.start()
+        buffer = ""
+        # Use a different color scheme for RolmOCR (purple-themed).
+        yield progress_bar_html("Processing with Qwen2.5VL (RolmOCR)")
+        for new_text in streamer:
+            buffer += new_text
+            buffer = buffer.replace("<|im_end|>", "")
+            time.sleep(0.01)
+            yield buffer
+        return
+
+    # Default Inference: Qwen2VL OCR
+    # Process files: support multiple images.
+    if len(files) > 1:
+        images = [load_image(image) for image in files]
+    elif len(files) == 1:
+        images = [load_image(files[0])]
+    else:
+        images = []
+    
+    if text == "" and not images:
+        yield "Error: Please input a text query and optionally image(s)."
+        return
+    if text == "" and images:
+        yield "Error: Please input a text query along with the image(s)."
+        return
+
+    messages = [{
+        "role": "user",
+        "content": [
+            *[{"type": "image", "image": image} for image in images],
+            {"type": "text", "text": text},
+        ],
+    }]
+    prompt_full = qwen_processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
+    inputs = qwen_processor(
+        text=[prompt_full],
+        images=images if images else None,
+        return_tensors="pt",
+        padding=True,
+    ).to("cuda")
+    streamer = TextIteratorStreamer(qwen_processor, skip_prompt=True, skip_special_tokens=True)
+    generation_kwargs = dict(inputs, streamer=streamer, max_new_tokens=1024)
+    thread = Thread(target=qwen_model.generate, kwargs=generation_kwargs)
+    thread.start()
+    buffer = ""
+    yield progress_bar_html("Processing with Qwen2VL OCR")
+    for new_text in streamer:
+        buffer += new_text
+        buffer = buffer.replace("<|im_end|>", "")
+        time.sleep(0.01)
+        yield buffer
+
+# Gradio Interface
+examples = [
+    [{"text": "@Prescription Extract medicines from this prescription", "files": ["examples/prescription1.jpg"]}],
+    [{"text": "@RolmOCR OCR the Text in the Image", "files": ["rolm/1.jpeg"]}],
+    [{"text": "@RolmOCR Explain the Ad in Detail", "files": ["examples/videoplayback.mp4"]}],
+    [{"text": "@RolmOCR OCR the Image", "files": ["rolm/3.jpeg"]}],
+    [{"text": "Extract as JSON table from the table", "files": ["examples/4.jpg"]}],
+]
+
+css = """
+.gradio-container {
+    font-family: 'Roboto', sans-serif;
+}
+.prescription-header {
+    background-color: #4B0082;
+    color: white;
+    padding: 10px;
+    border-radius: 5px;
+    margin-bottom: 10px;
+}
+"""
+
+description = """
+# **Multimodal OCR with Medicine Extraction**
+
+## Modes:
+- **@Prescription** - Upload a prescription image to extract medicine names
+- **@RolmOCR** - Use RolmOCR for general text extraction
+- **Default** - Use Qwen2VL OCR for general purposes
+
+Upload your medical prescription images and get the medicine names extracted automatically!
+"""
+
+demo = gr.ChatInterface(
+    fn=model_inference,
+    description=description,
+    examples=examples,
+    textbox=gr.MultimodalTextbox(
+        label="Query Input", 
+        file_types=["image", "video"], 
+        file_count="multiple", 
+        placeholder="Use @Prescription to extract medicines, @RolmOCR for RolmOCR, or leave blank for default Qwen2VL OCR"
+    ),
+    stop_btn="Stop Generation",
+    multimodal=True,
+    cache_examples=False,
+    css=css
+)
+
+demo.launch(debug=True)