Update README.md

README.md

@@ -3,159 +3,49 @@ library_name: transformers
 pipeline_tag: image-text-to-text
 ---
 
-## How to Use the *ferret-gemma* Model
+Ferret-UI is the first UI-centric multimodal large language model (MLLM) designed for referring, grounding, and reasoning tasks.
+Built on Gemma-2B and Llama-3-8B, it is capable of executing complex UI tasks.
+This is the Gemma-2B version of ferret-ui. It follows from [this paper](https://arxiv.org/pdf/2404.05719) by Apple.
 
-Please download and save `builder.py`, `conversation.py` locally.
 
-```python
-import torch
-from PIL import Image
-from conversation import conv_templates
-from builder import load_pretrained_model  # Assuming this is your custom model loader
-from functools import partial
-import numpy as np
-
-# define the task categories
-box_in_tasks = ['widgetcaptions', 'taperception', 'ocr', 'icon_recognition', 'widget_classification', 'example_0']
-box_out_tasks = ['widget_listing', 'find_text', 'find_icons', 'find_widget', 'conversation_interaction']
-no_box_tasks = ['screen2words', 'detailed_description', 'conversation_perception', 'gpt4']
-
-# function to generate the mask
-def generate_mask_for_feature(coor, raw_w, raw_h, mask=None):
-    """
-    Generates a region mask based on provided coordinates.
-    Handles both point and box input.
-    """
-    if mask is not None:
-        assert mask.shape[0] == raw_w and mask.shape[1] == raw_h
-    coor_mask = np.zeros((raw_w, raw_h))
-
-    # if it's a point (2 coordinates)
-    if len(coor) == 2:
-        span = 5  # Define the span for the point
-        x_min = max(0, coor[0] - span)
-        x_max = min(raw_w, coor[0] + span + 1)
-        y_min = max(0, coor[1] - span)
-        y_max = min(raw_h, coor[1] + span + 1)
-        coor_mask[int(x_min):int(x_max), int(y_min):int(y_max)] = 1
-        assert (coor_mask == 1).any(), f"coor: {coor}, raw_w: {raw_w}, raw_h: {raw_h}"
+## How to Use the *Ferret-UI-Gemma2b* Model
 
-    # if it's a box (4 coordinates)
-    elif len(coor) == 4:
-        coor_mask[coor[0]:coor[2]+1, coor[1]:coor[3]+1] = 1
-        if mask is not None:
-            coor_mask = coor_mask * mask
+You will need first to download `builder.py`, `conversation.py`, and `inference.py` locally.
 
-    # Convert to torch tensor and ensure it contains non-zero values
-    coor_mask = torch.from_numpy(coor_mask)
-    assert len(coor_mask.nonzero()) != 0, "Generated mask is empty :("
-
-    return coor_mask
-```
-### Now, define the infer function
-```python
-def infer_single_prompt(image_path, prompt, model_path, region=None, model_name="ferret_gemma", conv_mode="ferret_gemma_instruct"):
-    img = Image.open(image_path).convert('RGB')
-
-    # this loads the model, image processor and tokenizer
-    tokenizer, model, image_processor, context_len = load_pretrained_model(model_path, None, model_name)
-
-    # define the image size (e.g., 224x224 or 336x336)
-    image_size = {"height": 336, "width": 336}
-
-    # process the image
-    image_tensor = image_processor.preprocess(
-        img,
-        return_tensors='pt',
-        do_resize=True,
-        do_center_crop=False,
-        size=(image_size['height'], image_size['width'])
-    )['pixel_values'][0].unsqueeze(0)
-
-    image_tensor = image_tensor.half().cuda()
-
-    # generate the prompt per template requirement
-    conv = conv_templates[conv_mode].copy()
-    conv.append_message(conv.roles[0], prompt)
-    conv.append_message(conv.roles[1], None)
-    prompt_input = conv.get_prompt()
-
-    # tokenize prompt
-    input_ids = tokenizer(prompt_input, return_tensors='pt')['input_ids'].cuda()
-
-    # region mask logic (if region is provided)
-    region_masks = None
-    if region is not None:
-        raw_w, raw_h = img.size
-        region_masks = generate_mask_for_feature(region, raw_w, raw_h).unsqueeze(0).cuda().half()
-        region_masks = [[region_masks]]  # Wrap the mask in lists as expected by the model
-
-    # generate model output
-    with torch.inference_mode():
-        # Use region_masks in model's forward call
-        model.orig_forward = model.forward
-        model.forward = partial(
-            model.orig_forward,
-            region_masks=region_masks
-        )
-        output_ids = model.generate(
-            input_ids,
-            images=image_tensor,
-            max_new_tokens=1024,
-            num_beams=1,
-            region_masks=region_masks,  # pass the region mask to the model
-            image_sizes=[img.size]
-        )
-        model.forward = model.orig_forward
-
-    # we decode the output
-    output_text = tokenizer.batch_decode(output_ids, skip_special_tokens=True)[0]
-    return output_text.strip()
-```
-
-# We also define a task-specific inference function
-```python
-def infer_ui_task(image_path, prompt, model_path, task, region=None):
-    """
-    Handles task types: box_in_tasks, box_out_tasks, no_box_tasks.
-    """
-    if task in box_in_tasks and region is None:
-        raise ValueError(f"Task {task} requires a bounding box region.")
-    
-    if task in box_in_tasks:
-        print(f"Processing {task} with bounding box region.")
-        return infer_single_prompt(image_path, prompt, model_path, region)
-    
-    elif task in box_out_tasks:
-        print(f"Processing {task} without bounding box region.")
-        return infer_single_prompt(image_path, prompt, model_path)
-    
-    elif task in no_box_tasks:
-        print(f"Processing {task} without image or bounding box.")
-        return infer_single_prompt(image_path, prompt, model_path)
-    
-    else:
-        raise ValueError(f"Unknown task type: {task}")
+```bash
+wget https://huggingface.co/jadechoghari/ferret-gemma/raw/main/conversation.py
+wget https://huggingface.co/jadechoghari/ferret-gemma/raw/main/builder.py
+wget https://huggingface.co/jadechoghari/ferret-gemma/raw/main/inference.py
 ```
 
 ### Usage:
 ```python
-# Example image and model paths
+from inference import infer_ui_task
+# Pass an image and the online model path
 image_path = 'image.jpg'
-model_path = 'jadechoghari/ferret-gemma'
+model_path = 'jadechoghari/Ferret-UI-Gemma2b'
+```
 
-# Task requiring bounding box
-task = 'widgetcaptions'
+### Task requiring bounding box
+Choose a task from ['widgetcaptions', 'taperception', 'ocr', 'icon_recognition', 'widget_classification', 'example_0']
+```python
+task = 'widgetcaptions' 
 region = (50, 50, 200, 200)
 result = infer_ui_task(image_path, "Describe the contents of the box.", model_path, task, region=region)
 print("Result:", result)
+```
 
-# Task not requiring bounding box
+### Task not requiring bounding box
+Choose a task from ['widget_listing', 'find_text', 'find_icons', 'find_widget', 'conversation_interaction']
+```python
 task = 'conversation_interaction'
 result = infer_ui_task(image_path, "How do I navigate to the Games tab?", model_path, task)
 print("Result:", result)
+```
 
-# Task with no image processing
+### Task with no image processing
+Choose a task from ['screen2words', 'detailed_description', 'conversation_perception', 'gpt4']
+```python
 task = 'detailed_description'
 result = infer_ui_task(image_path, "Please describe the screen in detail.", model_path, task)
 print("Result:", result)