Create app.py

app.py

@@ -0,0 +1,211 @@
+import gradio as gr
+from PIL import Image
+
+from transformers import Qwen2VLForConditionalGeneration, GenerationConfig
+import base64
+from io import BytesIO
+from PIL import Image
+import re
+
+
+def extract_answer_content(text: str) -> str:
+    """
+    Extracts the content between <answer> and </answer> tags.
+    If no closing tag is found, extracts everything after the first <answer>.
+
+    Returns:
+        str: The extracted content.
+    """
+    text = text.replace("```", " ").replace("json", " ").strip()
+
+    # Try to find full <answer>...</answer>
+    match = re.search(r"<answer>(.*?)</answer>", text, re.DOTALL)
+    if match:
+        return match.group(1).strip()
+    # Fallback: everything after the first <answer>
+    match = re.search(r"<answer>(.*)", text, re.DOTALL)
+    return match.group(1).strip() if match else text
+
+
+def encode_image_to_base64(image: Image.Image, format: str = "PNG") -> str:
+    """
+    Encode a PIL Image to a base64 string.
+
+    Args:
+        image (PIL.Image): The image to encode.
+        format (str): Image format to use (e.g., "PNG", "JPEG"). Default is "PNG".
+
+    Returns:
+        str: Base64-encoded string of the image.
+    """
+    buffer = BytesIO()
+    image.save(buffer, format=format)
+    buffer.seek(0)
+    encoded_string = base64.b64encode(buffer.read()).decode("utf-8")
+    return encoded_string
+
+
+
+SYSTEM_PROMPT = (
+    "A conversation between User and Assistant. The user asks a question, and the Assistant solves it. The assistant "
+    "first thinks about the reasoning process in the mind and then provides the user with the answer. The reasoning "
+    "process and answer are enclosed within <think> </think> and <answer> </answer> tags, respectively, i.e., "
+    "<think> reasoning process here </think><answer> answer here </answer>"
+)
+
+
+
+# Load model and processor
+processor = AutoProcessor.from_pretrained("JosephZ/qwen2vl-7b-sft-grpo-close-sgg", max_pixels=1024*28*28)
+
+device='cuda' if torch.cuda.is_available() else "cpu"
+model = Qwen2VLForConditionalGeneration.from_pretrained("JosephZ/qwen2vl-7b-sft-grpo-close-sgg",
+             torch_dtype=torch.bfloat16,
+             device_map=device)
+
+generation_config=GenerationConfig(
+        do_sample=True,
+        temperature=0.01,
+        top_k=1,
+        top_p=0.001,
+        repetition_penalty=1.0,
+        max_new_tokens=2048,
+)
+
+def build_prompt(image, use_text):
+    #base64_image = encode_image_to_base64(image)
+    messages = [
+        {
+            "role": "system",
+            "content": SYSTEM_PROMPT
+        },
+        {
+            "role": "user",
+            "content": [
+                #{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}},
+                 {"type": "image"},
+                {"type": "text", "text": user_text},
+            ],
+        },
+    ]    
+    return messages
+
+PROMPT_CLOSE="""Generate a structured scene graph for an image using the specified object and relationship categories.
+
+### **Output Format:**
+```json
+{
+  "objects": [
+    {"id": "object_name.number", "bbox": [x1, y1, x2, y2]},
+    ...
+  ],
+  "relationships": [
+    {"subject": "object_name.number", "predicate": "relationship_type", "object": "object_name.number"},
+    ...
+  ]
+}
+```
+
+### **Guidelines:**
+- **Objects:**
+  - Assign unique IDs in the format `"object_name.number"` (e.g., `"person.1"`). The **object_name** must belong to the predefined object set: `["airplane", "animal", "arm", "bag", "banana", "basket", "beach", "bear", "bed", "bench", "bike", "bird", "board", "boat", "book", "boot", "bottle", "bowl", "box", "boy", "branch", "building", "bus", "cabinet", "cap", "car", "cat", "chair", "child", "clock", "coat", "counter", "cow", "cup", "curtain", "desk", "dog", "door", "drawer", "ear", "elephant", "engine", "eye", "face", "fence", "finger", "flag", "flower", "food", "fork", "fruit", "giraffe", "girl", "glass", "glove", "guy", "hair", "hand", "handle", "hat", "head", "helmet", "hill", "horse", "house", "jacket", "jean", "kid", "kite", "lady", "lamp", "laptop", "leaf", "leg", "letter", "light", "logo", "man", "men", "motorcycle", "mountain", "mouth", "neck", "nose", "number", "orange", "pant", "paper", "paw", "people", "person", "phone", "pillow", "pizza", "plane", "plant", "plate", "player", "pole", "post", "pot", "racket", "railing", "rock", "roof", "room", "screen", "seat", "sheep", "shelf", "shirt", "shoe", "short", "sidewalk", "sign", "sink", "skateboard", "ski", "skier", "sneaker", "snow", "sock", "stand", "street", "surfboard", "table", "tail", "tie", "tile", "tire", "toilet", "towel", "tower", "track", "train", "tree", "truck", "trunk", "umbrella", "vase", "vegetable", "vehicle", "wave", "wheel", "window", "windshield", "wing", "wire", "woman", "zebra"]`.
+  - Provide a bounding box `[x1, y1, x2, y2]` in integer pixel format.
+  - Include all visible objects, even if they have no relationships.
+
+- **Relationships:**
+  - Define relationships using `"subject"`, `"predicate"`, and `"object"`.
+  - The **predicate** must belong to the predefined relationship set: `["above", "across", "against", "along", "and", "at", "attached to", "behind", "belonging to", "between", "carrying", "covered in", "covering", "eating", "flying in", "for", "from", "growing on", "hanging from", "has", "holding", "in", "in front of", "laying on", "looking at", "lying on", "made of", "mounted on", "near", "of", "on", "on back of", "over", "painted on", "parked on", "part of", "playing", "riding", "says", "sitting on", "standing on", "to", "under", "using", "walking in", "walking on", "watching", "wearing", "wears", "with"]`.
+  - Omit relationships for orphan objects.
+
+### **Example Output:**
+```json
+{
+  "objects": [
+    {"id": "person.1", "bbox": [120, 200, 350, 700]},
+    {"id": "bike.2", "bbox": [100, 600, 400, 800]},
+    {"id": "helmet.3", "bbox": [150, 150, 280, 240]},
+    {"id": "tree.4", "bbox": [500, 100, 750, 700]}
+  ],
+  "relationships": [
+    {"subject": "person.1", "predicate": "riding", "object": "bike.2"},
+    {"subject": "person.1", "predicate": "wearing", "object": "helmet.3"}
+  ]
+}
+```
+
+Now, generate the complete scene graph for the provided image:
+"""
+
+def is_box(item):
+    return (
+        isinstance(item, (list, tuple)) and
+        len(item) == 4 and
+        all(isinstance(e, (int, float)) for e in item)
+    )
+
+def scale_box(box, scale):
+    sw, sh = scale
+    return [int(box[0]*sw), int(box[1]*sh), int(box[2]*sw), int(box[3]*sh)]
+
+
+def generate_sgg(image):
+    iw, ih = image.size
+    scale_factors = (iw / 1000.0, ih / 1000.0)
+
+    conversation = build_prompt(image, PROMPT_CLOSE)
+    text_prompt = processor.apply_chat_template(conversation, add_generation_prompt=True)
+
+    inputs = processor(
+        text=[text_prompt], images=[image], padding=True, return_tensors="pt"
+    )
+    inputs = inputs.to(model.device)
+    
+    output_ids = model.generate(**inputs, generation_config=generation_config)
+    generated_ids = [
+        output_ids[len(input_ids) :]
+        for input_ids, output_ids in zip(inputs.input_ids, output_ids)
+    ]
+    output_text = processor.batch_decode(
+        generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True
+    )[0]    
+
+    resp = extract_answer_content(output_text)
+    
+    try:
+        resp = json.loads(resp)
+        pred_objs = resp['objects']
+        pred_rels = resp['relationships']
+        new_objs = []
+        for obj in pred_objs:
+            assert len(obj['bbox']) == 4, "len(obj['bbox']) != 4"
+            assert is_box(obj['bbox']), "invalid box :{}".format(obj['bbox'])
+            assert 'id' in obj, "invalid obj:{}".format(obj)        
+            obj['bbox'] = scale_box(obj['bbox'], scale_factors)
+            new_objs.append(obj)
+        pred_objs = new_objs
+        resp = {"objects": pred_objs, "relationships": pred_rels}
+
+        # visualize pred_objs
+        draw = ImageDraw.Draw(image)
+        try:
+            font = ImageFont.truetype("arial.ttf", 16)
+        except:
+            font = ImageFont.load_default()
+
+        for obj in pred_objs:
+            bbox = obj["bbox"]
+            label = obj["id"]
+            draw.rectangle(bbox, outline="red", width=3)
+            draw.text((bbox[0], bbox[1] - 15), label, fill="red", font=font)
+    except:
+        pass
+
+    return json.dumps(resp) if not isinstance(resp, str) else resp
+
+gr.Interface(
+    fn=generate_sgg,
+    inputs=gr.Image(type="pil"),
+    outputs="text",
+    title="R1-SGG: Compile Scene Graphs with Reinforcement Learning",
+    description="Upload an image and generate a structured scene graph in JSON format."
+).launch()